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 types
import urllib.request

import pytest

from meshnet_node.model_backend import (
    InsufficientVRAMError,
    TensorPayload,
    _call_layer,
    _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


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():
    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"},
            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()


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"


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