main
py 78 lines 2.71 KB
Raw
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 # The `is` check above narrows DummyStatLogger to a Callable
27 assert isinstance(instance, DummyStatLogger)
28
29
30 def test_no_plugins_loaded_if_env_empty(monkeypatch: pytest.MonkeyPatch):
31 with monkeypatch.context() as m:
32 m.setenv("VLLM_PLUGINS", "")
33
34 factories = load_stat_logger_plugin_factories()
35 assert factories == []
36
37
38 def test_invalid_stat_logger_plugin_raises(monkeypatch: pytest.MonkeyPatch):
39 def fake_plugin_loader(group: str):
40 assert group == "vllm.stat_logger_plugins"
41 return {"bad": object()}
42
43 with monkeypatch.context() as m:
44 m.setattr(
45 "vllm.v1.metrics.loggers.load_plugins_by_group",
46 fake_plugin_loader,
47 )
48 with pytest.raises(
49 TypeError,
50 match="Stat logger plugin 'bad' must be a subclass of StatLoggerBase",
51 ):
52 load_stat_logger_plugin_factories()
53
54
55 @pytest.mark.asyncio
56 async def test_stat_logger_plugin_integration_with_engine(
57 monkeypatch: pytest.MonkeyPatch,
58 ):
59 with monkeypatch.context() as m:
60 m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
61
62 engine_args = AsyncEngineArgs(
63 model="facebook/opt-125m",
64 enforce_eager=True, # reduce test time
65 disable_log_stats=True, # disable default loggers
66 )
67
68 engine = AsyncLLM.from_engine_args(engine_args=engine_args)
69
70 assert engine.logger_manager is not None
71 assert len(engine.logger_manager.stat_loggers) == 2
72 assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1 # type: ignore[attr-defined]
73 assert isinstance(
74 engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0], # type: ignore[attr-defined]
75 DummyStatLogger,
76 )
77
78 engine.shutdown()