A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[V1][Metrics] Allow V1 AsyncLLM to use custom logger (#14661)
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Zijing Liu committed
Apr 25, 2025 at 22:05 UTC
53e8cf53a4d1631c49962c9b5ddf172b0bccd3ec
4 files changed
+118
-30
tests/v1/engine/test_async_llm.py
+33
@@ -3,16 +3,19 @@
3
import asyncio
4
from contextlib import ExitStack
5
from typing import Optional
6
+from unittest.mock import MagicMock
7
8
import pytest
9
10
from vllm import SamplingParams
11
from vllm.assets.image import ImageAsset
12
+from vllm.config import VllmConfig
13
from vllm.engine.arg_utils import AsyncEngineArgs
14
from vllm.inputs import PromptType
15
from vllm.platforms import current_platform
16
from vllm.sampling_params import RequestOutputKind
17
from vllm.v1.engine.async_llm import AsyncLLM
18
+from vllm.v1.metrics.loggers import LoggingStatLogger
19
20
if not current_platform.is_cuda():
21
pytest.skip(reason="V1 currently only supported on CUDA.",
@@ -216,3 +219,33 @@ async def test_finished_flag(monkeypatch: pytest.MonkeyPatch, n: int,
219
# Assert only the last output has the finished flag set
220
assert all(not out.finished for out in outputs[:-1])
221
assert outputs[-1].finished
222
+
223
+
224
+class MockLoggingStatLogger(LoggingStatLogger):
225
+
226
+ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
227
+ super().__init__(vllm_config, engine_index)
228
+ self.log = MagicMock()
229
+
230
+
231
+@pytest.mark.asyncio
232
+async def test_customize_loggers(monkeypatch):
233
+ """Test that we can customize the loggers.
234
+ If a customized logger is provided at the init, it should
235
+ be used directly.
236
+ """
237
+
238
+ with monkeypatch.context() as m, ExitStack() as after:
239
+ m.setenv("VLLM_USE_V1", "1")
240
+
241
+ engine = AsyncLLM.from_engine_args(
242
+ TEXT_ENGINE_ARGS,
243
+ stat_loggers=[MockLoggingStatLogger],
244
+ )
245
+ after.callback(engine.shutdown)
246
+
247
+ await engine.do_log_stats()
248
+
249
+ assert len(engine.stat_loggers) == 1
250
+ assert len(engine.stat_loggers[0]) == 1
251
+ engine.stat_loggers[0][0].log.assert_called_once()
vllm/v1/engine/async_llm.py
+33
-19
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
import asyncio
3
-import logging
3
from collections.abc import AsyncGenerator, Mapping
4
from copy import copy
5
from typing import Optional, Union
@@ -33,8 +32,8 @@ from vllm.v1.engine.output_processor import (OutputProcessor,
32
from vllm.v1.engine.parallel_sampling import ParentRequest
33
from vllm.v1.engine.processor import Processor
34
from vllm.v1.executor.abstract import Executor
36
-from vllm.v1.metrics.loggers import (LoggingStatLogger, PrometheusStatLogger,
37
- StatLoggerBase)
35
+from vllm.v1.metrics.loggers import (StatLoggerBase, StatLoggerFactory,
36
+ setup_default_loggers)
37
from vllm.v1.metrics.stats import IterationStats, SchedulerStats
38
39
logger = init_logger(__name__)
@@ -52,7 +51,28 @@ class AsyncLLM(EngineClient):
51
use_cached_outputs: bool = False,
52
log_requests: bool = True,
53
start_engine_loop: bool = True,
54
+ stat_loggers: Optional[list[StatLoggerFactory]] = None,
55
) -> None:
56
+ """
57
+ Create an AsyncLLM.
58
+
59
+ Args:
60
+ vllm_config: global configuration.
61
+ executor_class: an Executor impl, e.g. MultiprocExecutor.
62
+ log_stats: Whether to log stats.
63
+ usage_context: Usage context of the LLM.
64
+ mm_registry: Multi-modal registry.
65
+ use_cached_outputs: Whether to use cached outputs.
66
+ log_requests: Whether to log requests.
67
+ start_engine_loop: Whether to start the engine loop.
68
+ stat_loggers: customized stat loggers for the engine.
69
+ If not provided, default stat loggers will be used.
70
+ PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE
71
+ IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE.
72
+
73
+ Returns:
74
+ None
75
+ """
76
if not envs.VLLM_USE_V1:
77
raise ValueError(
78
"Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
@@ -66,15 +86,12 @@ class AsyncLLM(EngineClient):
86
self.log_stats = log_stats
87
88
# Set up stat loggers; independent set for each DP rank.
69
- self.stat_loggers: list[list[StatLoggerBase]] = []
70
- if self.log_stats:
71
- for i in range(vllm_config.parallel_config.data_parallel_size):
72
- loggers: list[StatLoggerBase] = []
73
- if logger.isEnabledFor(logging.INFO):
74
- loggers.append(LoggingStatLogger(engine_index=i))
75
- loggers.append(
76
- PrometheusStatLogger(vllm_config, engine_index=i))
77
- self.stat_loggers.append(loggers)
89
+ self.stat_loggers: list[list[StatLoggerBase]] = setup_default_loggers(
90
+ vllm_config=vllm_config,
91
+ log_stats=self.log_stats,
92
+ engine_num=vllm_config.parallel_config.data_parallel_size,
93
+ custom_stat_loggers=stat_loggers,
94
+ )
95
96
# Tokenizer (+ ensure liveness if running in another process).
97
self.tokenizer = init_tokenizer_from_configs(
@@ -118,7 +135,7 @@ class AsyncLLM(EngineClient):
135
vllm_config: VllmConfig,
136
start_engine_loop: bool = True,
137
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
121
- stat_loggers: Optional[dict[str, StatLoggerBase]] = None,
138
+ stat_loggers: Optional[list[StatLoggerFactory]] = None,
139
disable_log_requests: bool = False,
140
disable_log_stats: bool = False,
141
) -> "AsyncLLM":
@@ -129,17 +146,12 @@ class AsyncLLM(EngineClient):
146
"AsyncLLMEngine.from_vllm_config(...) or explicitly set "
147
"VLLM_USE_V1=0 or 1 and report this issue on Github.")
148
132
- # FIXME(rob): refactor VllmConfig to include the StatLoggers
133
- # include StatLogger in the Oracle decision.
134
- if stat_loggers is not None:
135
- raise ValueError("Custom StatLoggers are not yet supported on V1. "
136
- "Explicitly set VLLM_USE_V1=0 to disable V1.")
137
-
149
# Create the LLMEngine.
150
return cls(
151
vllm_config=vllm_config,
152
executor_class=Executor.get_class(vllm_config),
153
start_engine_loop=start_engine_loop,
154
+ stat_loggers=stat_loggers,
155
log_requests=not disable_log_requests,
156
log_stats=not disable_log_stats,
157
usage_context=usage_context,
@@ -151,6 +163,7 @@ class AsyncLLM(EngineClient):
163
engine_args: AsyncEngineArgs,
164
start_engine_loop: bool = True,
165
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
166
+ stat_loggers: Optional[list[StatLoggerFactory]] = None,
167
) -> "AsyncLLM":
168
"""Create an AsyncLLM from the EngineArgs."""
169
@@ -166,6 +179,7 @@ class AsyncLLM(EngineClient):
179
log_stats=not engine_args.disable_log_stats,
180
start_engine_loop=start_engine_loop,
181
usage_context=usage_context,
182
+ stat_loggers=stat_loggers,
183
)
184
185
def __del__(self):
vllm/v1/engine/llm_engine.py
+9
-9
@@ -10,7 +10,6 @@ import vllm.envs as envs
10
from vllm.config import ParallelConfig, VllmConfig
11
from vllm.distributed import stateless_destroy_torch_distributed_process_group
12
from vllm.engine.arg_utils import EngineArgs
13
-from vllm.engine.metrics_types import StatLoggerBase
13
from vllm.inputs import PromptType
14
from vllm.logger import init_logger
15
from vllm.lora.request import LoRARequest
@@ -28,6 +27,7 @@ from vllm.v1.engine.output_processor import OutputProcessor
27
from vllm.v1.engine.parallel_sampling import ParentRequest
28
from vllm.v1.engine.processor import Processor
29
from vllm.v1.executor.abstract import Executor
30
+from vllm.v1.metrics.loggers import StatLoggerFactory
31
32
logger = init_logger(__name__)
33
@@ -43,7 +43,7 @@ class LLMEngine:
43
executor_class: type[Executor],
44
log_stats: bool,
45
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
46
- stat_loggers: Optional[dict[str, StatLoggerBase]] = None,
46
+ stat_loggers: Optional[list[StatLoggerFactory]] = None,
47
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
48
use_cached_outputs: bool = False,
49
multiprocess_mode: bool = False,
@@ -55,6 +55,11 @@ class LLMEngine:
55
"LLMEngine.from_vllm_config(...) or explicitly set "
56
"VLLM_USE_V1=0 or 1 and report this issue on Github.")
57
58
+ if stat_loggers is not None:
59
+ raise NotImplementedError(
60
+ "Passing StatLoggers to LLMEngine in V1 is not yet supported. "
61
+ "Set VLLM_USE_V1=0 and file and issue on Github.")
62
+
63
self.vllm_config = vllm_config
64
self.model_config = vllm_config.model_config
65
self.cache_config = vllm_config.cache_config
@@ -101,14 +106,9 @@ class LLMEngine:
106
cls,
107
vllm_config: VllmConfig,
108
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
104
- stat_loggers: Optional[dict[str, StatLoggerBase]] = None,
109
+ stat_loggers: Optional[list[StatLoggerFactory]] = None,
110
disable_log_stats: bool = False,
111
) -> "LLMEngine":
107
- if stat_loggers is not None:
108
- raise NotImplementedError(
109
- "Passing StatLoggers to V1 is not yet supported. "
110
- "Set VLLM_USE_V1=0 and file and issue on Github.")
111
-
112
return cls(vllm_config=vllm_config,
113
executor_class=Executor.get_class(vllm_config),
114
log_stats=(not disable_log_stats),
@@ -121,7 +121,7 @@ class LLMEngine:
121
cls,
122
engine_args: EngineArgs,
123
usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
124
- stat_loggers: Optional[dict[str, StatLoggerBase]] = None,
124
+ stat_loggers: Optional[list[StatLoggerFactory]] = None,
125
enable_multiprocessing: bool = False,
126
) -> "LLMEngine":
127
"""Creates an LLM engine from the engine arguments."""
vllm/v1/metrics/loggers.py
+43
-2
@@ -1,8 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import logging
4
import time
5
from abc import ABC, abstractmethod
5
-from typing import Optional
6
+from typing import Callable, Optional
7
8
import numpy as np
9
import prometheus_client
@@ -18,8 +19,20 @@ logger = init_logger(__name__)
19
20
_LOCAL_LOGGING_INTERVAL_SEC = 5.0
21
22
+StatLoggerFactory = Callable[[VllmConfig, int], "StatLoggerBase"]
23
+
24
25
class StatLoggerBase(ABC):
26
+ """Interface for logging metrics.
27
+
28
+ API users may define custom loggers that implement this interface.
29
+ However, note that the `SchedulerStats` and `IterationStats` classes
30
+ are not considered stable interfaces and may change in future versions.
31
+ """
32
+
33
+ @abstractmethod
34
+ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
35
+ ...
36
37
@abstractmethod
38
def record(self, scheduler_stats: SchedulerStats,
@@ -32,7 +45,7 @@ class StatLoggerBase(ABC):
45
46
class LoggingStatLogger(StatLoggerBase):
47
35
- def __init__(self, engine_index: int = 0):
48
+ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
49
self.engine_index = engine_index
50
self._reset(time.monotonic())
51
self.last_scheduler_stats = SchedulerStats()
@@ -462,3 +475,31 @@ def build_cudagraph_buckets(vllm_config: VllmConfig) -> list[int]:
475
return buckets
476
else:
477
return [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8096]
478
+
479
+
480
+def setup_default_loggers(
481
+ vllm_config: VllmConfig,
482
+ log_stats: bool,
483
+ engine_num: int,
484
+ custom_stat_loggers: Optional[list[StatLoggerFactory]] = None,
485
+) -> list[list[StatLoggerBase]]:
486
+ """Setup logging and prometheus metrics."""
487
+ if not log_stats:
488
+ return []
489
+
490
+ factories: list[StatLoggerFactory]
491
+ if custom_stat_loggers is not None:
492
+ factories = custom_stat_loggers
493
+ else:
494
+ factories = [PrometheusStatLogger]
495
+ if logger.isEnabledFor(logging.INFO):
496
+ factories.append(LoggingStatLogger)
497
+
498
+ stat_loggers: list[list[StatLoggerBase]] = []
499
+ for i in range(engine_num):
500
+ per_engine_stat_loggers: list[StatLoggerBase] = []
501
+ for logger_factory in factories:
502
+ per_engine_stat_loggers.append(logger_factory(vllm_config, i))
503
+ stat_loggers.append(per_engine_stat_loggers)
504
+
505
+ return stat_loggers