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