inference working

tests/test_real_model_backend.py

@@ -12,6 +12,9 @@ import pytest
 from meshnet_node.model_backend import (
     InsufficientVRAMError,
     TensorPayload,
+    _call_layer,
+    _decoder_attention_mask,
+    _int_tensor_header,
     build_quantization_config,
     validate_quantization,
 )
@@ -52,6 +55,32 @@ class _FakeTailBackend(_FakeBackend):
         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"
@@ -145,6 +174,65 @@ def test_tail_forward_returns_text_completion_from_binary_activations():
         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"):