main
py 218 lines 8.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from http import HTTPStatus
5 from unittest.mock import MagicMock
6
7 import pytest
8
9 from vllm import PoolingParams
10 from vllm.config import ModelConfig
11 from vllm.engine.protocol import EngineClient
12 from vllm.entrypoints.openai.models.protocol import BaseModelPath
13 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
14 from vllm.entrypoints.pooling.base.serving import PoolingBaseServing
15 from vllm.entrypoints.pooling.typing import PoolingServeContext
16 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
17 from vllm.entrypoints.serve.lora.protocol import (
18 LoadLoRAAdapterRequest,
19 UnloadLoRAAdapterRequest,
20 )
21 from vllm.exceptions import VLLMNotFoundError
22 from vllm.lora.request import LoRARequest
23
24 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
25 BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
26 LORA_LOADING_SUCCESS_MESSAGE = "Success: LoRA adapter '{lora_name}' added successfully."
27 LORA_UNLOADING_SUCCESS_MESSAGE = (
28 "Success: LoRA adapter '{lora_name}' removed successfully."
29 )
30
31
32 async def _async_serving_models_init_with_mock() -> tuple[
33 OpenAIServingModels, MagicMock
34 ]:
35 mock_engine_client = MagicMock(spec=EngineClient)
36 # Set the max_model_len attribute to avoid missing attribute
37 mock_model_config = MagicMock(spec=ModelConfig)
38 mock_model_config.max_model_len = 2048
39 mock_engine_client.model_config = mock_model_config
40 mock_engine_client.input_processor = MagicMock()
41 mock_engine_client.renderer = MagicMock()
42
43 serving_models = OpenAIServingModels(
44 engine_client=mock_engine_client,
45 base_model_paths=BASE_MODEL_PATHS,
46 lora_modules=None,
47 )
48 await serving_models.init_static_loras()
49
50 # The mock is returned separately: `engine_client` is typed as the protocol,
51 # so assertion helpers are not visible through it.
52 return serving_models, mock_engine_client
53
54
55 async def _async_serving_models_init() -> OpenAIServingModels:
56 serving_models, _ = await _async_serving_models_init_with_mock()
57 return serving_models
58
59
60 @pytest.mark.asyncio
61 async def test_serving_model_name():
62 serving_models = await _async_serving_models_init()
63 assert serving_models.model_name(None) == MODEL_NAME
64 request = LoRARequest(
65 lora_name="adapter", lora_path="/path/to/adapter2", lora_int_id=1
66 )
67 assert serving_models.model_name(request) == request.lora_name
68
69
70 @pytest.mark.asyncio
71 async def test_load_lora_adapter_success():
72 serving_models = await _async_serving_models_init()
73 request = LoadLoRAAdapterRequest(lora_name="adapter", lora_path="/path/to/adapter2")
74 response = await serving_models.load_lora_adapter(request)
75 assert response == LORA_LOADING_SUCCESS_MESSAGE.format(lora_name="adapter")
76 assert len(serving_models.lora_requests) == 1
77 assert "adapter" in serving_models.lora_requests
78 assert serving_models.lora_requests["adapter"].lora_name == "adapter"
79
80
81 @pytest.mark.asyncio
82 async def test_load_lora_adapter_missing_fields():
83 serving_models = await _async_serving_models_init()
84 request = LoadLoRAAdapterRequest(lora_name="", lora_path="")
85 response = await serving_models.load_lora_adapter(request)
86 assert isinstance(response, ErrorResponse)
87 assert response.error.type == "InvalidUserInput"
88 assert response.error.code == HTTPStatus.BAD_REQUEST
89
90
91 @pytest.mark.asyncio
92 async def test_load_lora_adapter_duplicate():
93 serving_models = await _async_serving_models_init()
94 request = LoadLoRAAdapterRequest(
95 lora_name="adapter1", lora_path="/path/to/adapter1"
96 )
97 response = await serving_models.load_lora_adapter(request)
98 assert response == LORA_LOADING_SUCCESS_MESSAGE.format(lora_name="adapter1")
99 assert len(serving_models.lora_requests) == 1
100
101 request = LoadLoRAAdapterRequest(
102 lora_name="adapter1", lora_path="/path/to/adapter1"
103 )
104 response = await serving_models.load_lora_adapter(request)
105 assert isinstance(response, ErrorResponse)
106 assert response.error.type == "InvalidUserInput"
107 assert response.error.code == HTTPStatus.BAD_REQUEST
108 assert len(serving_models.lora_requests) == 1
109
110
111 @pytest.mark.asyncio
112 @pytest.mark.parametrize("load_inplace", [False, True])
113 async def test_load_lora_adapter_base_model_name(load_inplace: bool):
114 """A LoRA named after a served model would shadow it for all requests."""
115 serving_models, mock_engine_client = await _async_serving_models_init_with_mock()
116 request = LoadLoRAAdapterRequest(
117 lora_name=MODEL_NAME, lora_path="/path/to/adapter", load_inplace=load_inplace
118 )
119 response = await serving_models.load_lora_adapter(request)
120 assert isinstance(response, ErrorResponse)
121 assert response.error.type == "InvalidUserInput"
122 assert response.error.code == HTTPStatus.BAD_REQUEST
123 assert len(serving_models.lora_requests) == 0
124 mock_engine_client.add_lora.assert_not_called()
125
126
127 @pytest.mark.asyncio
128 async def test_unload_lora_adapter_success():
129 serving_models = await _async_serving_models_init()
130 request = LoadLoRAAdapterRequest(
131 lora_name="adapter1", lora_path="/path/to/adapter1"
132 )
133 response = await serving_models.load_lora_adapter(request)
134 assert len(serving_models.lora_requests) == 1
135
136 unload_request = UnloadLoRAAdapterRequest(lora_name="adapter1")
137 response = await serving_models.unload_lora_adapter(unload_request)
138 assert response == LORA_UNLOADING_SUCCESS_MESSAGE.format(lora_name="adapter1")
139 assert len(serving_models.lora_requests) == 0
140
141
142 @pytest.mark.asyncio
143 async def test_unload_lora_adapter_missing_fields():
144 serving_models = await _async_serving_models_init()
145 request = UnloadLoRAAdapterRequest(lora_name="", lora_int_id=None)
146 response = await serving_models.unload_lora_adapter(request)
147 assert isinstance(response, ErrorResponse)
148 assert response.error.type == "InvalidUserInput"
149 assert response.error.code == HTTPStatus.BAD_REQUEST
150
151
152 @pytest.mark.asyncio
153 async def test_unload_lora_adapter_not_found():
154 serving_models = await _async_serving_models_init()
155 request = UnloadLoRAAdapterRequest(lora_name="nonexistent_adapter")
156 response = await serving_models.unload_lora_adapter(request)
157 assert isinstance(response, ErrorResponse)
158 assert response.error.type == "NotFoundError"
159 assert response.error.code == HTTPStatus.NOT_FOUND
160
161
162 class _ConcretePoolingServing(PoolingBaseServing):
163 """Minimal concrete subclass used only in these unit tests."""
164
165 request_id_prefix = "test"
166
167 def get_io_processor(self, request):
168 raise NotImplementedError
169
170 def _build_response(self, ctx):
171 raise NotImplementedError
172
173
174 def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
175 lora_request = LoRARequest(
176 lora_name=lora_name, lora_int_id=1, lora_path="/path/to/lora"
177 )
178 mock_models = MagicMock()
179 mock_models.lora_requests = {lora_name: lora_request}
180 mock_models.is_base_model.side_effect = lambda name: name == MODEL_NAME
181
182 serving = object.__new__(_ConcretePoolingServing)
183 serving.models = mock_models
184 return serving
185
186
187 def _make_pooling_ctx(
188 model_name: str, serving: PoolingBaseServing
189 ) -> PoolingServeContext:
190 mock_request = MagicMock()
191 mock_request.model = model_name
192 return PoolingServeContext(
193 request=mock_request,
194 model_name=MODEL_NAME,
195 request_id="test-id",
196 pooling_params=PoolingParams(),
197 lora_request=serving._maybe_get_adapters(mock_request),
198 priorities=0,
199 prompt_extras=None,
200 )
201
202
203 def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
204 """LoRA adapter name must populate ctx.lora_request without raising."""
205 lora_name = "bot-embed-lora"
206 serving = _make_pooling_serving(lora_name)
207 ctx = _make_pooling_ctx(lora_name, serving)
208
209 assert ctx.lora_request is not None
210 assert ctx.lora_request.lora_name == lora_name
211
212
213 def test_pooling_maybe_get_adapters_unknown_model_raises():
214 """An unrecognised model name must still raise VLLMNotFoundError."""
215 serving = _make_pooling_serving("some-lora")
216
217 with pytest.raises(VLLMNotFoundError):
218 _make_pooling_ctx("unknown-model", serving)