134 lines
4.1 KiB
Python
134 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import sys
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
import pytest
|
|
from torch import nn
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.load import LoadConfig
|
|
from vllm.model_executor.model_loader import get_model_loader
|
|
from vllm.model_executor.model_loader.modelexpress_loader import (
|
|
ModelExpressModelLoader,
|
|
)
|
|
|
|
|
|
class FakeModelexpressLoader:
|
|
calls: list[tuple[str, tuple, dict]] = []
|
|
loaded_model: nn.Module
|
|
|
|
def __init__(self, load_config: LoadConfig):
|
|
self.load_config = load_config
|
|
|
|
def download_model(self, *args, **kwargs):
|
|
self.calls.append(("download_model", args, kwargs))
|
|
|
|
def load_weights(self, *args, **kwargs):
|
|
self.calls.append(("load_weights", args, kwargs))
|
|
|
|
def load_model(self, *args, **kwargs):
|
|
self.calls.append(("load_model", args, kwargs))
|
|
return self.loaded_model
|
|
|
|
|
|
def _install_fake_modelexpress(monkeypatch):
|
|
FakeModelexpressLoader.calls = []
|
|
FakeModelexpressLoader.loaded_model = nn.Module()
|
|
|
|
for name in [
|
|
"modelexpress",
|
|
"modelexpress.engines",
|
|
"modelexpress.engines.vllm",
|
|
]:
|
|
monkeypatch.setitem(sys.modules, name, ModuleType(name))
|
|
|
|
module = ModuleType("modelexpress.engines.vllm.loader")
|
|
module.__dict__["MxModelLoader"] = FakeModelexpressLoader
|
|
monkeypatch.setitem(sys.modules, module.__name__, module)
|
|
|
|
|
|
def test_modelexpress_load_format_resolves_to_modelexpress_loader(monkeypatch):
|
|
_install_fake_modelexpress(monkeypatch)
|
|
|
|
loader = get_model_loader(LoadConfig(load_format="modelexpress"))
|
|
|
|
assert isinstance(loader, ModelExpressModelLoader)
|
|
|
|
|
|
def test_modelexpress_loader_delegates_to_modelexpress(monkeypatch):
|
|
_install_fake_modelexpress(monkeypatch)
|
|
loader = ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
|
|
model = nn.Module()
|
|
model_config = SimpleNamespace()
|
|
vllm_config = SimpleNamespace()
|
|
|
|
loader.download_model(model_config)
|
|
loader.load_weights(model, model_config)
|
|
FakeModelexpressLoader.loaded_model.train()
|
|
result = loader.load_model(
|
|
vllm_config=vllm_config,
|
|
model_config=model_config,
|
|
prefix="model",
|
|
)
|
|
|
|
assert result is FakeModelexpressLoader.loaded_model
|
|
assert not result.training
|
|
assert FakeModelexpressLoader.calls == [
|
|
("download_model", (model_config,), {}),
|
|
("load_weights", (model, model_config), {}),
|
|
(
|
|
"load_model",
|
|
(),
|
|
{
|
|
"vllm_config": vllm_config,
|
|
"model_config": model_config,
|
|
"prefix": "model",
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
def test_modelexpress_loader_missing_modelexpress_error(monkeypatch):
|
|
import importlib
|
|
|
|
def missing_modelexpress(name):
|
|
raise ModuleNotFoundError(name=name)
|
|
|
|
monkeypatch.setattr(importlib, "import_module", missing_modelexpress)
|
|
|
|
with pytest.raises(ImportError, match="requires the ModelExpress Python package"):
|
|
ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
|
|
|
|
|
|
def test_modelexpress_loader_preserves_internal_import_errors(monkeypatch):
|
|
import importlib
|
|
|
|
def missing_dependency(name):
|
|
raise ModuleNotFoundError(name="not_modelexpress_dependency")
|
|
|
|
monkeypatch.setattr(importlib, "import_module", missing_dependency)
|
|
|
|
with pytest.raises(ModuleNotFoundError) as exc_info:
|
|
ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
|
|
assert exc_info.value.name == "not_modelexpress_dependency"
|
|
|
|
|
|
def test_modelexpress_load_format_allows_object_storage_model_weights():
|
|
model_config = SimpleNamespace(
|
|
architecture="UnknownForTest",
|
|
config_updated=False,
|
|
convert_type=None,
|
|
is_hybrid=False,
|
|
model="test-model",
|
|
model_weights="s3://bucket/model",
|
|
)
|
|
vllm_config = object.__new__(VllmConfig)
|
|
vllm_config.model_config = model_config
|
|
vllm_config.load_config = LoadConfig(load_format="modelexpress")
|
|
|
|
vllm_config.try_verify_and_update_config()
|
|
|
|
assert vllm_config.load_config.load_format == "modelexpress"
|