diff --git a/tests/v1/metrics/test_engine_logger_apis.py b/tests/v1/metrics/test_engine_logger_apis.py new file mode 100644 index 00000000000..c633a2d4a4f --- /dev/null +++ b/tests/v1/metrics/test_engine_logger_apis.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM +from vllm.v1.metrics.loggers import PrometheusStatLogger + + +@pytest.mark.asyncio +async def test_async_llm_add_logger(): + # Minimal model config for test + model_name = "distilbert/distilgpt2" + dtype = "half" + engine_args = AsyncEngineArgs( + model=model_name, + dtype=dtype, + disable_log_stats=False, + enforce_eager=True, + ) + + # Force empty list to avoid default loggers + engine = AsyncLLM.from_engine_args(engine_args, stat_loggers=[]) + + # Add PrometheusStatLogger and verify no exception is raised + await engine.add_logger(PrometheusStatLogger) + + # Verify that logger is present in the first DP rank + assert len(engine.stat_loggers[0]) == 1 + assert isinstance(engine.stat_loggers[0][0], PrometheusStatLogger) \ No newline at end of file diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 3754570dfaa..30e01a43903 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -608,6 +608,19 @@ async def collective_rpc(self, return await self.engine_core.collective_rpc_async( method, timeout, args, kwargs) + async def add_logger(self, logger_factory: StatLoggerFactory) -> None: + if not self.log_stats: + raise RuntimeError( + "Stat logging is disabled. Set `disable_log_stats=False` " + "argument to enable.") + + engine_num = self.vllm_config.parallel_config.data_parallel_size + if len(self.stat_loggers) == 0: + self.stat_loggers = [[] for _ in range(engine_num)] + + for i, logger_list in enumerate(self.stat_loggers): + logger_list.append(logger_factory(self.vllm_config, i)) + @property def is_running(self) -> bool: # Is None before the loop is started.