import sys
from contextlib import contextmanager
from types import ModuleType, SimpleNamespace
import pytest
import torch
from modelexpress.refit.reshard.types import IncompleteRefit
from modelexpress.refit.timing import RefitTimingRecorder, use_refit_timing
from modelexpress_rl.inference.engines.vllm.installer import (
_update_mla_absorbed_weights,
_VllmInstaller,
)
from modelexpress_rl.inference.plan import (
PreparedCheckpointArtifact,
PreparedEngineTensors,
PreparedRuntimeTensors,
)
from modelexpress_rl.inference.receiver import PreparedCheckpoint
from torch import nn
def _install_fake_vllm(monkeypatch, initialize):
@contextmanager
def current_config(_config):
yield
class QuantizeMethodBase:
pass
modules = {
"vllm": ModuleType("vllm"),
"vllm.config": ModuleType("vllm.config"),
"vllm.model_executor": ModuleType("vllm.model_executor"),
"vllm.model_executor.layers": ModuleType("vllm.model_executor.layers"),
"vllm.model_executor.layers.quantization": ModuleType(
"vllm.model_executor.layers.quantization"
),
"vllm.model_executor.layers.quantization.base_config": ModuleType(
"vllm.model_executor.layers.quantization.base_config"
),
"vllm.model_executor.model_loader": ModuleType(
"vllm.model_executor.model_loader"
),
"vllm.model_executor.model_loader.default_loader": ModuleType(
"vllm.model_executor.model_loader.default_loader"
),
"vllm.model_executor.model_loader.reload": ModuleType(
"vllm.model_executor.model_loader.reload"
),
"vllm.model_executor.model_loader.reload.layerwise": ModuleType(
"vllm.model_executor.model_loader.reload.layerwise"
),
}
modules["vllm.config"].set_current_vllm_config = current_config
modules[
"vllm.model_executor.layers.quantization.base_config"
].QuantizeMethodBase = QuantizeMethodBase
layerwise = modules["vllm.model_executor.model_loader.reload.layerwise"]
layerwise.LAYERWISE_INFO = {}
layerwise.initialize_layerwise_reload = initialize
layerwise.finalize_layerwise_reload = lambda _model, _config: None
layerwise._copy_and_restore_kernel_tensors = lambda _layer, _info: None
for name, module in modules.items():
monkeypatch.setitem(sys.modules, name, module)
def test_installer_resolves_load_time_parameters_after_layerwise_reload(monkeypatch):
model = nn.Module()
model.register_parameter("packed", nn.Parameter(torch.zeros(1)))
def initialize(target):
del target._parameters["packed"]
target.register_parameter("weight", nn.Parameter(torch.empty(1, device="meta")))
_install_fake_vllm(monkeypatch, initialize)
installer = _VllmInstaller(
model=model,
vllm_config=object(),
model_config=object(),
device=torch.device("cpu"),
)
installer._process_and_commit({"weight": torch.tensor([7.0])})
assert model.weight.item() == 7.0
def test_installer_rejects_parameters_left_on_meta(monkeypatch):
model = nn.Module()
def initialize(target):
target.register_parameter("weight", nn.Parameter(torch.empty(1, device="meta")))
target.register_parameter("orphan", nn.Parameter(torch.empty(1, device="meta")))
_install_fake_vllm(monkeypatch, initialize)
installer = _VllmInstaller(
model=model,
vllm_config=object(),
model_config=object(),
device=torch.device("cpu"),
)
with pytest.raises(IncompleteRefit, match="left parameters on the meta device"):
installer._process_and_commit({"weight": torch.tensor([7.0])})
def test_installer_loads_prepared_checkpoint_inside_vllm_config(monkeypatch, tmp_path):
active_config = [None]
events = []
@contextmanager
def current_config(config):
active_config[0] = config
try:
yield
finally:
active_config[0] = None
def initialize(_model):
events.append(("initialize", active_config[0]))
_install_fake_vllm(monkeypatch, initialize)
sys.modules["vllm.config"].set_current_vllm_config = current_config
layerwise = sys.modules["vllm.model_executor.model_loader.reload.layerwise"]
layerwise.finalize_layerwise_reload = lambda _model, _config: events.append(
("finalize", active_config[0])
)
class DefaultModelLoader:
def __init__(self, load_config):
events.append(("loader", load_config.load_format))
def load_weights(self, model, model_config):
events.append(
(
"load",
active_config[0],
model_config.model,
model_config.revision,
)
)
model.weight.data.fill_(7.0)
sys.modules[
"vllm.model_executor.model_loader.default_loader"
].DefaultModelLoader = DefaultModelLoader
synchronized = []
monkeypatch.setattr(torch.cuda, "synchronize", synchronized.append)
model = nn.Linear(1, 1, bias=False)
model_config = SimpleNamespace(model="/launch", revision="main")
vllm_config = SimpleNamespace(
load_config=SimpleNamespace(load_format="modelexpress"),
quant_config=None,
)
installer = _VllmInstaller(
model=model,
vllm_config=vllm_config,
model_config=model_config,
device=torch.device("cpu"),
)
prepared_path = tmp_path / "prepared"
prepared = PreparedCheckpointArtifact(
PreparedCheckpoint(
target_version="version-a",
path=prepared_path,
metrics={"bytes_received": 7.0},
)
)
recorder = RefitTimingRecorder(backend="test", version="version-a")
with use_refit_timing(recorder):
metrics = installer.install(prepared)
assert events == [
("loader", "safetensors"),
("initialize", vllm_config),
("load", vllm_config, str(prepared_path), None),
("finalize", vllm_config),
]
assert model.weight.item() == 7.0
assert model_config.model == "/launch"
assert model_config.revision == "main"
assert vllm_config.load_config.load_format == "modelexpress"
assert synchronized == [torch.device("cpu")]
assert metrics["bytes_received"] == 7.0
assert metrics["perf/mx_receive_install_time"] >= 0
assert recorder.as_dict()["stages"]["post_install"]["count"] == 1
def test_installer_includes_prepared_engine_tensor_metrics(monkeypatch):
installer = _VllmInstaller(
model=nn.Module(),
vllm_config=object(),
model_config=object(),
device=torch.device("cpu"),
)
monkeypatch.setattr(installer, "install_tensors", lambda _tensors: None)
staged = SimpleNamespace(tensors={}, metrics={"bytes_received": 7.0})
metrics = installer.install(PreparedEngineTensors(staged=staged))
assert metrics["bytes_received"] == 7.0
assert metrics["perf/mx_receive_install_time"] >= 0
def test_installer_accepts_runtime_tensors_written_directly_in_place():
live = {
"weight": torch.tensor([7.0, 8.0]),
"runtime_buffer": torch.tensor([9.0]),
}
original_pointers = {name: tensor.data_ptr() for name, tensor in live.items()}
installer = _VllmInstaller(
model=nn.Module(),
vllm_config=object(),
model_config=object(),
device=torch.device("cpu"),
runtime_tensors=live,
)
staged = type(
"Staged", (), {"tensors": live, "metrics": {"bytes_received": 0}}
)()
metrics = installer.install(PreparedRuntimeTensors(staged=staged))
assert torch.equal(live["weight"], torch.tensor([7.0, 8.0]))
assert torch.equal(live["runtime_buffer"], torch.tensor([9.0]))
assert {name: tensor.data_ptr() for name, tensor in live.items()} == original_pointers
assert metrics["bytes_received"] == 0
def test_installer_rejects_a_runtime_staging_copy():
live = {"weight": torch.tensor([1.0])}
installer = _VllmInstaller(
model=nn.Module(),
vllm_config=object(),
model_config=object(),
device=torch.device("cpu"),
runtime_tensors=live,
)
staged = type(
"Staged",
(),
{"tensors": {"weight": torch.tensor([7.0])}, "metrics": {}},
)()
with pytest.raises(IncompleteRefit, match="directly into live storage"):
installer.install(PreparedRuntimeTensors(staged=staged))
def test_installer_rejects_quantized_mla_derived_weight_refresh():
model = nn.Module()
mla = nn.Module()
mla.kv_b_proj = nn.Linear(1, 1, bias=False)
mla.W_UV = torch.zeros(1)
model.add_module("mla", mla)
with pytest.raises(IncompleteRefit, match="quantized kv_b_proj"):
_update_mla_absorbed_weights(model, quantized=True)