main
py 256 lines 8.34 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from contextlib import suppress
5 from dataclasses import dataclass, field
6 from http import HTTPStatus
7 from unittest.mock import AsyncMock, MagicMock
8
9 import pytest
10
11 from vllm.config.multimodal import MultiModalConfig
12 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
13 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
14 from vllm.entrypoints.openai.models.protocol import BaseModelPath
15 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
16 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
17 from vllm.lora.request import LoRARequest
18 from vllm.lora.resolver import LoRAResolver, LoRAResolverRegistry
19 from vllm.renderers.hf import HfRenderer
20 from vllm.renderers.online_renderer import OnlineRenderer
21 from vllm.tokenizers.registry import cached_tokenizer_from_config
22 from vllm.v1.engine.async_llm import AsyncLLM
23
24 MODEL_NAME = "openai-community/gpt2"
25 BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
26
27 MOCK_RESOLVER_NAME = "mock_test_resolver"
28
29
30 @dataclass
31 class MockHFConfig:
32 model_type: str = "any"
33
34
35 @dataclass
36 class MockModelConfig:
37 """Minimal mock ModelConfig for testing."""
38
39 model: str = MODEL_NAME
40 runner_type = "generate"
41 tokenizer: str = MODEL_NAME
42 trust_remote_code: bool = False
43 tokenizer_mode: str = "auto"
44 max_model_len: int = 100
45 tokenizer_revision: str | None = None
46 multimodal_config: MultiModalConfig = field(default_factory=MultiModalConfig)
47 hf_config: MockHFConfig = field(default_factory=MockHFConfig)
48 logits_processors: list[str] | None = None
49 diff_sampling_param: dict | None = None
50 allowed_local_media_path: str = ""
51 allowed_media_domains: list[str] | None = None
52 encoder_config = None
53 generation_config: str = "auto"
54 skip_tokenizer_init: bool = False
55 is_encoder_decoder: bool = False
56 is_multimodal_model: bool = False
57 supports_multimodal_inputs: bool = False
58 renderer_num_workers: int = 1
59
60 def get_diff_sampling_param(self):
61 return self.diff_sampling_param or {}
62
63
64 @dataclass
65 class MockParallelConfig:
66 _api_process_rank: int = 0
67
68
69 @dataclass
70 class MockVllmConfig:
71 model_config: MockModelConfig
72 parallel_config: MockParallelConfig
73
74
75 class MockLoRAResolver(LoRAResolver):
76 async def resolve_lora(
77 self, base_model_name: str, lora_name: str
78 ) -> LoRARequest | None:
79 if lora_name == "test-lora":
80 return LoRARequest(
81 lora_name="test-lora",
82 lora_int_id=1,
83 lora_path="/fake/path/test-lora",
84 )
85 elif lora_name == "invalid-lora":
86 return LoRARequest(
87 lora_name="invalid-lora",
88 lora_int_id=2,
89 lora_path="/fake/path/invalid-lora",
90 )
91 return None
92
93
94 @pytest.fixture(autouse=True)
95 def register_mock_resolver():
96 """Fixture to register and unregister the mock LoRA resolver."""
97 resolver = MockLoRAResolver()
98 LoRAResolverRegistry.register_resolver(MOCK_RESOLVER_NAME, resolver)
99 yield
100 # Cleanup: remove the resolver after the test runs
101 if MOCK_RESOLVER_NAME in LoRAResolverRegistry.resolvers:
102 del LoRAResolverRegistry.resolvers[MOCK_RESOLVER_NAME]
103
104
105 def _build_renderer(model_config: MockModelConfig):
106 return HfRenderer(
107 MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
108 cached_tokenizer_from_config(model_config),
109 )
110
111
112 @pytest.fixture
113 def mock_serving_setup():
114 """Provides a mocked engine and serving completion instance."""
115 mock_engine = MagicMock(spec=AsyncLLM)
116 mock_engine.errored = False
117
118 async def mock_add_lora_side_effect(lora_request: LoRARequest):
119 """Simulate engine behavior when adding LoRAs."""
120 if lora_request.lora_name == "test-lora":
121 # Simulate successful addition
122 return True
123 if lora_request.lora_name == "invalid-lora":
124 # Simulate failure during addition (e.g. invalid format)
125 raise ValueError(f"Simulated failure adding LoRA: {lora_request.lora_name}")
126 return True
127
128 mock_engine.add_lora = AsyncMock(side_effect=mock_add_lora_side_effect)
129
130 async def mock_generate(*args, **kwargs):
131 for _ in []:
132 yield _
133
134 mock_engine.generate = MagicMock(spec=AsyncLLM.generate, side_effect=mock_generate)
135
136 mock_engine.generate.reset_mock()
137 mock_engine.add_lora.reset_mock()
138
139 mock_engine.model_config = MockModelConfig()
140 mock_engine.input_processor = MagicMock()
141 mock_engine.renderer = _build_renderer(mock_engine.model_config)
142
143 models = OpenAIServingModels(
144 engine_client=mock_engine,
145 base_model_paths=BASE_MODEL_PATHS,
146 )
147
148 online_renderer = OnlineRenderer(
149 model_config=mock_engine.model_config,
150 renderer=mock_engine.renderer,
151 request_logger=None,
152 chat_template=None,
153 chat_template_content_format="auto",
154 )
155 serving_completion = OpenAIServingCompletion(
156 mock_engine, models, online_renderer=online_renderer, request_logger=None
157 )
158
159 return mock_engine, serving_completion
160
161
162 @pytest.mark.asyncio
163 async def test_serving_completion_with_lora_resolver(mock_serving_setup, monkeypatch):
164 monkeypatch.setenv("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "true")
165
166 mock_engine, serving_completion = mock_serving_setup
167
168 lora_model_name = "test-lora"
169 req_found = CompletionRequest(
170 model=lora_model_name,
171 prompt="Generate with LoRA",
172 )
173
174 # Suppress potential errors during the mocked generate call,
175 # as we are primarily checking for add_lora and generate calls
176 with suppress(Exception):
177 await serving_completion.create_completion(req_found)
178
179 mock_engine.add_lora.assert_awaited_once()
180 called_lora_request = mock_engine.add_lora.call_args[0][0]
181 assert isinstance(called_lora_request, LoRARequest)
182 assert called_lora_request.lora_name == lora_model_name
183
184 mock_engine.generate.assert_called_once()
185 called_lora_request = mock_engine.generate.call_args[1]["lora_request"]
186 assert isinstance(called_lora_request, LoRARequest)
187 assert called_lora_request.lora_name == lora_model_name
188
189
190 @pytest.mark.asyncio
191 async def test_serving_completion_resolver_not_found(mock_serving_setup, monkeypatch):
192 monkeypatch.setenv("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "true")
193
194 mock_engine, serving_completion = mock_serving_setup
195
196 non_existent_model = "non-existent-lora-adapter"
197 req = CompletionRequest(
198 model=non_existent_model,
199 prompt="what is 1+1?",
200 )
201
202 response = await serving_completion.create_completion(req)
203
204 mock_engine.add_lora.assert_not_awaited()
205 mock_engine.generate.assert_not_called()
206
207 assert isinstance(response, ErrorResponse)
208 assert response.error.code == HTTPStatus.NOT_FOUND.value
209 assert non_existent_model in response.error.message
210
211
212 @pytest.mark.asyncio
213 async def test_serving_completion_resolver_add_lora_fails(
214 mock_serving_setup, monkeypatch
215 ):
216 monkeypatch.setenv("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "true")
217
218 mock_engine, serving_completion = mock_serving_setup
219
220 invalid_model = "invalid-lora"
221 req = CompletionRequest(
222 model=invalid_model,
223 prompt="what is 1+1?",
224 )
225
226 response = await serving_completion.create_completion(req)
227
228 # Assert add_lora was called before the failure
229 mock_engine.add_lora.assert_awaited_once()
230 called_lora_request = mock_engine.add_lora.call_args[0][0]
231 assert isinstance(called_lora_request, LoRARequest)
232 assert called_lora_request.lora_name == invalid_model
233
234 # Assert generate was *not* called due to the failure
235 mock_engine.generate.assert_not_called()
236
237 # Assert the correct error response
238 assert isinstance(response, ErrorResponse)
239 assert response.error.code == HTTPStatus.BAD_REQUEST.value
240 assert invalid_model in response.error.message
241
242
243 @pytest.mark.asyncio
244 async def test_serving_completion_flag_not_set(mock_serving_setup):
245 mock_engine, serving_completion = mock_serving_setup
246
247 lora_model_name = "test-lora"
248 req_found = CompletionRequest(
249 model=lora_model_name,
250 prompt="Generate with LoRA",
251 )
252
253 await serving_completion.create_completion(req_found)
254
255 mock_engine.add_lora.assert_not_called()
256 mock_engine.generate.assert_not_called()