Allow v1 async llm engine to use custom logger

Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>

Author: liuzijing2014

tests/v1/engine/test_async_llm.py

@@ -3,6 +3,7 @@
 import asyncio
 from contextlib import ExitStack
 from typing import Optional
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -13,11 +14,16 @@
 from vllm.platforms import current_platform
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.metrics.loggers import (PROMETHEUS_LOGGING_LOGGER_NAME,
+                                     STANDARD_LOGGING_LOGGER_NAME,
+                                     LoggingStatLogger, StatLoggerBase)
 
 if not current_platform.is_cuda():
     pytest.skip(reason="V1 currently only supported on CUDA.",
                 allow_module_level=True)
 
+TEST_LOGGER_NAME = "test_logger"
+
 TEXT_ENGINE_ARGS = AsyncEngineArgs(model="meta-llama/Llama-3.2-1B-Instruct",
                                    enforce_eager=True,
                                    disable_log_requests=True)
@@ -220,3 +226,47 @@ async def test_finished_flag(monkeypatch, n: int,
         # Assert only the last output has the finished flag set
         assert all(not out.finished for out in outputs[:-1])
         assert outputs[-1].finished
+
+
+def get_customized_logger_mock() -> StatLoggerBase:
+    logger = LoggingStatLogger()
+    logger.log = MagicMock()
+    return logger
+
+
+@pytest.mark.parametrize(
+    "loggers",
+    [{
+        TEST_LOGGER_NAME: get_customized_logger_mock()
+    }, None],
+)
+@pytest.mark.asyncio
+async def test_customize_loggers(
+    monkeypatch,
+    loggers: Optional[dict[str, StatLoggerBase]],
+):
+    """Test that we can customize the loggers.
+    Test case #1: Not customized logger is provided at the init, default loggers
+    would be initialized. Thus, we should be able to remove those and add a
+    customized one later.
+    Test case #2: If a customized logger is provided at the init, it should
+    be used directly.
+    """
+
+    with monkeypatch.context() as m, ExitStack() as after:
+        m.setenv("VLLM_USE_V1", "1")
+
+        engine = AsyncLLM.from_engine_args(
+            TEXT_ENGINE_ARGS,
+            stat_loggers=loggers,
+        )
+        after.callback(engine.shutdown)
+
+        if loggers is None:
+            engine.remove_logger(PROMETHEUS_LOGGING_LOGGER_NAME)
+            engine.remove_logger(STANDARD_LOGGING_LOGGER_NAME)
+            engine.add_logger(TEST_LOGGER_NAME, get_customized_logger_mock())
+
+        await engine.do_log_stats()
+        for logger in engine.stat_loggers.values():
+            logger.log.assert_called_once()

vllm/third_party/pynvml.py

@@ -30,16 +30,16 @@
 # THE POSSIBILITY OF SUCH DAMAGE.
 #####
 
+import os
+import string
+import sys
+import threading
 ##
 # Python bindings for the NVML library
 ##
 from ctypes import *
 from ctypes.util import find_library
 from functools import wraps
-import sys
-import os
-import threading
-import string
 
 ## C Type mappings ##
 ## Enums

vllm/v1/engine/async_llm.py

@@ -29,7 +29,9 @@
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.engine.processor import Processor
 from vllm.v1.executor.abstract import Executor
-from vllm.v1.metrics.loggers import (LoggingStatLogger, PrometheusStatLogger,
+from vllm.v1.metrics.loggers import (PROMETHEUS_LOGGING_LOGGER_NAME,
+                                     STANDARD_LOGGING_LOGGER_NAME,
+                                     LoggingStatLogger, PrometheusStatLogger,
                                      StatLoggerBase)
 from vllm.v1.metrics.stats import IterationStats, SchedulerStats
 
@@ -48,6 +50,7 @@ def __init__(
         use_cached_outputs: bool = False,
         log_requests: bool = True,
         start_engine_loop: bool = True,
+        stat_loggers: Optional[dict[str, StatLoggerBase]] = None,
     ) -> None:
 
         assert start_engine_loop
@@ -56,11 +59,16 @@ def __init__(
 
         self.log_requests = log_requests
         self.log_stats = log_stats
-        self.stat_loggers: list[StatLoggerBase] = []
+        self.stat_loggers: dict[str, StatLoggerBase] = dict()
         if self.log_stats:
-            if logger.isEnabledFor(logging.INFO):
-                self.stat_loggers.append(LoggingStatLogger())
-            self.stat_loggers.append(PrometheusStatLogger(vllm_config))
+            if stat_loggers is not None:
+                self.stat_loggers = stat_loggers
+            else:
+                if logger.isEnabledFor(logging.INFO):
+                    self.stat_loggers[STANDARD_LOGGING_LOGGER_NAME] = (
+                        LoggingStatLogger())
+                self.stat_loggers[PROMETHEUS_LOGGING_LOGGER_NAME] = (
+                    PrometheusStatLogger(vllm_config))
 
         # Tokenizer (+ ensure liveness if running in another process).
         self.tokenizer = init_tokenizer_from_configs(
@@ -99,6 +107,7 @@ def from_engine_args(
         engine_config: Optional[VllmConfig] = None,
         start_engine_loop: bool = True,
         usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
+        stat_loggers: Optional[dict[str, StatLoggerBase]] = None,
     ) -> "AsyncLLM":
         """Create an AsyncLLM from the EngineArgs."""
 
@@ -118,6 +127,7 @@ def from_engine_args(
             log_stats=not engine_args.disable_log_stats,
             start_engine_loop=start_engine_loop,
             usage_context=usage_context,
+            stat_loggers=stat_loggers,
         )
 
     def shutdown(self):
@@ -313,7 +323,7 @@ def _record_stats(
             return
 
         assert scheduler_stats is not None
-        for stat_logger in self.stat_loggers:
+        for stat_logger in self.stat_loggers.values():
             stat_logger.record(scheduler_stats=scheduler_stats,
                                iteration_stats=iteration_stats)
 
@@ -351,9 +361,27 @@ async def do_log_stats(
         scheduler_outputs=None,
         model_output=None,
     ) -> None:
-        for stat_logger in self.stat_loggers:
+        for stat_logger in self.stat_loggers.values():
             stat_logger.log()
 
+    def add_logger(self, logger_name: str, logger: StatLoggerBase) -> None:
+        if not self.log_stats:
+            raise RuntimeError(
+                "Stat logging is disabled. Set `disable_log_stats=False` "
+                "argument to enable.")
+        if logger_name in self.stat_loggers:
+            raise KeyError(f"Logger with name {logger_name} already exists.")
+        self.stat_loggers[logger_name] = logger
+
+    def remove_logger(self, logger_name: str) -> None:
+        if not self.log_stats:
+            raise RuntimeError(
+                "Stat logging is disabled. Set `disable_log_stats=False` "
+                "argument to enable.")
+        if logger_name not in self.stat_loggers:
+            raise KeyError(f"Logger with name {logger_name} does not exist.")
+        del self.stat_loggers[logger_name]
+
     async def check_health(self) -> None:
         logger.debug("Called check_health.")
 

vllm/v1/metrics/loggers.py

@@ -17,6 +17,9 @@
 
 _LOCAL_LOGGING_INTERVAL_SEC = 5.0
 
+STANDARD_LOGGING_LOGGER_NAME = "logging"
+PROMETHEUS_LOGGING_LOGGER_NAME = "prometheus"
+
 
 class StatLoggerBase(ABC):