modelexpress-client 0.7.0

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Tests for the SGLang ModelExpress adapter and loader entrypoint."""

import os
import sys
import weakref
from contextlib import contextmanager
from types import ModuleType
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
import torch
import torch.nn as nn

from modelexpress import p2p_pb2
from modelexpress.engines.sglang.adapter import (
    SglangAdapter,
    build_sglang_load_context,
)
from modelexpress.engines.sglang.loader import MxModelLoader
from modelexpress.load_strategy.context import LoadResult


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)


@pytest.fixture(autouse=True)
def _stub_accelerator_backend_selection(monkeypatch, mock_accelerator_backend_cls):
    monkeypatch.setattr(
        "modelexpress.engines.sglang.adapter.accelerator_backend_for",
        lambda device: mock_accelerator_backend_cls(),
    )


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):
    """Keep SGLang's engine worker rank separate from the distributed rank."""
    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
    distributed_mod.get_world_group = lambda: SimpleNamespace(local_rank=3)
    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.local_rank == 3
    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(
    mock_accelerator_backend_cls,
):
    backend = mock_accelerator_backend_cls(torch_device_type="cpu")
    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    adapter.accelerator_backend = backend
    model = nn.Module()
    model.weight_t = nn.Parameter(torch.randn(4, 3).T)

    tensors = adapter.discover_tensors(SimpleNamespace(model=model))

    assert "weight_t.__storage" in tensors
    assert tensors["weight_t.__storage"].dtype == torch.uint8


def test_collect_sglang_tensors_deduplicates_tied_parameters(
    mock_accelerator_backend_cls,
):
    backend = mock_accelerator_backend_cls(torch_device_type="cpu")
    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    adapter.accelerator_backend = backend
    model = nn.Module()
    shared = nn.Parameter(torch.randn(4, 3))
    model.first = shared
    model.second = shared

    tensors = adapter.discover_tensors(SimpleNamespace(model=model))

    assert list(tensors) == ["first"]


def test_collect_sglang_tensors_includes_named_buffers(
    mock_accelerator_backend_cls,
):
    # capture_tensor_attrs promotes bare tensor assigns (e.g. the DeepSeek MLA
    # w_kc/w_vc) to non-persistent buffers; discovery must include
    # named_buffers so they register on both roles and transfer via RDMA.
    backend = mock_accelerator_backend_cls(torch_device_type="cpu")
    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    adapter.accelerator_backend = backend
    model = nn.Module()
    model.weight = nn.Parameter(torch.randn(4, 3))
    model.register_buffer("w_kc", torch.randn(3, 2), persistent=False)

    tensors = adapter.discover_tensors(SimpleNamespace(model=model))

    assert "weight" in tensors
    assert "w_kc" in tensors


def test_capture_then_collect_registers_bare_tensor_assign(
    mock_accelerator_backend_cls,
):
    # End-to-end of the Class H-direct fix: a bare tensor assign made under
    # capture_tensor_attrs (as the DeepSeek MLA mixin does self_attn.w_kc = t)
    # is promoted to a non-persistent buffer and then discovered by
    # discover_tensors' named_buffers() pass.
    from modelexpress.tensor_utils import capture_tensor_attrs

    backend = mock_accelerator_backend_cls(torch_device_type="cpu")
    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    adapter.accelerator_backend = backend
    model = nn.Module()
    model.self_attn = nn.Module()

    with capture_tensor_attrs(backend):
        model.self_attn.w_kc = torch.randn(3, 2)

    assert "w_kc" in dict(model.self_attn.named_buffers())
    assert "self_attn.w_kc" in adapter.discover_tensors(SimpleNamespace(model=model))


def test_sglang_adapter_discovery_uses_backend_predicate(
    monkeypatch,
    mock_accelerator_backend_cls,
):
    backend = mock_accelerator_backend_cls(torch_device_type="cpu")
    monkeypatch.setattr(
        "modelexpress.engines.sglang.adapter.accelerator_backend_for",
        lambda device: backend,
    )
    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    model = nn.Module()
    model.weight = nn.Parameter(torch.randn(4, 3))

    tensors = adapter.discover_tensors(SimpleNamespace(model=model))

    assert list(tensors) == ["weight"]


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._post_load_weights(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._post_load_weights(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_sglang_retry_initializes_model_with_configured_dtype(monkeypatch):
    original_dtype = torch.get_default_dtype()
    sglang_mod = ModuleType("sglang")
    srt_mod = ModuleType("sglang.srt")
    model_loader_mod = ModuleType("sglang.srt.model_loader")
    loader_mod = ModuleType("sglang.srt.model_loader.loader")
    model_loader_utils_mod = ModuleType("sglang.srt.model_loader.utils")
    observed_dtypes = []
    initial_model = nn.Linear(2, 2)
    initial_weight_ref = weakref.ref(initial_model.weight)

    @contextmanager
    def set_default_torch_dtype(dtype):
        previous_dtype = torch.get_default_dtype()
        torch.set_default_dtype(dtype)
        try:
            yield
        finally:
            torch.set_default_dtype(previous_dtype)

    loader_mod._get_quantization_config = lambda *_: None

    def initialize_model(*_):
        assert initial_weight_ref() is None
        observed_dtypes.append(torch.get_default_dtype())
        return nn.Linear(2, 2)

    loader_mod._initialize_model = initialize_model
    model_loader_utils_mod.set_default_torch_dtype = set_default_torch_dtype
    monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader", model_loader_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader.loader", loader_mod)
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.model_loader.utils",
        model_loader_utils_mod,
    )

    model_config = _model_config(dtype=torch.bfloat16)
    adapter = SglangAdapter(_load_config(), model_config, _device_config())
    result = LoadResult(
        value=initial_model,
        model=initial_model,
        publishable=True,
    )

    retried = adapter.reinit_for_retry(result)

    assert observed_dtypes == [torch.bfloat16]
    assert torch.get_default_dtype() == original_dtype
    assert retried.value is initial_model
    assert retried.model is initial_model
    assert list(initial_model.parameters())


def test_sglang_retry_failure_restores_envelope_and_original_error(monkeypatch):
    sglang_mod = ModuleType("sglang")
    srt_mod = ModuleType("sglang.srt")
    model_loader_mod = ModuleType("sglang.srt.model_loader")
    loader_mod = ModuleType("sglang.srt.model_loader.loader")
    model_loader_utils_mod = ModuleType("sglang.srt.model_loader.utils")

    @contextmanager
    def set_default_torch_dtype(_dtype):
        yield

    failure = RuntimeError("fresh initialization failed")
    loader_mod._get_quantization_config = lambda *_: None
    loader_mod._initialize_model = MagicMock(side_effect=failure)
    model_loader_utils_mod.set_default_torch_dtype = set_default_torch_dtype
    monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader", model_loader_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader.loader", loader_mod)
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.model_loader.utils",
        model_loader_utils_mod,
    )

    model = nn.Linear(2, 2)
    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    result = LoadResult(value=model, model=model, metadata={"attempt": 1})

    with pytest.raises(RuntimeError) as exc:
        adapter.reinit_for_retry(result)

    assert exc.value is failure
    assert result.value is model
    assert result.model is model
    assert result.metadata == {"attempt": 1}
    assert model.__dict__ == {}


def test_sglang_retry_reuses_root_for_native_fallback(monkeypatch):
    sglang_mod = ModuleType("sglang")
    srt_mod = ModuleType("sglang.srt")
    model_loader_mod = ModuleType("sglang.srt.model_loader")
    loader_mod = ModuleType("sglang.srt.model_loader.loader")
    model_loader_utils_mod = ModuleType("sglang.srt.model_loader.utils")
    configs_mod = ModuleType("sglang.srt.configs")
    load_config_mod = ModuleType("sglang.srt.configs.load_config")

    @contextmanager
    def set_default_torch_dtype(_dtype):
        yield

    initial_model = nn.Linear(2, 2)
    initial_weight_ref = weakref.ref(initial_model.weight)
    native_roots = []

    loader_mod._get_quantization_config = lambda *_: None

    def initialize_model(*_):
        assert initial_weight_ref() is None
        return nn.Linear(2, 2)

    class DefaultModelLoader:
        def __init__(self, _load_config):
            pass

        def _get_all_weights(self, _model_config, model):
            native_roots.append(model)
            return iter([])

        @staticmethod
        def load_weights_and_postprocess(model, _weights, _target_device):
            model.weight.data.fill_(7)

    loader_mod._initialize_model = initialize_model
    loader_mod.DefaultModelLoader = DefaultModelLoader
    model_loader_utils_mod.set_default_torch_dtype = set_default_torch_dtype
    load_config_mod.LoadFormat = SimpleNamespace(AUTO="auto")
    monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader", model_loader_mod)
    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader.loader", loader_mod)
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.model_loader.utils",
        model_loader_utils_mod,
    )
    monkeypatch.setitem(sys.modules, "sglang.srt.configs", configs_mod)
    monkeypatch.setitem(
        sys.modules,
        "sglang.srt.configs.load_config",
        load_config_mod,
    )

    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
    result = LoadResult(value=initial_model, model=initial_model)

    retried = adapter.reinit_for_retry(result)
    loaded = adapter.load_via_native(retried)

    assert loaded.model is initial_model
    assert native_roots == [initial_model]
    assert torch.all(initial_model.weight == 7)


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(),
    )


@pytest.mark.parametrize(
    ("ready_url", "health_gated"),
    [("", False), ("http://127.0.0.1:30000/health", True)],
)
def test_mx_model_loader_nixl_path_delegates_to_shared_strategy_chain(
    ready_url, health_gated
):
    from modelexpress.engines.sglang import loader as loader_mod

    model = nn.Linear(2, 2)
    loader = MxModelLoader(_load_config(modelexpress_transport="nixl"))

    with patch.dict(os.environ, {"MX_ARTIFACT_READY_URL": ready_url}), patch(
        "modelexpress.engines.sglang.loader.run_load_strategy_chain",
        return_value=model,
    ) as run, patch(
        "modelexpress.engines.sglang.loader.install_sglang_cache_artifacts",
    ) as install_artifacts, patch(
        "modelexpress.engines.sglang.loader.schedule_sglang_cache_artifact_publish",
    ) as schedule_artifacts:
        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
    if health_gated:
        # Bound to ctx so the URL resolves against this worker's node_rank
        # and head address, so identity is not asserted.
        assert callable(ctx.source_ready_fn)
    else:
        assert ctx.source_ready_fn is None
    install_artifacts.assert_called_once_with(ctx)
    schedule_artifacts.assert_called_once_with(ctx)


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_transfer_engine_receive_failure_reinitializes_before_native_fallback():
    transfer_engine = MagicMock()
    transfer_engine.register_memory.return_value = 0
    transfer_engine.batch_transfer_sync_read.return_value = -1
    load_config = _load_config(
        modelexpress_transport="transfer_engine",
        remote_instance_weight_loader_transfer_engine=transfer_engine,
        remote_instance_weight_loader_transfer_engine_session_id="target-session",
    )
    loader = MxModelLoader(load_config)
    initial_model = nn.Linear(2, 2)
    prepared_model = nn.Linear(2, 2)
    fresh_model = nn.Linear(2, 2)
    native_model = nn.Linear(2, 2)
    target_tensor = torch.randn(2, 3)
    native_tensor = torch.randn(3, 2)
    prepared_result = SimpleNamespace(value=prepared_model, model=prepared_model)
    fresh_result = SimpleNamespace(value=fresh_model, model=fresh_model)
    native_result = SimpleNamespace(value=native_model, model=native_model)
    adapter = MagicMock()
    adapter.before_rdma_receive.return_value = prepared_result
    adapter.discover_tensors.side_effect = [
        {"target": target_tensor},
        {"native": native_tensor},
    ]
    adapter.reinit_for_retry.return_value = fresh_result
    adapter.load_via_native.return_value = native_result
    ctx = SimpleNamespace(
        global_rank=0,
        identity=SimpleNamespace(model_name="model"),
        adapter=adapter,
        tensors={},
    )
    source_worker = p2p_pb2.WorkerMetadata(
        transfer_engine_session_id="source-session",
        tensors=[
            p2p_pb2.TensorDescriptor(
                name="target",
                addr=1234,
                size=target_tensor.numel() * target_tensor.element_size(),
            )
        ],
    )

    with patch(
        "modelexpress.engines.sglang.loader.build_sglang_load_context",
        return_value=ctx,
    ), patch.object(
        loader,
        "_find_transfer_engine_source",
        return_value=source_worker,
    ), patch.object(
        loader,
        "_publish_transfer_engine_source",
        return_value=True,
    ) as publish:
        loaded = loader._load_model_via_transfer_engine(
            model=initial_model,
            model_config=_model_config(),
            device_config=_device_config(),
        )

    assert loaded is native_model
    transfer_engine.unregister_memory.assert_called_once_with(
        target_tensor.data_ptr()
    )
    adapter.reinit_for_retry.assert_called_once_with(prepared_result)
    adapter.load_via_native.assert_called_once_with(fresh_result)
    publish.assert_called_once()
    assert publish.call_args.kwargs["weight_info"] == {
        "native": (
            native_tensor.data_ptr(),
            native_tensor.numel(),
            native_tensor.element_size(),
        )
    }


def test_transfer_engine_source_registration_failure_keeps_native_model():
    transfer_engine = MagicMock()
    transfer_engine.register_memory.side_effect = [0, -1]
    load_config = _load_config(
        modelexpress_transport="transfer_engine",
        remote_instance_weight_loader_transfer_engine=transfer_engine,
        remote_instance_weight_loader_transfer_engine_session_id="source-session",
    )
    loader = MxModelLoader(load_config)
    initial_model = nn.Linear(2, 2)
    native_model = nn.Linear(2, 2)
    first = torch.randn(2, 3)
    second = torch.randn(3, 2)
    native_result = SimpleNamespace(value=native_model, model=native_model)
    adapter = MagicMock()
    adapter.load_via_native.return_value = native_result
    adapter.discover_tensors.return_value = {"first": first, "second": second}
    ctx = SimpleNamespace(
        global_rank=0,
        identity=SimpleNamespace(model_name="model"),
        adapter=adapter,
        tensors={},
    )

    with patch(
        "modelexpress.engines.sglang.loader.build_sglang_load_context",
        return_value=ctx,
    ), patch.object(
        loader,
        "_find_transfer_engine_source",
        return_value=None,
    ), patch.object(
        loader,
        "_publish_transfer_engine_source",
    ) as publish:
        loaded = loader._load_model_via_transfer_engine(
            model=initial_model,
            model_config=_model_config(),
            device_config=_device_config(),
        )

    assert loaded is native_model
    transfer_engine.unregister_memory.assert_called_once_with(first.data_ptr())
    adapter.reinit_for_retry.assert_not_called()
    adapter.load_via_native.assert_called_once()
    publish.assert_not_called()


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)
    empty = torch.empty(0)
    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, "empty": empty},
        FakeTransferEngine(),
    )

    assert calls == [(tensor.data_ptr(), tensor.numel() * tensor.element_size())]
    assert weight_info == {
        "weight.__storage": (
            tensor.data_ptr(),
            tensor.numel(),
            tensor.element_size(),
        ),
        "empty": (empty.data_ptr(), empty.numel(), empty.element_size()),
    }


def test_transfer_engine_receive_uses_discovered_tensor_map():
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
    tensor = torch.randn(2, 3)
    empty = torch.empty(0)
    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,
            ),
            p2p_pb2.TensorDescriptor(name="empty", addr=0, size=0, 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, "empty": empty},
        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()],
    }


@pytest.mark.parametrize(
    ("ready_url", "health_gated"),
    [("", False), ("http://127.0.0.1:30000/health", True)],
)
def test_transfer_engine_publish_starts_non_nixl_heartbeat(
    ready_url, health_gated
):
    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(),
        accelerator_backend=SimpleNamespace(name="cuda"),
    )
    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 FakePublisher:
        def __init__(self, **kwargs):
            published["heartbeat"] = kwargs

        def start(self):
            published["heartbeat_started"] = True

    with patch.dict(os.environ, {"MX_ARTIFACT_READY_URL": ready_url}), patch(
        "modelexpress.engines.sglang.loader.PublisherThread",
        FakePublisher,
    ):
        published_ok = loader._publish_transfer_engine_source(
            ctx=ctx,
            session_id="te-session",
            weight_info={"weight": (1000, 4, 2)},
        )

    assert published_ok
    assert "identity" not in published
    assert "status" not in published
    assert published["heartbeat"]["nixl_manager"] is None
    assert callable(published["heartbeat"]["publish_fn"])
    if health_gated:
        assert callable(published["heartbeat"]["ready_fn"])
    else:
        assert published["heartbeat"]["ready_fn"] is None
    assert published["heartbeat_started"]

    mx_source_id = published["heartbeat"]["publish_fn"]()

    assert mx_source_id == "mx-source-id"
    assert published["worker"].transfer_engine_session_id == "te-session"
    assert published["worker"].accelerator == "cuda"
    assert "status" not in published


def test_transfer_engine_publish_failure_is_deferred_to_publisher():
    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"),
        accelerator_backend=SimpleNamespace(name="cuda"),
        mx_client=SimpleNamespace(
            publish_metadata=lambda *args: (_ for _ in ()).throw(
                RuntimeError("metadata down")
            )
        ),
    )

    scheduled = {}

    class FakePublisher:
        def __init__(self, **kwargs):
            scheduled.update(kwargs)

        def start(self):
            pass

    with patch(
        "modelexpress.engines.sglang.loader.PublisherThread",
        FakePublisher,
    ):
        assert loader._publish_transfer_engine_source(
            ctx=ctx,
            session_id="te-session",
            weight_info={"weight": (1000, 4, 2)},
        )

    with pytest.raises(RuntimeError, match="metadata down"):
        scheduled["publish_fn"]()


# ---------------------------------------------------------------------------
# _find_transfer_engine_source: source-selector wiring
# ---------------------------------------------------------------------------


def _te_ref(sid, wid, rank=0):
    return p2p_pb2.SourceInstanceRef(
        mx_source_id=sid, worker_id=wid, model_name="m", worker_rank=rank
    )


def _te_ctx(instances):
    ctx = SimpleNamespace()
    ctx.global_rank = 0
    ctx.worker_rank = 0
    ctx.worker_id = "tgt-0"
    ctx.identity = SimpleNamespace(model_name="m")
    ctx.mx_client = MagicMock()
    ctx.mx_client.list_sources.return_value = p2p_pb2.ListSourcesResponse(
        instances=instances
    )
    return ctx


def _te_meta(found=True, transfer_engine=False):
    if not found:
        return SimpleNamespace(found=False, worker=None)
    worker = (
        p2p_pb2.WorkerMetadata(worker_rank=0, transfer_engine_session_id="te")
        if transfer_engine
        else p2p_pb2.WorkerMetadata(worker_rank=0)
    )
    return SimpleNamespace(found=True, worker=worker)


def test_te_find_source_filters_rank_and_returns_transfer_engine(monkeypatch):
    monkeypatch.delenv("MX_P2P_SOURCE_SELECTOR", raising=False)
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
    ctx = _te_ctx(
        [
            _te_ref("s0aaaaaaaaaaaaaa", "w0", rank=0),
            _te_ref("s1aaaaaaaaaaaaaa", "w1", rank=1),  # wrong rank -> filtered out
            _te_ref("s2aaaaaaaaaaaaaa", "w2", rank=0),
        ]
    )
    ctx.mx_client.get_metadata.side_effect = lambda mx_source_id, worker_id: _te_meta(
        transfer_engine=(worker_id == "w2")
    )

    worker = loader._find_transfer_engine_source(ctx)
    assert worker is not None
    assert worker.WhichOneof("backend_metadata") == "transfer_engine_session_id"
    queried = {c.kwargs["worker_id"] for c in ctx.mx_client.get_metadata.call_args_list}
    assert "w1" not in queried  # rank-mismatched source never queried


def test_te_find_source_iterates_in_selector_order_and_none_when_no_match(monkeypatch):
    class _ReverseSelector:
        name = "reverse"

        def order(self, candidates, context):
            return list(reversed(candidates))

    monkeypatch.setattr(
        "modelexpress.engines.sglang.loader.get_configured_selector",
        lambda: _ReverseSelector(),
    )
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
    ctx = _te_ctx([_te_ref(f"s{i}aaaaaaaaaaaaaa", f"w{i}", rank=0) for i in range(3)])
    ctx.mx_client.get_metadata.side_effect = lambda mx_source_id, worker_id: _te_meta(
        transfer_engine=False
    )

    # No transfer_engine source -> None, and iteration follows the selector order.
    assert loader._find_transfer_engine_source(ctx) is None
    order = [c.kwargs["worker_id"] for c in ctx.mx_client.get_metadata.call_args_list]
    assert order == ["w2", "w1", "w0"]


def test_te_find_source_skips_not_found(monkeypatch):
    # Force identity order so w0 (not-found) is queried first and the
    # `if not metadata.found: continue` branch is actually exercised.
    class _IdentitySelector:
        name = "identity"

        def order(self, candidates, context):
            return list(candidates)

    monkeypatch.setattr(
        "modelexpress.engines.sglang.loader.get_configured_selector",
        lambda: _IdentitySelector(),
    )
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
    ctx = _te_ctx(
        [_te_ref("s0aaaaaaaaaaaaaa", "w0"), _te_ref("s1aaaaaaaaaaaaaa", "w1")]
    )
    ctx.mx_client.get_metadata.side_effect = lambda mx_source_id, worker_id: _te_meta(
        found=(worker_id == "w1"), transfer_engine=(worker_id == "w1")
    )
    worker = loader._find_transfer_engine_source(ctx)
    assert worker is not None
    assert worker.WhichOneof("backend_metadata") == "transfer_engine_session_id"
    assert [
        c.kwargs["worker_id"] for c in ctx.mx_client.get_metadata.call_args_list
    ] == ["w0", "w1"]


# ---------------------------------------------------------------------------
# TransferEngine source discovery -> selection metrics
#
# This transport bypasses LoadStrategyChain and RdmaStrategy, so the D8
# instrumentation added to RdmaStrategy._find_source_instances does not reach
# it. Without its own recording, a mooncake/transfer_engine pod exports
# mx_build_info and nothing else in every state, and a metadata-backend outage
# is indistinguishable from a cluster that has published no peers.
# ---------------------------------------------------------------------------


def _patched_te_metrics(monkeypatch):
    m = MagicMock()
    monkeypatch.setattr("modelexpress.engines.sglang.loader.selection_metrics", m)
    return m


def test_te_find_source_records_a_list_sources_error(monkeypatch):
    m = _patched_te_metrics(monkeypatch)
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
    ctx = _te_ctx([])
    ctx.mx_client.list_sources.side_effect = RuntimeError("grpc down")

    assert loader._find_transfer_engine_source(ctx) is None

    assert m.record_list_sources.call_args.args[1] == "error"


def test_te_find_source_records_a_zero_funnel_when_no_peers(monkeypatch):
    m = _patched_te_metrics(monkeypatch)
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))

    assert loader._find_transfer_engine_source(_te_ctx([])) is None

    assert m.record_list_sources.call_args.args[1] == "empty"
    observed = {
        call.args[1]: call.args[2] for call in m.observe_candidates.call_args_list
    }
    assert observed == {"listed": 0, "rank_matched": 0}


def test_te_find_source_records_the_funnel_on_success(monkeypatch):
    m = _patched_te_metrics(monkeypatch)
    loader = MxModelLoader(_load_config(modelexpress_transport="transfer_engine"))
    ctx = _te_ctx(
        [_te_ref("s0aaaaaaaaaaaaaa", "w0"), _te_ref("s1aaaaaaaaaaaaaa", "w1")]
    )

    loader._find_transfer_engine_source(ctx)

    assert m.record_list_sources.call_args.args[1] == "ok"
    observed = {
        call.args[1]: call.args[2] for call in m.observe_candidates.call_args_list
    }
    assert observed["listed"] == 2