feat: add real PyTorch model backend

tests/test_real_model_backend.py

@@ -0,0 +1,208 @@
+"""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()