feat: checkpoint distributed gguf runtime stories

tests/test_gguf_backend.py

@@ -0,0 +1,186 @@
+"""Tests for the GGUF backend adapter and recipe-gated startup seam."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+from meshnet_node.gguf_backend import GgufNodeBackend, build_gguf_backend
+from meshnet_node.model_backend import TailTokenResult, TensorPayload
+from meshnet_node.recipe_manifest import DEFAULT_RECIPE_ID, load_recipe_manifest
+from meshnet_node.startup import _gguf_backend_for_recipe
+
+
+class _RecordingTransport:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, tuple, dict]] = []
+
+    def encode_prompt(self, prompt: str, session_id: str | None = None):
+        self.calls.append(("encode_prompt", (prompt, session_id), {}))
+        return TensorPayload(
+            body=b"\x00" * 16,
+            shape=[1, 2, 4],
+            attention_mask_header=None,
+            position_ids_header=None,
+        )
+
+    def encode_next_token(self, token_id: int, session_id: str):
+        self.calls.append(("encode_next_token", (token_id, session_id), {}))
+        return TensorPayload(
+            body=b"\x00" * 8,
+            shape=[1, 1, 4],
+            attention_mask_header=None,
+            position_ids_header=None,
+            past_len=2,
+        )
+
+    def forward_bytes(
+        self,
+        body: bytes,
+        shape: list[int],
+        attention_mask_header: str | None,
+        position_ids_header: str | None,
+        *,
+        start_layer: int | None = None,
+        session_id: str | None = None,
+        cache_mode: str | None = None,
+        past_len: int | None = None,
+    ):
+        self.calls.append(
+            (
+                "forward_bytes",
+                (body, tuple(shape), attention_mask_header, position_ids_header),
+                {
+                    "start_layer": start_layer,
+                    "session_id": session_id,
+                    "cache_mode": cache_mode,
+                    "past_len": past_len,
+                },
+            )
+        )
+        if cache_mode == "decode":
+            return TailTokenResult(text=" done", token_id=17)
+        return TensorPayload(
+            body=b"\x00" * 16,
+            shape=[1, 2, 4],
+            attention_mask_header=attention_mask_header,
+            position_ids_header=position_ids_header,
+            past_len=past_len,
+        )
+
+    def decode_tail_token(self, hidden_states):
+        self.calls.append(("decode_tail_token", (hidden_states.shape,), {}))
+        return TailTokenResult(text=" tail", token_id=19)
+
+    def generate_text(self, messages, max_new_tokens=5120, temperature=1.0, top_p=1.0):
+        self.calls.append(("generate_text", (tuple(messages), max_new_tokens, temperature, top_p), {}))
+        return "ok"
+
+    def generate_text_streaming(self, messages, max_new_tokens=5120, temperature=1.0, top_p=1.0):
+        self.calls.append(("generate_text_streaming", (tuple(messages), max_new_tokens, temperature, top_p), {}))
+        yield "ok"
+
+    def count_prompt_tokens(self, messages):
+        self.calls.append(("count_prompt_tokens", (tuple(messages),), {}))
+        return 3
+
+    def count_text_tokens(self, text):
+        self.calls.append(("count_text_tokens", (text,), {}))
+        return 2
+
+    def eos_token_ids(self):
+        self.calls.append(("eos_token_ids", (), {}))
+        return [19]
+
+    def release_session(self, session_id: str) -> None:
+        self.calls.append(("release_session", (session_id,), {}))
+
+
+def test_build_gguf_backend_delegates_to_transport():
+    transport = _RecordingTransport()
+    backend = build_gguf_backend(
+        model_id="meshnet/native-model",
+        shard_start=0,
+        shard_end=1,
+        quantization="bfloat16",
+        transport=transport,
+        total_layers=2,
+        device_type="cpu",
+    )
+
+    assert isinstance(backend, GgufNodeBackend)
+    assert backend.backend_id == "llama.cpp"
+    assert backend.is_head is True
+    assert backend.is_tail is True
+    assert backend.model.config.to_dict()["architecture_adapter"] == "dense-llama"
+    assert backend.loaded_tensor_names[0] == "blk.0.weight"
+
+    prompt = backend.encode_prompt("hello", session_id="session-1")
+    assert prompt.shape == [1, 2, 4]
+
+    decode = backend.forward_bytes(
+        b"\x00" * 16,
+        [1, 2, 4],
+        None,
+        None,
+        session_id="session-1",
+        cache_mode="decode",
+        past_len=2,
+    )
+    assert isinstance(decode, TailTokenResult)
+    assert decode.token_id == 17
+
+    backend.release_session("session-1")
+
+    assert [call[0] for call in transport.calls] == [
+        "encode_prompt",
+        "forward_bytes",
+        "release_session",
+    ]
+    assert transport.calls[0][1] == ("hello", "session-1")
+    assert transport.calls[1][2]["cache_mode"] == "decode"
+    assert transport.calls[1][2]["past_len"] == 2
+
+
+def test_recipe_gates_native_backend_selection(monkeypatch):
+    manifest = load_recipe_manifest()
+    torch_recipe = manifest.require(DEFAULT_RECIPE_ID)
+    native_recipe = manifest.require("llama-cpp-native")
+
+    sentinel_backend = object()
+    calls: list[dict] = []
+
+    def fake_build_gguf_backend(**kwargs):
+        calls.append(kwargs)
+        return sentinel_backend
+
+    monkeypatch.setattr(
+        "meshnet_node.startup.build_gguf_backend",
+        fake_build_gguf_backend,
+    )
+
+    assert _gguf_backend_for_recipe(
+        torch_recipe,
+        model_id="meshnet/native-model",
+        shard_start=0,
+        shard_end=1,
+        quantization="bfloat16",
+        total_layers=2,
+        device="cpu",
+    ) is None
+
+    backend = _gguf_backend_for_recipe(
+        native_recipe,
+        model_id="meshnet/native-model",
+        shard_start=0,
+        shard_end=1,
+        quantization="bfloat16",
+        total_layers=2,
+        device="cpu",
+    )
+
+    assert backend is sentinel_backend
+    assert calls[0]["model_id"] == "meshnet/native-model"
+    assert calls[0]["shard_start"] == 0
+    assert calls[0]["shard_end"] == 1
+    assert calls[0]["quantization"] == "bfloat16"
+    assert calls[0]["total_layers"] == 2