neuron-tai/tests/test_real_model_backend.py

"""US-012 tests for the real PyTorch node backend."""

import json
import os
from pathlib import Path
import sys
import threading
import time
import types
import urllib.request

import pytest

from meshnet_node.model_backend import (
    InsufficientVRAMError,
    PartialModelLoadUnsupported,
    TensorPayload,
    TorchModelShard,
    _call_layer,
    _checkpoint_tensor_name_for_model,
    _load_partial_model_from_snapshot,
    _should_partial_materialize_shard,
    _decoder_attention_mask,
    _int_tensor_header,
    build_quantization_config,
    validate_quantization,
)
from meshnet_node.torch_server import TorchNodeServer


class _FakeBackend:
    model_id = "fake-model"
    total_layers = 12
    is_head = True
    is_tail = False

    def encode_prompt(self, prompt: str) -> TensorPayload:
        assert prompt == "The capital of France is"
        return TensorPayload(
            body=b"\x00" * (1 * 6 * 8 * 2),
            shape=[1, 6, 8],
            attention_mask_header=None,
            position_ids_header=None,
        )

    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
        assert shape == [1, 6, 8]
        return TensorPayload(
            body=body,
            shape=shape,
            attention_mask_header=attention_mask_header,
            position_ids_header=position_ids_header,
        )


class _FakeTailBackend(_FakeBackend):
    is_head = False
    is_tail = True

    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
        assert len(body) == 1 * 6 * 8 * 2
        return " Paris"


class _FakeFullBackend(_FakeBackend):
    is_head = True
    is_tail = True

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 16,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str:
        assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
        assert max_new_tokens == 7
        assert temperature == 1.0
        assert top_p == 1.0
        return "56"

    def count_prompt_tokens(self, messages: list[dict]) -> int:
        assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
        return 8

    def count_text_tokens(self, text: str) -> int:
        assert text == "56"
        return 1


class _FakeChatTokenizer:
    eos_token = ""

    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
        assert add_generation_prompt is True
        assert tokenize is False
        return "debug prompt"


class _FakePipelineHeadBackend(_FakeBackend):
    tokenizer = _FakeChatTokenizer()

    def encode_prompt(self, prompt: str) -> TensorPayload:
        assert prompt.startswith("debug prompt")
        return TensorPayload(
            body=b"\x00" * (1 * 6 * 8 * 2),
            shape=[1, 6, 8],
            attention_mask_header=None,
            position_ids_header=None,
        )


class _FakePipelineTailBackend(_FakeTailBackend):
    def __init__(self) -> None:
        self.start_layers: list[int | None] = []

    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
        self.start_layers.append(start_layer)
        assert len(body) == 1 * 6 * 8 * 2
        return " token"


class _BlockingStreamingTailBackend(_FakeTailBackend):
    def __init__(self, second_token_release: threading.Event) -> None:
        self._release = second_token_release
        self.calls = 0

    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
        self.calls += 1
        if self.calls == 1:
            return " first"
        self._release.wait(timeout=3.0)
        return " second"


def test_quantization_flag_validation():
    assert validate_quantization("bfloat16") == "bfloat16"
    assert validate_quantization("int8") == "int8"
    assert validate_quantization("nf4") == "nf4"
    with pytest.raises(ValueError, match="quantization"):
        validate_quantization("float32")


def test_node_package_declares_torch_dependency():
    pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")

    assert '"torch>=' in pyproject


def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
    calls = []

    class FakeBitsAndBytesConfig:
        def __init__(self, **kwargs):
            calls.append(kwargs)

    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16"))
    monkeypatch.setitem(
        sys.modules,
        "transformers",
        types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig),
    )

    assert build_quantization_config("bfloat16") is None
    build_quantization_config("int8")
    build_quantization_config("nf4")

    assert calls == [
        {"load_in_8bit": True},
        {
            "load_in_4bit": True,
            "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_compute_dtype": "bf16",
        },
    ]


def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
    node = TorchNodeServer(backend=_FakeBackend())
    port = node.start()
    try:
        payload = json.dumps({"prompt": "The capital of France is"}).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/forward",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = resp.read()
            headers = {key.lower(): value for key, value in resp.headers.items()}

        assert len(body) == 1 * 6 * 8 * 2
        assert headers["x-meshnet-shape"] == "1,6,8"
        assert headers["x-meshnet-dtype"] == "bfloat16"
        assert headers["x-meshnet-wire"] == "2"
    finally:
        node.stop()


def test_tail_forward_returns_text_completion_from_binary_activations():
    node = TorchNodeServer(backend=_FakeTailBackend())
    port = node.start()
    try:
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/forward",
            data=b"\x00" * (1 * 6 * 8 * 2),
            headers={
                "Content-Type": "application/octet-stream",
                "X-Meshnet-Shape": "1,6,8",
                "X-Meshnet-Dtype": "bfloat16",
                "X-Meshnet-Session": "session-1",
                "X-Meshnet-Chunk-Index": "0",
                "X-Meshnet-Chunk-Total": "1",
                "X-Meshnet-Hop-Index": "1",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = json.loads(resp.read())

        assert body == {"text": " Paris"}
        assert node.received_activations
        assert node.forward_chunk_count == 1
    finally:
        node.stop()


def test_full_model_chat_completion_uses_generation_not_single_token_decode(capsys):
    node = TorchNodeServer(backend=_FakeFullBackend())
    port = node.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "What is 7 times 8?"}],
            "max_tokens": 7,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Request-Id": "req-test-123",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = json.loads(resp.read())

        assert body["choices"][0]["message"]["content"] == "56"
        assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}
    finally:
        node.stop()

    out = capsys.readouterr().out
    assert "  [node] processing chat model='fake-model' stream=False max_tokens=7 request_id=req-test-123" in out
    assert "  [node] chat complete tokens=1 elapsed_s=" in out


def test_pipeline_hop_logs_are_suppressed_without_debug(capsys):
    tail_backend = _FakePipelineTailBackend()
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=tail_backend)
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 1,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            body = json.loads(resp.read())
    finally:
        head.stop()
        tail.stop()

    out = capsys.readouterr().out
    assert body["choices"][0]["message"]["content"] == " token"
    assert tail_backend.start_layers == [22]
    assert "pipeline hop 0:" not in out
    assert "pipeline hop 0 returned text" not in out


def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True)
    tail = TorchNodeServer(backend=_FakePipelineTailBackend())
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 1,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            json.loads(resp.read())
    finally:
        head.stop()
        tail.stop()

    out = capsys.readouterr().out
    assert f"  [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out
    assert "  [node] pipeline hop 0 returned text=' token'" in out


def test_split_shard_chat_streams_each_generated_token_incrementally():
    release_second = threading.Event()
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
    head_port = head.start()
    tail_port = tail.start()
    response = None
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "stream": True,
            "max_tokens": 2,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        response = urllib.request.urlopen(req, timeout=5)

        first_token_line = ""
        deadline = time.time() + 2.0
        while time.time() < deadline:
            line = response.readline().decode()
            if '"content": " first"' in line:
                first_token_line = line
                break

        assert first_token_line
        assert not release_second.is_set()
        release_second.set()
        rest = response.read().decode()
    finally:
        release_second.set()
        if response is not None:
            response.close()
        head.stop()
        tail.stop()

    assert '"content": " second"' in rest
    assert "data: [DONE]" in rest


def test_current_requests_snapshot_while_generating():
    release_second = threading.Event()
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
    head_port = head.start()
    tail_port = tail.start()
    response = None
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "stream": True,
            "max_tokens": 2,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Request-Id": "req-live-1",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        response = urllib.request.urlopen(req, timeout=5)
        deadline = time.time() + 2.0
        while time.time() < deadline:
            live = head.current_requests
            if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1:
                break
            time.sleep(0.02)
        assert head.current_requests
        snap = head.current_requests[0]
        assert snap["request_id"] == "req-live-1"
        assert snap["tokens"] >= 1
        assert snap["tokens_per_sec"] >= 0
        assert snap["routing_complete"] is True
        release_second.set()
        response.read()
    finally:
        release_second.set()
        if response is not None:
            response.close()
        head.stop()
        tail.stop()

    assert head.current_requests == []


def test_distributed_generating_log_includes_tps(capsys):
    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
    tail = TorchNodeServer(backend=_FakePipelineTailBackend())
    head_port = head.start()
    tail_port = tail.start()
    try:
        payload = json.dumps({
            "model": "fake-model",
            "messages": [{"role": "user", "content": "hello"}],
            "max_tokens": 1,
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/v1/chat/completions",
            data=payload,
            headers={
                "Content-Type": "application/json",
                "X-Meshnet-Route": json.dumps([
                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
                ]),
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            json.loads(resp.read())
    finally:
        head.stop()
        tail.stop()

    out = capsys.readouterr().out
    assert "generating step=1/1" in out
    assert " tps=" in out
    assert "generation complete tokens=1" in out
    assert out.count("generating step=1/1") == 1


def test_int_tensor_header_serializes_torch_tensors():
    torch = pytest.importorskip("torch")

    header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))

    assert header.startswith("1,3:")


def test_decoder_attention_mask_is_causal_float_mask():
    torch = pytest.importorskip("torch")

    hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16)
    mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch)

    assert mask.shape == (1, 1, 3, 3)
    assert mask.dtype == torch.bfloat16
    assert mask[0, 0, 0, 1] < 0
    assert mask[0, 0, 2, 0] == 0


def test_call_layer_passes_rotary_position_embeddings():
    class NeedsPositionEmbeddings:
        def __call__(self, hidden_states, **kwargs):
            assert kwargs["position_embeddings"] == "rotary"
            return hidden_states

    assert _call_layer(
        NeedsPositionEmbeddings(),
        "hidden",
        attention_mask=None,
        position_ids="positions",
        position_embeddings="rotary",
    ) == "hidden"


def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')

    assert _should_partial_materialize_shard(
        str(snapshot_dir),
        4,
        7,
        total_layers_hint=40,
        uses_quantized_weights=False,
    ) is True
    assert _should_partial_materialize_shard(
        str(snapshot_dir),
        0,
        39,
        total_layers_hint=40,
        uses_quantized_weights=False,
    ) is True
    assert _should_partial_materialize_shard(
        str(snapshot_dir),
        4,
        7,
        total_layers_hint=40,
        uses_quantized_weights=True,
    ) is False
    assert _should_partial_materialize_shard(
        "repo/model",
        4,
        7,
        total_layers_hint=40,
        uses_quantized_weights=False,
    ) is False


def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm():
    class TextOnlyModel:
        def __init__(self):
            self.model = types.SimpleNamespace(layers=[])

    model = TextOnlyModel()
    assert _checkpoint_tensor_name_for_model(
        model,
        "model.language_model.layers.0.mlp.gate.weight",
    ) == "model.layers.0.mlp.gate.weight"
    assert _checkpoint_tensor_name_for_model(
        model,
        "model.language_model.embed_tokens.weight",
    ) == "model.embed_tokens.weight"


def test_checkpoint_tensor_name_kept_for_multimodal_backbone():
    class MultimodalModel:
        def __init__(self):
            self.model = types.SimpleNamespace(language_model=types.SimpleNamespace())

    model = MultimodalModel()
    name = "model.language_model.layers.0.mlp.gate.weight"
    assert _checkpoint_tensor_name_for_model(model, name) == name


def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path):
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text(json.dumps({
        "text_config": {"num_hidden_layers": 3},
    }))
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
        }
    }))
    (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")

    class FakeModule:
        def __init__(self):
            self.to_calls = []

        def to(self, device):
            self.to_calls.append(device)
            return self

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                layers=[FakeModule(), FakeModule(), FakeModule()],
                rotary_emb=FakeModule(),
            )

        def tie_weights(self):
            pass

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            return types.SimpleNamespace(
                text_config=types.SimpleNamespace(num_hidden_layers=3),
                get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
            )

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            return FakeModel()

    set_calls = []

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        set_calls.append(tensor_name)

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            self.filename = Path(filename).name

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return tensor_name

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        1,
        1,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=UnusedContext,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )

    assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]


def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path):
    # Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.*
    # tensors that the text-only CausalLM never builds — they must be skipped,
    # not assigned (assignment raises AttributeError: 'mtp' / 'visual').
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text(json.dumps({
        "text_config": {"num_hidden_layers": 3},
    }))
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
            "mtp.layers.1.input_layernorm.weight": "shard-2.safetensors",
            "model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors",
        }
    }))
    (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")

    class FakeModule:
        def to(self, device):
            return self

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                layers=[FakeModule(), FakeModule(), FakeModule()],
                rotary_emb=FakeModule(),
            )

        def tie_weights(self):
            pass

        def state_dict(self):
            return {"model.layers.1.self_attn.q_proj.weight": None}

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            return types.SimpleNamespace(
                text_config=types.SimpleNamespace(num_hidden_layers=3),
                get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
            )

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            return FakeModel()

    set_calls = []

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        set_calls.append(tensor_name)

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            pass

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return tensor_name

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        1,
        1,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=UnusedContext,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )

    assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]


def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {
            "model.embed_tokens.weight": "shard-1.safetensors",
            "model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors",
            "model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
            "model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors",
            "model.norm.weight": "shard-3.safetensors",
            "lm_head.weight": "shard-3.safetensors",
        }
    }))
    for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"):
        (snapshot_dir / rel).write_bytes(b"stub")

    class FakeModule:
        def __init__(self, name):
            self.name = name
            self.to_calls = []

        def to(self, device):
            self.to_calls.append(device)
            return self

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                embed_tokens=FakeModule("embed"),
                layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")],
                rotary_emb=FakeModule("rotary"),
                norm=FakeModule("norm"),
            )
            self.lm_head = FakeModule("lm_head")
            self.tie_weights_called = 0

        def tie_weights(self):
            self.tie_weights_called += 1

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            assert model_id == str(snapshot_dir)
            return types.SimpleNamespace(num_hidden_layers=3)

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            assert cfg.num_hidden_layers == 3
            assert torch_dtype == "bf16"
            return FakeModel()

    class EmptyWeights:
        def __init__(self):
            self.entered = 0
            self.exited = 0

        def __call__(self):
            return self

        def __enter__(self):
            self.entered += 1
            return None

        def __exit__(self, exc_type, exc, tb):
            self.exited += 1
            return False

    init_empty_weights = EmptyWeights()
    set_calls = []

    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
        set_calls.append((tensor_name, device, value, dtype))

    tensors = {
        "shard-1.safetensors": {
            "model.embed_tokens.weight": "embed",
            "model.layers.0.self_attn.q_proj.weight": "layer0",
        },
        "shard-2.safetensors": {
            "model.layers.1.self_attn.q_proj.weight": "layer1",
        },
        "shard-3.safetensors": {
            "model.layers.2.self_attn.q_proj.weight": "layer2",
            "model.norm.weight": "norm",
            "lm_head.weight": "lm_head",
        },
    }

    class FakeSafeOpen:
        def __init__(self, filename, framework, device):
            assert framework == "pt"
            assert device == "cpu"
            self.filename = Path(filename).name

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def get_tensor(self, tensor_name):
            return tensors[self.filename][tensor_name]

    model = _load_partial_model_from_snapshot(
        AutoConfigStub,
        AutoModelStub,
        types.SimpleNamespace(),
        str(snapshot_dir),
        1,
        1,
        "bf16",
        "cpu:0",
        init_empty_weights_fn=init_empty_weights,
        set_tensor_fn=fake_set_tensor,
        safe_open_fn=FakeSafeOpen,
    )

    assert init_empty_weights.entered == 1
    assert init_empty_weights.exited == 1
    assert model.tie_weights_called == 1
    assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"]
    assert model.model.layers[1].to_calls == ["cpu:0"]
    assert model.model.layers[0].to_calls == []
    assert model.model.layers[2].to_calls == []
    assert model.model.embed_tokens.to_calls == []
    assert model.model.norm.to_calls == []
    assert model.lm_head.to_calls == []
    assert model.model.rotary_emb.to_calls == ["cpu:0"]


def test_partial_snapshot_loader_requires_known_layer_count(tmp_path):
    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
        "weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"}
    }))
    (snapshot_dir / "shard.safetensors").write_bytes(b"stub")

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id):
            return types.SimpleNamespace()

    class AutoModelStub:
        @staticmethod
        def from_config(cfg, torch_dtype=None):
            raise AssertionError("from_config should not run without a known layer count")

    class UnusedContext:
        def __enter__(self):
            return None

        def __exit__(self, exc_type, exc, tb):
            return False

    with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"):
        _load_partial_model_from_snapshot(
            AutoConfigStub,
            AutoModelStub,
            types.SimpleNamespace(),
            str(snapshot_dir),
            0,
            0,
            "bf16",
            "cpu:0",
            init_empty_weights_fn=lambda: UnusedContext(),
            set_tensor_fn=lambda *args, **kwargs: None,
            safe_open_fn=lambda *args, **kwargs: None,
        )


def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch):
    import meshnet_node.model_backend as backend

    snapshot_dir = tmp_path / "snapshot"
    snapshot_dir.mkdir()
    (snapshot_dir / "config.json").write_text("{}")
    (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')

    class FakeModel:
        def __init__(self):
            self.model = types.SimpleNamespace(
                layers=[object(), object(), object()],
                embed_tokens=object(),
            )
            self.config = types.SimpleNamespace(hidden_size=8)
            self.eval_called = 0

        def eval(self):
            self.eval_called += 1

    fake_model = FakeModel()
    partial_calls = []

    class AutoConfigStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16"))

    class AutoModelStub:
        @staticmethod
        def from_pretrained(*args, **kwargs):
            raise AssertionError("full model load should not run for partial local shards")

    class AutoTokenizerStub:
        @staticmethod
        def from_pretrained(model_id, cache_dir=None):
            assert model_id == str(snapshot_dir)
            return types.SimpleNamespace()

    monkeypatch.setitem(
        sys.modules,
        "torch",
        types.SimpleNamespace(
            cuda=types.SimpleNamespace(is_available=lambda: False),
            device=lambda value: value,
            bfloat16="bf16",
        ),
    )
    monkeypatch.setitem(
        sys.modules,
        "transformers",
        types.SimpleNamespace(
            AutoConfig=AutoConfigStub,
            AutoModelForCausalLM=AutoModelStub,
            AutoTokenizer=AutoTokenizerStub,
        ),
    )
    monkeypatch.setattr(
        backend,
        "_load_partial_model_from_snapshot",
        lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model,
    )

    shard = TorchModelShard(
        "repo/model",
        1,
        1,
        quantization="auto",
        cache_dir=snapshot_dir,
    )

    assert len(partial_calls) == 1
    assert shard.model is fake_model
    assert fake_model.eval_called == 1
    assert shard.total_layers == 3
    assert shard.is_head is False
    assert shard.is_tail is False


@pytest.mark.integration
def test_two_node_gpt2_completion_is_deterministic():
    if os.environ.get("CI"):
        pytest.skip("GPT-2 integration test is skipped in CI")
    torch = pytest.importorskip("torch")
    pytest.importorskip("transformers")
    pytest.importorskip("safetensors")
    pytest.importorskip("accelerate")
    pytest.importorskip("bitsandbytes")
    if not torch.cuda.is_available():
        pytest.skip("GPT-2 integration test requires a CUDA GPU")

    head = TorchNodeServer(
        model_id="openai-community/gpt2",
        shard_start=0,
        shard_end=6,
        quantization="bfloat16",
    )
    tail = TorchNodeServer(
        model_id="openai-community/gpt2",
        shard_start=6,
        shard_end=12,
        quantization="bfloat16",
    )
    head_port = head.start()
    tail_port = tail.start()
    try:
        prompt_req = urllib.request.Request(
            f"http://127.0.0.1:{head_port}/forward",
            data=json.dumps({"prompt": "The capital of France is"}).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(prompt_req, timeout=60) as resp:
            activation = resp.read()
            head_headers = resp.headers

        tail_req = urllib.request.Request(
            f"http://127.0.0.1:{tail_port}/forward",
            data=activation,
            headers={
                "Content-Type": "application/octet-stream",
                "X-Meshnet-Shape": head_headers["X-Meshnet-Shape"],
                "X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"],
                "X-Meshnet-Session": "gpt2-session",
                "X-Meshnet-Chunk-Index": "0",
                "X-Meshnet-Chunk-Total": "1",
                "X-Meshnet-Hop-Index": "1",
                "X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"],
                "X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"],
            },
            method="POST",
        )
        with urllib.request.urlopen(tail_req, timeout=60) as resp:
            body = json.loads(resp.read())

        assert body["text"].strip()
        assert body["text"] == " Paris"
    finally:
        head.stop()
        tail.stop()