main
py 236 lines 7.89 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the `vllm.endpoint_plugins` framework (RFC #46565).
4
5 Uses the worked in repo example plugin (`vllm_add_dummy_endpoint_plugin`,
6 installed via `tests/plugins/vllm_add_dummy_endpoint_plugin`) exercising both
7 `EndpointPlugin` hooks against a fake `EngineClient`, unit tests for the
8 `load_endpoint_plugins` gating matrix and an e2e test that drives a real HTTP
9 request through the plugin's route.
10 """
11
12 from argparse import Namespace
13 from typing import Any
14
15 import httpx
16 import pytest
17 from fastapi import FastAPI
18 from vllm_add_dummy_endpoint_plugin import DummyAdminEndpointPlugin
19
20 from vllm.entrypoints.launchers.app import build_app
21 from vllm.entrypoints.launchers.cli_args import make_arg_parser
22 from vllm.plugins import load_endpoint_plugins
23 from vllm.plugins.endpoint_plugins.interface import (
24 EndpointPlugin,
25 attach_endpoint_plugins,
26 init_endpoint_plugins_state,
27 )
28 from vllm.utils.argparse_utils import FlexibleArgumentParser
29
30
31 class _RaisingEndpointPlugin:
32 """Factory that raises to exercise the "instantiation fails" path."""
33
34 name = "raising_endpoint_plugin"
35 required_tasks = None
36
37 def __init__(self):
38 raise RuntimeError("boom")
39
40
41 class _FakeEngineClient:
42 """Minimal stand in exercising `collective_rpc`. Not a real engine."""
43
44 def __init__(self, rpc_result: Any = None):
45 self.rpc_result = rpc_result
46 self.rpc_calls: list[tuple[str, tuple, dict]] = []
47
48 async def collective_rpc(self, method, timeout=None, args=(), kwargs=None):
49 self.rpc_calls.append((method, args, kwargs or {}))
50 return self.rpc_result
51
52
53 def _build_args() -> Namespace:
54 parser = FlexibleArgumentParser()
55 subparsers = parser.add_subparsers()
56 serve_parser = subparsers.add_parser("serve")
57 make_arg_parser(serve_parser)
58 return serve_parser.parse_args([])
59
60
61 def _fake_loader(factories: dict[str, Any]):
62 def _load_plugins_by_group(group: str) -> dict[str, Any]:
63 assert group == "vllm.endpoint_plugins"
64 return factories
65
66 return _load_plugins_by_group
67
68
69 def test_dummy_plugin_satisfies_protocol():
70 assert isinstance(DummyAdminEndpointPlugin(), EndpointPlugin)
71
72
73 def test_no_plugins_loaded_when_allowlist_unset(monkeypatch: pytest.MonkeyPatch):
74 monkeypatch.delenv("VLLM_PLUGINS", raising=False)
75
76 assert load_endpoint_plugins(("generate",)) == []
77
78
79 def test_no_plugins_loaded_when_allowlist_is_empty_string(
80 monkeypatch: pytest.MonkeyPatch,
81 ):
82 """`VLLM_PLUGINS=""` parses to `[""]`, not `None` (see `vllm.envs`), so it
83 must be treated as a (non strict) allowlist matching no plugin name, not
84 as "unset"."""
85 monkeypatch.setenv("VLLM_PLUGINS", "")
86
87 assert load_endpoint_plugins(("generate",)) == []
88
89
90 def test_plugin_loaded_when_allowlisted_and_task_matches(
91 monkeypatch: pytest.MonkeyPatch,
92 ):
93 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
94
95 plugins = load_endpoint_plugins(("generate",))
96
97 assert len(plugins) == 1
98 assert isinstance(plugins[0], DummyAdminEndpointPlugin)
99
100
101 def test_plugin_skipped_when_required_tasks_miss(monkeypatch: pytest.MonkeyPatch):
102 class _GenerateOnlyPlugin(DummyAdminEndpointPlugin):
103 required_tasks = ("generate",)
104
105 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
106 monkeypatch.setattr(
107 "vllm.plugins.load_plugins_by_group",
108 _fake_loader({"dummy_admin_endpoint_plugin": _GenerateOnlyPlugin}),
109 )
110
111 assert load_endpoint_plugins(("embed",)) == []
112 assert len(load_endpoint_plugins(("generate",))) == 1
113
114
115 def test_plugin_loaded_when_required_tasks_is_none(monkeypatch: pytest.MonkeyPatch):
116 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
117
118 assert len(load_endpoint_plugins(supported_tasks=None)) == 1
119
120
121 def test_plugin_skipped_when_required_tasks_set_but_supported_tasks_none(
122 monkeypatch: pytest.MonkeyPatch,
123 ):
124 class _GenerateOnlyPlugin(DummyAdminEndpointPlugin):
125 required_tasks = ("generate",)
126
127 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
128 monkeypatch.setattr(
129 "vllm.plugins.load_plugins_by_group",
130 _fake_loader({"dummy_admin_endpoint_plugin": _GenerateOnlyPlugin}),
131 )
132
133 assert load_endpoint_plugins(supported_tasks=None) == []
134
135
136 def test_factory_raising_is_logged_and_skipped(monkeypatch: pytest.MonkeyPatch):
137 monkeypatch.setenv(
138 "VLLM_PLUGINS", "raising_endpoint_plugin,dummy_admin_endpoint_plugin"
139 )
140 monkeypatch.setattr(
141 "vllm.plugins.load_plugins_by_group",
142 _fake_loader(
143 {
144 "raising_endpoint_plugin": _RaisingEndpointPlugin,
145 "dummy_admin_endpoint_plugin": DummyAdminEndpointPlugin,
146 }
147 ),
148 )
149
150 plugins = load_endpoint_plugins(("generate",))
151
152 assert len(plugins) == 1
153 assert isinstance(plugins[0], DummyAdminEndpointPlugin)
154
155
156 def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch):
157 monkeypatch.delenv("VLLM_PLUGINS", raising=False)
158
159 app = FastAPI()
160 attach_endpoint_plugins(app, ("generate",))
161
162 assert app.state.endpoint_plugins == []
163
164
165 @pytest.mark.asyncio
166 async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatch):
167 """`init_app_state` callers that never ran `build_app` (e.g.
168 `run_batch.py`, which builds a bare `State()`) must not crash just
169 because `state.endpoint_plugins` was never set."""
170 from starlette.datastructures import State
171
172 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
173
174 state = State()
175 await init_endpoint_plugins_state(_FakeEngineClient(), state, _build_args())
176
177 assert not hasattr(state, "dummy_engine_client")
178
179
180 @pytest.mark.asyncio
181 async def test_render_server_attaches_endpoint_plugins_with_no_engine_client(
182 monkeypatch: pytest.MonkeyPatch,
183 ):
184 """The CPU only render server has no `EngineClient` but a plugin eligible
185 for the `render` task (`required_tasks` is `None` or includes `"render"`)
186 still gets its routes attached at Phase A. Phase B passes `None` for
187 `engine_client` and it's up to the plugin to handle that."""
188 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
189
190 args = _build_args()
191 app = build_app(args, ("render",))
192
193 assert len(app.state.endpoint_plugins) == 1
194 assert any(
195 getattr(route, "path", None) == "/v1/admin/scheduler_config"
196 for route in app.routes
197 )
198
199 await init_endpoint_plugins_state(None, app.state, args)
200
201 assert app.state.dummy_engine_client is None
202
203 transport = httpx.ASGITransport(app=app)
204 async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
205 response = await client.get("/v1/admin/scheduler_config")
206
207 assert response.status_code == 503
208
209
210 @pytest.mark.asyncio
211 async def test_endpoint_plugin_end_to_end(monkeypatch: pytest.MonkeyPatch):
212 """Phase A (attach) + Phase B (init) wired through `build_app` then
213 exercised with a real HTTP request against the worked example plugin."""
214 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")
215
216 args = _build_args()
217 app = build_app(args, supported_tasks=())
218
219 assert len(app.state.endpoint_plugins) == 1
220 assert any(
221 getattr(route, "path", None) == "/v1/admin/scheduler_config"
222 for route in app.routes
223 )
224
225 fake_engine_client = _FakeEngineClient(rpc_result=["cfg-a", "cfg-b"])
226 await init_endpoint_plugins_state(fake_engine_client, app.state, args)
227
228 assert app.state.dummy_engine_client is fake_engine_client
229
230 transport = httpx.ASGITransport(app=app)
231 async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
232 response = await client.get("/v1/admin/scheduler_config")
233
234 assert response.status_code == 200
235 assert response.json() == {"scheduler_config": ["cfg-a", "cfg-b"]}
236 assert fake_engine_client.rpc_calls == [("get_scheduler_config", (), {})]