import sys
from types import ModuleType
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import torch
import torch.nn as nn
from modelexpress import p2p_pb2
from modelexpress.engines.sglang.adapter import (
SglangAdapter,
build_sglang_load_context,
collect_sglang_tensors,
)
from modelexpress.engines.sglang.loader import MxModelLoader
def _load_config(**overrides):
defaults = dict(
tp_rank=3,
modelexpress_url="modelexpress-server:8001",
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
def _model_config(**overrides):
defaults = dict(
model_path="deepseek-ai/DeepSeek-V3",
dtype=torch.bfloat16,
quantization="fp8",
revision="abc123",
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
def _device_config(**overrides):
defaults = dict(device="cpu", gpu_id=0)
defaults.update(overrides)
return SimpleNamespace(**defaults)
def test_sglang_adapter_builds_identity_from_sglang_configs():
adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
with patch(
"modelexpress.engines.sglang.adapter._get_parallel_size",
side_effect=lambda name: {
"get_tensor_model_parallel_world_size": 8,
"get_pipeline_model_parallel_world_size": 2,
"get_moe_expert_parallel_world_size": 4,
}[name],
):
identity = adapter.build_identity()
assert identity.model_name == "deepseek-ai/DeepSeek-V3"
assert identity.backend_framework == p2p_pb2.BACKEND_FRAMEWORK_SGLANG
assert identity.tensor_parallel_size == 8
assert identity.pipeline_parallel_size == 2
assert identity.expert_parallel_size == 4
assert identity.dtype == "bfloat16"
assert identity.quantization == "fp8"
assert identity.revision == "abc123"
def test_sglang_context_uses_tp_rank_for_matching_and_url_override():
ctx = build_sglang_load_context(
_load_config(tp_rank=5, modelexpress_url="mx.example:9000"),
_model_config(),
_device_config(),
)
assert ctx.worker_rank == 5
assert ctx.global_rank == 5
assert ctx.mx_client.server_url == "mx.example:9000"
def test_sglang_context_separates_worker_rank_from_global_rank(monkeypatch):
sglang_mod = ModuleType("sglang")
srt_mod = ModuleType("sglang.srt")
distributed_mod = ModuleType("sglang.srt.distributed")
distributed_mod.get_tensor_model_parallel_rank = lambda: 1
distributed_mod.get_pipeline_model_parallel_rank = lambda: 2
distributed_mod.get_tensor_model_parallel_world_size = lambda: 4
srt_mod.distributed = distributed_mod
monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
monkeypatch.setitem(sys.modules, "sglang.srt.distributed", distributed_mod)
with patch("torch.distributed.is_available", return_value=True), patch(
"torch.distributed.is_initialized", return_value=True,
), patch("torch.distributed.get_rank", return_value=17):
ctx = build_sglang_load_context(
_load_config(tp_rank=5, modelexpress_url="mx.example:9000"),
_model_config(),
_device_config(),
)
assert ctx.worker_rank == 9
assert ctx.global_rank == 17
assert ctx.mx_client.server_url == "mx.example:9000"
def test_sglang_is_cuda_alike_uses_sglang_platform_helper(monkeypatch):
adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
sglang_mod = ModuleType("sglang")
srt_mod = ModuleType("sglang.srt")
utils_mod = ModuleType("sglang.srt.utils")
utils_mod.is_cuda_alike = lambda: True
monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
monkeypatch.setitem(sys.modules, "sglang.srt.utils", utils_mod)
assert adapter.is_cuda_alike() is True
def test_collect_sglang_tensors_preserves_non_contiguous_storage_names():
model = nn.Module()
model.weight_t = nn.Parameter(torch.randn(4, 3).T)
tensors = collect_sglang_tensors(model)
assert "weight_t.__storage" in tensors
assert tensors["weight_t.__storage"].dtype == torch.uint8
def test_collect_sglang_tensors_deduplicates_tied_parameters():
model = nn.Module()
shared = nn.Parameter(torch.randn(4, 3))
model.first = shared
model.second = shared
tensors = collect_sglang_tensors(model)
assert list(tensors) == ["first"]
def test_sglang_adapter_post_load_delegates_to_child_module():
adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
class ChildModel(nn.Module):
def __init__(self):
super().__init__()
self.post_load_called = False
def post_load_weights(self):
self.post_load_called = True
model = nn.Module()
model.child = ChildModel()
adapter.after_rdma_receive(SimpleNamespace(model=model))
assert model.child.post_load_called
def test_sglang_adapter_post_load_prefers_top_level_hook():
adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
class TopLevelModel(nn.Module):
def __init__(self):
super().__init__()
self.child = nn.Module()
self.child.post_load_called = False
self.post_load_called = False
def post_load_weights(self):
self.post_load_called = True
def child_post_load_weights():
model.child.post_load_called = True
model = TopLevelModel()
model.child.post_load_weights = child_post_load_weights
adapter.after_rdma_receive(SimpleNamespace(model=model))
assert model.post_load_called
assert not model.child.post_load_called
def _install_sglang_runai_loader_modules(monkeypatch, loader_cls, load_format):
sglang_mod = ModuleType("sglang")
srt_mod = ModuleType("sglang.srt")
configs_mod = ModuleType("sglang.srt.configs")
load_config_mod = ModuleType("sglang.srt.configs.load_config")
model_loader_mod = ModuleType("sglang.srt.model_loader")
loader_mod = ModuleType("sglang.srt.model_loader.loader")
load_config_mod.LoadFormat = load_format
loader_mod.RunaiModelStreamerLoader = loader_cls
monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
monkeypatch.setitem(sys.modules, "sglang.srt.configs", configs_mod)
monkeypatch.setitem(
sys.modules,
"sglang.srt.configs.load_config",
load_config_mod,
)
monkeypatch.setitem(sys.modules, "sglang.srt.model_loader", model_loader_mod)
monkeypatch.setitem(sys.modules, "sglang.srt.model_loader.loader", loader_mod)
def test_sglang_adapter_uses_native_model_streamer_loader(monkeypatch):
tensor = torch.randn(2, 2)
loader_instance = MagicMock()
loader_instance._get_all_weights.return_value = iter([("w", tensor)])
loader_cls = MagicMock(return_value=loader_instance)
load_format = SimpleNamespace(RUNAI_STREAMER="runai_streamer")
_install_sglang_runai_loader_modules(monkeypatch, loader_cls, load_format)
load_config = _load_config(model_loader_extra_config={"concurrency": 4})
model_config = _model_config()
adapter = SglangAdapter(
load_config,
model_config,
_device_config(device="cuda:2", gpu_id=2),
)
model = nn.Linear(2, 2)
weights = list(
adapter.build_model_streamer_weight_iter(
"az://models/deepseek-ai/DeepSeek-V3",
model=model,
)
)
stream_config = loader_cls.call_args.args[0]
assert stream_config is not load_config
assert stream_config.load_format == "runai_streamer"
assert stream_config.model_loader_extra_config == {"concurrency": 4}
stream_model_config = loader_instance._get_all_weights.call_args.args[0]
assert stream_model_config is not model_config
assert stream_model_config.model_weights == "az://models/deepseek-ai/DeepSeek-V3"
assert loader_instance._get_all_weights.call_args.args[1] is model
assert loader_instance.target_device_str == "cuda:2"
assert weights == [("w", tensor)]
def test_sglang_adapter_enables_distributed_model_streamer(monkeypatch):
loader_instance = MagicMock()
loader_instance._get_all_weights.return_value = iter([])
loader_cls = MagicMock(return_value=loader_instance)
load_format = SimpleNamespace(RUNAI_STREAMER="runai_streamer")
_install_sglang_runai_loader_modules(monkeypatch, loader_cls, load_format)
adapter = SglangAdapter(
_load_config(model_loader_extra_config={"concurrency": 4}),
_model_config(),
_device_config(device="cuda:0", gpu_id=0),
)
with patch(
"modelexpress.engines.sglang.adapter._get_parallel_size",
return_value=8,
), patch.object(adapter, "is_cuda_alike", return_value=True), patch.dict(
"os.environ", {"MX_MS_DISTRIBUTED": "1"}
):
list(
adapter.build_model_streamer_weight_iter(
"s3://bucket/deepseek-ai/DeepSeek-V3",
model=nn.Linear(2, 2),
)
)
stream_config = loader_cls.call_args.args[0]
assert stream_config.model_loader_extra_config == {
"concurrency": 4,
"distributed": True,
}
def test_sglang_model_streamer_requires_initialized_model(monkeypatch):
loader_cls = MagicMock()
load_format = SimpleNamespace(RUNAI_STREAMER="runai_streamer")
_install_sglang_runai_loader_modules(monkeypatch, loader_cls, load_format)
adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
try:
list(adapter.build_model_streamer_weight_iter("s3://bucket/model"))
except RuntimeError as exc:
assert "requires result.model" in str(exc)
else:
raise AssertionError("Expected missing model to fail")
def test_mx_model_loader_delegates_to_shared_strategy_chain():
model = nn.Linear(2, 2)
loader = MxModelLoader(_load_config(modelexpress_transport="nixl"))
with patch.object(
MxModelLoader,
"_load_model_via_nixl",
return_value=model,
) as load_via_nixl:
loaded = loader.load_model(
model=model,
model_config=_model_config(),
device_config=_device_config(),
)
assert loaded is model
load_via_nixl.assert_called_once_with(
model=model,
model_config=_model_config(),
device_config=_device_config(),
)
def test_mx_model_loader_nixl_path_delegates_to_shared_strategy_chain():
model = nn.Linear(2, 2)
loader = MxModelLoader(_load_config(modelexpress_transport="nixl"))
with patch(
"modelexpress.engines.sglang.loader.LoadStrategyChain.run",
return_value=model,
) as run:
loaded = loader._load_model_via_nixl(
model=model,
model_config=_model_config(),
device_config=_device_config(),
)
assert loaded is model
run.assert_called_once()
assert run.call_args.args[0] is model
ctx = run.call_args.args[1]
assert ctx.adapter.__class__ is SglangAdapter
assert ctx.identity.backend_framework == p2p_pb2.BACKEND_FRAMEWORK_SGLANG
def test_mx_model_loader_delegates_transfer_engine_transport_in_mx_package():
model = nn.Linear(2, 2)
loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
with patch.object(
MxModelLoader,
"_load_model_via_transfer_engine",
return_value=model,
) as load_via_transfer_engine:
loaded = loader.load_model(
model=model,
model_config=_model_config(),
device_config=_device_config(),
)
assert loaded is model
load_via_transfer_engine.assert_called_once_with(
model=model,
model_config=_model_config(),
device_config=_device_config(),
)
def test_mx_model_loader_rejects_unknown_transport_in_mx_package():
loader = MxModelLoader(_load_config(modelexpress_transport="unknown"))
try:
loader.load_model(
model=nn.Linear(2, 2),
model_config=_model_config(),
device_config=_device_config(),
)
except ValueError as exc:
assert "unknown" in str(exc)
else:
raise AssertionError("Expected unsupported transport to fail")
def test_transfer_engine_registers_discovered_tensor_map():
loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
tensor = torch.randn(2, 3)
calls = []
class FakeTransferEngine:
def register_memory(self, addr, size):
calls.append((addr, size))
return 0
weight_info = loader._register_transfer_engine_tensors(
{"weight.__storage": tensor},
FakeTransferEngine(),
)
assert calls == [(tensor.data_ptr(), tensor.numel() * tensor.element_size())]
assert weight_info == {
"weight.__storage": (
tensor.data_ptr(),
tensor.numel(),
tensor.element_size(),
)
}
def test_transfer_engine_receive_uses_discovered_tensor_map():
loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
tensor = torch.randn(2, 3)
ctx = SimpleNamespace(global_rank=0)
transferred = {}
source_worker = p2p_pb2.WorkerMetadata(
transfer_engine_session_id="te-session",
tensors=[
p2p_pb2.TensorDescriptor(
name="weight.__storage",
addr=1234,
size=tensor.numel() * tensor.element_size(),
device_id=0,
)
],
)
class FakeTransferEngine:
def batch_transfer_sync_read(
self,
session_id,
client_ptr_list,
seed_ptr_list,
client_len_list,
):
transferred["session_id"] = session_id
transferred["client_ptr_list"] = client_ptr_list
transferred["seed_ptr_list"] = seed_ptr_list
transferred["client_len_list"] = client_len_list
return 0
loader._receive_via_transfer_engine(
{"weight.__storage": tensor},
FakeTransferEngine(),
source_worker,
ctx,
)
assert transferred == {
"session_id": "te-session",
"client_ptr_list": [tensor.data_ptr()],
"seed_ptr_list": [1234],
"client_len_list": [tensor.numel() * tensor.element_size()],
}
def test_transfer_engine_publish_starts_non_nixl_heartbeat():
loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
ctx = SimpleNamespace(
global_rank=9,
worker_rank=3,
worker_id="worker-id",
device_id=1,
identity=p2p_pb2.SourceIdentity(model_name="sglang-model"),
mx_client=SimpleNamespace(),
)
published = {}
def publish_metadata(identity, worker, worker_id):
published["identity"] = identity
published["worker"] = worker
published["worker_id"] = worker_id
return "mx-source-id"
def update_status(**kwargs):
published["status"] = kwargs
return True
ctx.mx_client.publish_metadata = publish_metadata
ctx.mx_client.update_status = update_status
class FakeHeartbeat:
def __init__(self, **kwargs):
published["heartbeat"] = kwargs
def start(self):
published["heartbeat_started"] = True
with patch(
"modelexpress.engines.sglang.loader.HeartbeatThread",
FakeHeartbeat,
):
published_ok = loader._publish_transfer_engine_source(
ctx=ctx,
session_id="te-session",
weight_info={"weight": (1000, 4, 2)},
)
assert published_ok
assert published["worker"].transfer_engine_session_id == "te-session"
assert published["status"]["status"] == p2p_pb2.SOURCE_STATUS_READY
assert published["heartbeat"]["nixl_manager"] is None
assert published["heartbeat_started"]
def test_transfer_engine_publish_failure_is_non_fatal():
loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
ctx = SimpleNamespace(
global_rank=9,
worker_rank=3,
worker_id="worker-id",
device_id=1,
identity=p2p_pb2.SourceIdentity(model_name="sglang-model"),
mx_client=SimpleNamespace(
publish_metadata=lambda *args: (_ for _ in ()).throw(
RuntimeError("metadata down")
)
),
)
assert not loader._publish_transfer_engine_source(
ctx=ctx,
session_id="te-session",
weight_info={"weight": (1000, 4, 2)},
)