main
py 133 lines 4.11 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import sys
5 from types import ModuleType, SimpleNamespace
6
7 import pytest
8 from torch import nn
9
10 from vllm.config import VllmConfig
11 from vllm.config.load import LoadConfig
12 from vllm.model_executor.model_loader import get_model_loader
13 from vllm.model_executor.model_loader.modelexpress_loader import (
14 ModelExpressModelLoader,
15 )
16
17
18 class FakeModelexpressLoader:
19 calls: list[tuple[str, tuple, dict]] = []
20 loaded_model: nn.Module
21
22 def __init__(self, load_config: LoadConfig):
23 self.load_config = load_config
24
25 def download_model(self, *args, **kwargs):
26 self.calls.append(("download_model", args, kwargs))
27
28 def load_weights(self, *args, **kwargs):
29 self.calls.append(("load_weights", args, kwargs))
30
31 def load_model(self, *args, **kwargs):
32 self.calls.append(("load_model", args, kwargs))
33 return self.loaded_model
34
35
36 def _install_fake_modelexpress(monkeypatch):
37 FakeModelexpressLoader.calls = []
38 FakeModelexpressLoader.loaded_model = nn.Module()
39
40 for name in [
41 "modelexpress",
42 "modelexpress.engines",
43 "modelexpress.engines.vllm",
44 ]:
45 monkeypatch.setitem(sys.modules, name, ModuleType(name))
46
47 module = ModuleType("modelexpress.engines.vllm.loader")
48 module.__dict__["MxModelLoader"] = FakeModelexpressLoader
49 monkeypatch.setitem(sys.modules, module.__name__, module)
50
51
52 def test_modelexpress_load_format_resolves_to_modelexpress_loader(monkeypatch):
53 _install_fake_modelexpress(monkeypatch)
54
55 loader = get_model_loader(LoadConfig(load_format="modelexpress"))
56
57 assert isinstance(loader, ModelExpressModelLoader)
58
59
60 def test_modelexpress_loader_delegates_to_modelexpress(monkeypatch):
61 _install_fake_modelexpress(monkeypatch)
62 loader = ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
63 model = nn.Module()
64 model_config = SimpleNamespace()
65 vllm_config = SimpleNamespace()
66
67 loader.download_model(model_config)
68 loader.load_weights(model, model_config)
69 FakeModelexpressLoader.loaded_model.train()
70 result = loader.load_model(
71 vllm_config=vllm_config,
72 model_config=model_config,
73 prefix="model",
74 )
75
76 assert result is FakeModelexpressLoader.loaded_model
77 assert not result.training
78 assert FakeModelexpressLoader.calls == [
79 ("download_model", (model_config,), {}),
80 ("load_weights", (model, model_config), {}),
81 (
82 "load_model",
83 (),
84 {
85 "vllm_config": vllm_config,
86 "model_config": model_config,
87 "prefix": "model",
88 },
89 ),
90 ]
91
92
93 def test_modelexpress_loader_missing_modelexpress_error(monkeypatch):
94 import importlib
95
96 def missing_modelexpress(name):
97 raise ModuleNotFoundError(name=name)
98
99 monkeypatch.setattr(importlib, "import_module", missing_modelexpress)
100
101 with pytest.raises(ImportError, match="requires the ModelExpress Python package"):
102 ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
103
104
105 def test_modelexpress_loader_preserves_internal_import_errors(monkeypatch):
106 import importlib
107
108 def missing_dependency(name):
109 raise ModuleNotFoundError(name="not_modelexpress_dependency")
110
111 monkeypatch.setattr(importlib, "import_module", missing_dependency)
112
113 with pytest.raises(ModuleNotFoundError) as exc_info:
114 ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
115 assert exc_info.value.name == "not_modelexpress_dependency"
116
117
118 def test_modelexpress_load_format_allows_object_storage_model_weights():
119 model_config = SimpleNamespace(
120 architecture="UnknownForTest",
121 config_updated=False,
122 convert_type=None,
123 is_hybrid=False,
124 model="test-model",
125 model_weights="s3://bucket/model",
126 )
127 vllm_config = object.__new__(VllmConfig)
128 vllm_config.model_config = model_config
129 vllm_config.load_config = LoadConfig(load_format="modelexpress")
130
131 vllm_config.try_verify_and_update_config()
132
133 assert vllm_config.load_config.load_format == "modelexpress"