[Misc] Have AsyncLLM custom_stat_loggers extend default logger list (#20952)

Signed-off-by: Seiji Eicher <[email protected]>
Signed-off-by: Seiji Eicher <[email protected]>
Co-authored-by: Nick Hill <[email protected]>

Author: eicherseiji

tests/v1/engine/test_async_llm.py

@@ -393,7 +393,7 @@ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
 async def test_customize_loggers(monkeypatch):
     """Test that we can customize the loggers.
     If a customized logger is provided at the init, it should
-    be used directly.
+    be added to the default loggers.
     """
 
     with monkeypatch.context() as m, ExitStack() as after:
@@ -410,7 +410,8 @@ async def test_customize_loggers(monkeypatch):
 
         stat_loggers = engine.logger_manager.per_engine_logger_dict
         assert len(stat_loggers) == 1
-        assert len(stat_loggers[0]) == 1
+        assert len(
+            stat_loggers[0]) == 2  # LoggingStatLogger + MockLoggingStatLogger
         stat_loggers[0][0].log.assert_called_once()
 
 

tests/v1/metrics/test_engine_logger_apis.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+
+import pytest
+
+from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM
+from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
+
+
+class DummyStatLogger:
+    """
+    A dummy stat logger for testing purposes.
+    Implements the minimal interface expected by StatLoggerManager.
+    """
+
+    def __init__(self, vllm_config, engine_idx):
+        self.vllm_config = vllm_config
+        self.engine_idx = engine_idx
+        self.recorded = []
+        self.logged = False
+        self.engine_initialized = False
+
+    def record(self, scheduler_stats, iteration_stats, engine_idx):
+        self.recorded.append((scheduler_stats, iteration_stats, engine_idx))
+
+    def log(self):
+        self.logged = True
+
+    def log_engine_initialized(self):
+        self.engine_initialized = True
+
+
+@pytest.fixture
+def log_stats_enabled_engine_args():
+    """
+    Shared fixture providing common AsyncEngineArgs configuration
+    used across multiple tests.
+    """
+    return AsyncEngineArgs(
+        model="distilbert/distilgpt2",
+        dtype="half",
+        disable_log_stats=False,
+        enforce_eager=True,
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_llm_replace_default_loggers(
+        log_stats_enabled_engine_args):
+    """
+    RayPrometheusStatLogger should replace the default PrometheusStatLogger
+    """
+
+    engine = AsyncLLM.from_engine_args(log_stats_enabled_engine_args,
+                                       stat_loggers=[RayPrometheusStatLogger])
+    assert isinstance(engine.logger_manager.prometheus_logger,
+                      RayPrometheusStatLogger)
+    engine.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_async_llm_add_to_default_loggers(log_stats_enabled_engine_args):
+    """
+    It's still possible to use custom stat loggers exclusively by passing 
+    disable_log_stats=True in addition to a list of custom stat loggers.
+    """
+    # Create engine_args with disable_log_stats=True for this test
+    disabled_log_engine_args = copy.deepcopy(log_stats_enabled_engine_args)
+    disabled_log_engine_args.disable_log_stats = True
+
+    # Disable default loggers; pass custom stat logger to the constructor
+    engine = AsyncLLM.from_engine_args(disabled_log_engine_args,
+                                       stat_loggers=[DummyStatLogger])
+
+    assert len(engine.logger_manager.per_engine_logger_dict[0]) == 1
+    assert isinstance(engine.logger_manager.per_engine_logger_dict[0][0],
+                      DummyStatLogger)
+
+    # log_stats is still True, since custom stat loggers are used
+    assert engine.log_stats
+
+    engine.shutdown()

vllm/v1/engine/async_llm.py

@@ -98,7 +98,12 @@ def __init__(
         self.model_config = vllm_config.model_config
         self.vllm_config = vllm_config
         self.log_requests = log_requests
-        self.log_stats = log_stats
+
+        self.log_stats = log_stats or (stat_loggers is not None)
+        if not log_stats and stat_loggers is not None:
+            logger.info(
+                "AsyncLLM created with log_stats=False and non-empty custom "
+                "logger list; enabling logging without default stat loggers")
 
         if self.model_config.skip_tokenizer_init:
             self.tokenizer = None
@@ -137,6 +142,7 @@ def __init__(
                 vllm_config=vllm_config,
                 engine_idxs=self.engine_core.engine_ranks_managed,
                 custom_stat_loggers=stat_loggers,
+                enable_default_loggers=log_stats,
             )
             self.logger_manager.log_engine_initialized()
 

vllm/v1/metrics/loggers.py

@@ -651,16 +651,16 @@ def __init__(
         vllm_config: VllmConfig,
         engine_idxs: Optional[list[int]] = None,
         custom_stat_loggers: Optional[list[StatLoggerFactory]] = None,
+        enable_default_loggers: bool = True,
     ):
         self.engine_idxs = engine_idxs if engine_idxs else [0]
 
-        factories: list[StatLoggerFactory]
+        factories: list[StatLoggerFactory] = []
         if custom_stat_loggers is not None:
-            factories = custom_stat_loggers
-        else:
-            factories = []
-            if logger.isEnabledFor(logging.INFO):
-                factories.append(LoggingStatLogger)
+            factories.extend(custom_stat_loggers)
+
+        if enable_default_loggers and logger.isEnabledFor(logging.INFO):
+            factories.append(LoggingStatLogger)
 
         # engine_idx: StatLogger
         self.per_engine_logger_dict: dict[int, list[StatLoggerBase]] = {}