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,
    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):
        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):
        assert len(body) == 1 * 6 * 8 * 2
        return " Paris"


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


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