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 14, 05:30 UTC
[V1][Metrics][Plugin] Add plugin support for custom `StatLoggerBase` implementations (#22456)
Signed-off-by: tovam <tovam@pliops.com>
Tova Movshovitz committed
Oct 19, 2025 at 01:12 UTC
83e760c57ddbfcaf2004eb96d32f06dca3f50e2f
8 files changed
+164
-30
.buildkite/test-pipeline.yaml
+5
@@ -1020,6 +1020,11 @@ steps:
1020
- pytest -v -s plugins_tests/test_io_processor_plugins.py
1021
- pip uninstall prithvi_io_processor_plugin -y
1022
# end io_processor plugins test
1023
+ # begin stat_logger plugins test
1024
+ - pip install -e ./plugins/vllm_add_dummy_stat_logger
1025
+ - pytest -v -s plugins_tests/test_stats_logger_plugins.py
1026
+ - pip uninstall dummy_stat_logger -y
1027
+ # end stat_logger plugins test
1028
# other tests continue here:
1029
- pytest -v -s plugins_tests/test_scheduler_plugins.py
1030
- pip install -e ./plugins/vllm_add_dummy_model
docs/design/plugin_system.md
+3
-1
@@ -41,7 +41,7 @@ Every plugin has three parts:
41
42
1. **Plugin group**: The name of the entry point group. vLLM uses the entry point group `vllm.general_plugins` to register general plugins. This is the key of `entry_points` in the `setup.py` file. Always use `vllm.general_plugins` for vLLM's general plugins.
43
2. **Plugin name**: The name of the plugin. This is the value in the dictionary of the `entry_points` dictionary. In the example above, the plugin name is `register_dummy_model`. Plugins can be filtered by their names using the `VLLM_PLUGINS` environment variable. To load only a specific plugin, set `VLLM_PLUGINS` to the plugin name.
44
-3. **Plugin value**: The fully qualified name of the function to register in the plugin system. In the example above, the plugin value is `vllm_add_dummy_model:register`, which refers to a function named `register` in the `vllm_add_dummy_model` module.
44
+3. **Plugin value**: The fully qualified name of the function or module to register in the plugin system. In the example above, the plugin value is `vllm_add_dummy_model:register`, which refers to a function named `register` in the `vllm_add_dummy_model` module.
45
46
## Types of supported plugins
47
@@ -51,6 +51,8 @@ Every plugin has three parts:
51
52
- **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre/post processing of the model prompt and model output for pooling models. The plugin function returns the IOProcessor's class fully qualified name.
53
54
+- **Stat logger plugins** (with group name `vllm.stat_logger_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree loggers into vLLM. The entry point should be a class that subclasses StatLoggerBase.
55
+
56
## Guidelines for Writing Plugins
57
58
- **Being re-entrant**: The function specified in the entry point should be re-entrant, meaning it can be called multiple times without causing issues. This is necessary because the function might be called multiple times in some processes.
tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py
new
+29
@@ -0,0 +1,29 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.v1.metrics.loggers import StatLoggerBase
5
+
6
+
7
+class DummyStatLogger(StatLoggerBase):
8
+ """
9
+ A dummy stat logger for testing purposes.
10
+ Implements the minimal interface expected by StatLoggerManager.
11
+ """
12
+
13
+ def __init__(self, vllm_config, engine_idx=0):
14
+ self.vllm_config = vllm_config
15
+ self.engine_idx = engine_idx
16
+ self.recorded = []
17
+ self.logged = False
18
+ self.engine_initialized = False
19
+
20
+ def record(self, scheduler_stats, iteration_stats, mm_cache_stats, engine_idx):
21
+ self.recorded.append(
22
+ (scheduler_stats, iteration_stats, mm_cache_stats, engine_idx)
23
+ )
24
+
25
+ def log(self):
26
+ self.logged = True
27
+
28
+ def log_engine_initialized(self):
29
+ self.engine_initialized = True
tests/plugins/vllm_add_dummy_stat_logger/setup.py
new
+15
@@ -0,0 +1,15 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from setuptools import setup
5
+
6
+setup(
7
+ name="dummy_stat_logger",
8
+ version="0.1",
9
+ packages=["dummy_stat_logger"],
10
+ entry_points={
11
+ "vllm.stat_logger_plugins": [
12
+ "dummy_stat_logger = dummy_stat_logger.dummy_stat_logger:DummyStatLogger" # noqa
13
+ ]
14
+ },
15
+)
tests/plugins_tests/test_stats_logger_plugins.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+from dummy_stat_logger.dummy_stat_logger import DummyStatLogger
6
+
7
+from vllm.config import VllmConfig
8
+from vllm.engine.arg_utils import AsyncEngineArgs
9
+from vllm.v1.engine.async_llm import AsyncLLM
10
+from vllm.v1.metrics.loggers import load_stat_logger_plugin_factories
11
+
12
+
13
+def test_stat_logger_plugin_is_discovered(monkeypatch: pytest.MonkeyPatch):
14
+ with monkeypatch.context() as m:
15
+ m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
16
+
17
+ factories = load_stat_logger_plugin_factories()
18
+ assert len(factories) == 1, f"Expected 1 factory, got {len(factories)}"
19
+ assert factories[0] is DummyStatLogger, (
20
+ f"Expected DummyStatLogger class, got {factories[0]}"
21
+ )
22
+
23
+ # instantiate and confirm the right type
24
+ vllm_config = VllmConfig()
25
+ instance = factories[0](vllm_config)
26
+ assert isinstance(instance, DummyStatLogger)
27
+
28
+
29
+def test_no_plugins_loaded_if_env_empty(monkeypatch: pytest.MonkeyPatch):
30
+ with monkeypatch.context() as m:
31
+ m.setenv("VLLM_PLUGINS", "")
32
+
33
+ factories = load_stat_logger_plugin_factories()
34
+ assert factories == []
35
+
36
+
37
+def test_invalid_stat_logger_plugin_raises(monkeypatch: pytest.MonkeyPatch):
38
+ def fake_plugin_loader(group: str):
39
+ assert group == "vllm.stat_logger_plugins"
40
+ return {"bad": object()}
41
+
42
+ with monkeypatch.context() as m:
43
+ m.setattr(
44
+ "vllm.v1.metrics.loggers.load_plugins_by_group",
45
+ fake_plugin_loader,
46
+ )
47
+ with pytest.raises(
48
+ TypeError,
49
+ match="Stat logger plugin 'bad' must be a subclass of StatLoggerBase",
50
+ ):
51
+ load_stat_logger_plugin_factories()
52
+
53
+
54
+@pytest.mark.asyncio
55
+async def test_stat_logger_plugin_integration_with_engine(
56
+ monkeypatch: pytest.MonkeyPatch,
57
+):
58
+ with monkeypatch.context() as m:
59
+ m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
60
+
61
+ engine_args = AsyncEngineArgs(
62
+ model="facebook/opt-125m",
63
+ enforce_eager=True, # reduce test time
64
+ disable_log_stats=True, # disable default loggers
65
+ )
66
+
67
+ engine = AsyncLLM.from_engine_args(engine_args=engine_args)
68
+
69
+ assert len(engine.logger_manager.stat_loggers) == 2
70
+ assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1
71
+ assert isinstance(
72
+ engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0],
73
+ DummyStatLogger,
74
+ )
75
+
76
+ engine.shutdown()
tests/v1/metrics/test_engine_logger_apis.py
+3
-23
@@ -4,33 +4,13 @@ import copy
4
5
import pytest
6
7
+from tests.plugins.vllm_add_dummy_stat_logger.dummy_stat_logger.dummy_stat_logger import ( # noqa E501
8
+ DummyStatLogger,
9
+)
10
from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM
11
from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
12
13
11
-class DummyStatLogger:
12
- """
13
- A dummy stat logger for testing purposes.
14
- Implements the minimal interface expected by StatLoggerManager.
15
- """
16
-
17
- def __init__(self, vllm_config, engine_idx):
18
- self.vllm_config = vllm_config
19
- self.engine_idx = engine_idx
20
- self.recorded = []
21
- self.logged = False
22
- self.engine_initialized = False
23
-
24
- def record(self, scheduler_stats, iteration_stats, engine_idx):
25
- self.recorded.append((scheduler_stats, iteration_stats, engine_idx))
26
-
27
- def log(self):
28
- self.logged = True
29
-
30
- def log_engine_initialized(self):
31
- self.engine_initialized = True
32
-
33
-
14
@pytest.fixture
15
def log_stats_enabled_engine_args():
16
"""
vllm/v1/engine/async_llm.py
+15
-6
@@ -40,7 +40,11 @@ from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollec
40
from vllm.v1.engine.parallel_sampling import ParentRequest
41
from vllm.v1.engine.processor import Processor
42
from vllm.v1.executor.abstract import Executor
43
-from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
43
+from vllm.v1.metrics.loggers import (
44
+ StatLoggerFactory,
45
+ StatLoggerManager,
46
+ load_stat_logger_plugin_factories,
47
+)
48
from vllm.v1.metrics.prometheus import shutdown_prometheus
49
from vllm.v1.metrics.stats import IterationStats
50
@@ -100,11 +104,16 @@ class AsyncLLM(EngineClient):
104
self.observability_config = vllm_config.observability_config
105
self.log_requests = log_requests
106
103
- self.log_stats = log_stats or (stat_loggers is not None)
104
- if not log_stats and stat_loggers is not None:
107
+ custom_stat_loggers = list(stat_loggers or [])
108
+ custom_stat_loggers.extend(load_stat_logger_plugin_factories())
109
+
110
+ has_custom_loggers = bool(custom_stat_loggers)
111
+ self.log_stats = log_stats or has_custom_loggers
112
+ if not log_stats and has_custom_loggers:
113
logger.info(
106
- "AsyncLLM created with log_stats=False and non-empty custom "
107
- "logger list; enabling logging without default stat loggers"
114
+ "AsyncLLM created with log_stats=False, "
115
+ "but custom stat loggers were found; "
116
+ "enabling logging without default stat loggers."
117
)
118
119
if self.model_config.skip_tokenizer_init:
@@ -144,7 +153,7 @@ class AsyncLLM(EngineClient):
153
self.logger_manager = StatLoggerManager(
154
vllm_config=vllm_config,
155
engine_idxs=self.engine_core.engine_ranks_managed,
147
- custom_stat_loggers=stat_loggers,
156
+ custom_stat_loggers=custom_stat_loggers,
157
enable_default_loggers=log_stats,
158
client_count=client_count,
159
aggregate_engine_logging=aggregate_engine_logging,
vllm/v1/metrics/loggers.py
+18
@@ -12,6 +12,7 @@ from prometheus_client import Counter, Gauge, Histogram
12
from vllm.config import SupportsMetricsInfo, VllmConfig
13
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorLogging
14
from vllm.logger import init_logger
15
+from vllm.plugins import load_plugins_by_group
16
from vllm.v1.engine import FinishReason
17
from vllm.v1.metrics.prometheus import unregister_vllm_metrics
18
from vllm.v1.metrics.stats import (
@@ -56,6 +57,23 @@ class StatLoggerBase(ABC):
57
pass
58
59
60
+def load_stat_logger_plugin_factories() -> list[StatLoggerFactory]:
61
+ factories: list[StatLoggerFactory] = []
62
+
63
+ for name, plugin_class in load_plugins_by_group("vllm.stat_logger_plugins").items():
64
+ if not isinstance(plugin_class, type) or not issubclass(
65
+ plugin_class, StatLoggerBase
66
+ ):
67
+ raise TypeError(
68
+ f"Stat logger plugin {name!r} must be a subclass of "
69
+ f"StatLoggerBase (got {plugin_class!r})."
70
+ )
71
+
72
+ factories.append(plugin_class)
73
+
74
+ return factories
75
+
76
+
77
class AggregateStatLoggerBase(StatLoggerBase):
78
"""Abstract base class for loggers that
79
aggregate across multiple DP engines."""