194 lines
7.0 KiB
Python
194 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from http import HTTPStatus
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from vllm import PoolingParams
|
|
from vllm.config import ModelConfig
|
|
from vllm.engine.protocol import EngineClient
|
|
from vllm.entrypoints.openai.engine.protocol import (
|
|
ErrorResponse,
|
|
)
|
|
from vllm.entrypoints.openai.models.protocol import BaseModelPath
|
|
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
|
from vllm.entrypoints.pooling.base.serving import PoolingBaseServing
|
|
from vllm.entrypoints.pooling.typing import PoolingServeContext
|
|
from vllm.entrypoints.serve.lora.protocol import (
|
|
LoadLoRAAdapterRequest,
|
|
UnloadLoRAAdapterRequest,
|
|
)
|
|
from vllm.exceptions import VLLMNotFoundError
|
|
from vllm.lora.request import LoRARequest
|
|
|
|
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
|
BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
|
|
LORA_LOADING_SUCCESS_MESSAGE = "Success: LoRA adapter '{lora_name}' added successfully."
|
|
LORA_UNLOADING_SUCCESS_MESSAGE = (
|
|
"Success: LoRA adapter '{lora_name}' removed successfully."
|
|
)
|
|
|
|
|
|
async def _async_serving_models_init() -> OpenAIServingModels:
|
|
mock_engine_client = MagicMock(spec=EngineClient)
|
|
# Set the max_model_len attribute to avoid missing attribute
|
|
mock_model_config = MagicMock(spec=ModelConfig)
|
|
mock_model_config.max_model_len = 2048
|
|
mock_engine_client.model_config = mock_model_config
|
|
mock_engine_client.input_processor = MagicMock()
|
|
mock_engine_client.renderer = MagicMock()
|
|
|
|
serving_models = OpenAIServingModels(
|
|
engine_client=mock_engine_client,
|
|
base_model_paths=BASE_MODEL_PATHS,
|
|
lora_modules=None,
|
|
)
|
|
await serving_models.init_static_loras()
|
|
|
|
return serving_models
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_serving_model_name():
|
|
serving_models = await _async_serving_models_init()
|
|
assert serving_models.model_name(None) == MODEL_NAME
|
|
request = LoRARequest(
|
|
lora_name="adapter", lora_path="/path/to/adapter2", lora_int_id=1
|
|
)
|
|
assert serving_models.model_name(request) == request.lora_name
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_load_lora_adapter_success():
|
|
serving_models = await _async_serving_models_init()
|
|
request = LoadLoRAAdapterRequest(lora_name="adapter", lora_path="/path/to/adapter2")
|
|
response = await serving_models.load_lora_adapter(request)
|
|
assert response == LORA_LOADING_SUCCESS_MESSAGE.format(lora_name="adapter")
|
|
assert len(serving_models.lora_requests) == 1
|
|
assert "adapter" in serving_models.lora_requests
|
|
assert serving_models.lora_requests["adapter"].lora_name == "adapter"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_load_lora_adapter_missing_fields():
|
|
serving_models = await _async_serving_models_init()
|
|
request = LoadLoRAAdapterRequest(lora_name="", lora_path="")
|
|
response = await serving_models.load_lora_adapter(request)
|
|
assert isinstance(response, ErrorResponse)
|
|
assert response.error.type == "InvalidUserInput"
|
|
assert response.error.code == HTTPStatus.BAD_REQUEST
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_load_lora_adapter_duplicate():
|
|
serving_models = await _async_serving_models_init()
|
|
request = LoadLoRAAdapterRequest(
|
|
lora_name="adapter1", lora_path="/path/to/adapter1"
|
|
)
|
|
response = await serving_models.load_lora_adapter(request)
|
|
assert response == LORA_LOADING_SUCCESS_MESSAGE.format(lora_name="adapter1")
|
|
assert len(serving_models.lora_requests) == 1
|
|
|
|
request = LoadLoRAAdapterRequest(
|
|
lora_name="adapter1", lora_path="/path/to/adapter1"
|
|
)
|
|
response = await serving_models.load_lora_adapter(request)
|
|
assert isinstance(response, ErrorResponse)
|
|
assert response.error.type == "InvalidUserInput"
|
|
assert response.error.code == HTTPStatus.BAD_REQUEST
|
|
assert len(serving_models.lora_requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unload_lora_adapter_success():
|
|
serving_models = await _async_serving_models_init()
|
|
request = LoadLoRAAdapterRequest(
|
|
lora_name="adapter1", lora_path="/path/to/adapter1"
|
|
)
|
|
response = await serving_models.load_lora_adapter(request)
|
|
assert len(serving_models.lora_requests) == 1
|
|
|
|
request = UnloadLoRAAdapterRequest(lora_name="adapter1")
|
|
response = await serving_models.unload_lora_adapter(request)
|
|
assert response == LORA_UNLOADING_SUCCESS_MESSAGE.format(lora_name="adapter1")
|
|
assert len(serving_models.lora_requests) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unload_lora_adapter_missing_fields():
|
|
serving_models = await _async_serving_models_init()
|
|
request = UnloadLoRAAdapterRequest(lora_name="", lora_int_id=None)
|
|
response = await serving_models.unload_lora_adapter(request)
|
|
assert isinstance(response, ErrorResponse)
|
|
assert response.error.type == "InvalidUserInput"
|
|
assert response.error.code == HTTPStatus.BAD_REQUEST
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unload_lora_adapter_not_found():
|
|
serving_models = await _async_serving_models_init()
|
|
request = UnloadLoRAAdapterRequest(lora_name="nonexistent_adapter")
|
|
response = await serving_models.unload_lora_adapter(request)
|
|
assert isinstance(response, ErrorResponse)
|
|
assert response.error.type == "NotFoundError"
|
|
assert response.error.code == HTTPStatus.NOT_FOUND
|
|
|
|
|
|
class _ConcretePoolingServing(PoolingBaseServing):
|
|
"""Minimal concrete subclass used only in these unit tests."""
|
|
|
|
request_id_prefix = "test"
|
|
|
|
def get_io_processor(self, request):
|
|
raise NotImplementedError
|
|
|
|
def _build_response(self, ctx):
|
|
raise NotImplementedError
|
|
|
|
|
|
def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
|
|
lora_request = LoRARequest(
|
|
lora_name=lora_name, lora_int_id=1, lora_path="/path/to/lora"
|
|
)
|
|
mock_models = MagicMock()
|
|
mock_models.lora_requests = {lora_name: lora_request}
|
|
mock_models.is_base_model.side_effect = lambda name: name == MODEL_NAME
|
|
|
|
serving = object.__new__(_ConcretePoolingServing)
|
|
serving.models = mock_models
|
|
return serving
|
|
|
|
|
|
def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
|
|
mock_request = MagicMock()
|
|
mock_request.model = model_name
|
|
return PoolingServeContext(
|
|
request=mock_request,
|
|
model_name=MODEL_NAME,
|
|
request_id="test-id",
|
|
pooling_params=PoolingParams(),
|
|
)
|
|
|
|
|
|
def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
|
|
"""LoRA adapter name must populate ctx.lora_request without raising."""
|
|
lora_name = "bot-embed-lora"
|
|
serving = _make_pooling_serving(lora_name)
|
|
ctx = _make_pooling_ctx(lora_name)
|
|
|
|
ctx.lora_request = serving._maybe_get_adapters(ctx.request)
|
|
|
|
assert ctx.lora_request is not None
|
|
assert ctx.lora_request.lora_name == lora_name
|
|
|
|
|
|
def test_pooling_maybe_get_adapters_unknown_model_raises():
|
|
"""An unrecognised model name must still raise VLLMNotFoundError."""
|
|
serving = _make_pooling_serving("some-lora")
|
|
ctx = _make_pooling_ctx("unknown-model")
|
|
|
|
with pytest.raises(VLLMNotFoundError):
|
|
serving._maybe_get_adapters(ctx.request)
|