feat: checkpoint distributed gguf runtime stories
This commit is contained in:
488
tests/test_boundary_adapter.py
Normal file
488
tests/test_boundary_adapter.py
Normal file
@@ -0,0 +1,488 @@
|
||||
"""Architecture-defined boundary input/output and dense-Llama parity (DGR-006).
|
||||
|
||||
These tests prove the boundary contract with a *pure-numpy* dense-Llama reference
|
||||
model: no download, no GPU, no torch, no API credit. The reference implements the
|
||||
same ``ShardComputation`` duck type the real llama.cpp/PyTorch backends expose, so
|
||||
whole-model execution and a two-range (or three-range) split are the exact same
|
||||
arithmetic applied to the exact same float32 residual stream. Splitting the layer
|
||||
stack at a seam and shipping the *unnormalized* residual bundle across a simulated
|
||||
process boundary must reproduce the whole-model tokens bit-for-bit.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from meshnet_node.boundary_adapter import (
|
||||
BOUNDARY_SCHEMA_VERSION,
|
||||
BoundaryAdapter,
|
||||
BoundaryBundle,
|
||||
BoundaryContractError,
|
||||
SamplingContract,
|
||||
ShardRole,
|
||||
TailOutput,
|
||||
UncertifiedArchitectureError,
|
||||
certified_architecture,
|
||||
is_certified_architecture,
|
||||
role_for_range,
|
||||
)
|
||||
|
||||
# Documented parity tolerance. The split path applies the identical layer
|
||||
# functions in the identical order to the identical float32 arrays, so the
|
||||
# residual seam is bit-exact in practice; the tolerance is a conservative guard.
|
||||
PARITY_ATOL = 1e-6
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Pure-numpy dense-Llama reference model (test fixture, not production).
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class _ReferenceDenseLlama:
|
||||
"""A tiny deterministic dense-Llama: RMSNorm, RoPE attention, SwiGLU MLP."""
|
||||
|
||||
architecture_adapter = "dense-llama"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
vocab: int = 48,
|
||||
hidden: int = 32,
|
||||
n_layers: int = 6,
|
||||
n_heads: int = 4,
|
||||
intermediate: int = 64,
|
||||
rms_eps: float = 1e-6,
|
||||
rope_theta: float = 10000.0,
|
||||
seed: int = 20260715,
|
||||
) -> None:
|
||||
assert hidden % n_heads == 0
|
||||
self.vocab = vocab
|
||||
self.hidden = hidden
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.head_dim = hidden // n_heads
|
||||
assert self.head_dim % 2 == 0
|
||||
self.rms_eps = rms_eps
|
||||
self.rope_theta = rope_theta
|
||||
|
||||
rng = np.random.default_rng(seed)
|
||||
|
||||
def w(*shape: int) -> np.ndarray:
|
||||
return (rng.standard_normal(shape) * 0.08).astype(np.float32)
|
||||
|
||||
self.embed = w(vocab, hidden)
|
||||
self.layers = []
|
||||
for _ in range(n_layers):
|
||||
self.layers.append(
|
||||
{
|
||||
"in_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
||||
"q": w(hidden, hidden),
|
||||
"k": w(hidden, hidden),
|
||||
"v": w(hidden, hidden),
|
||||
"o": w(hidden, hidden),
|
||||
"post_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
||||
"gate": w(intermediate, hidden),
|
||||
"up": w(intermediate, hidden),
|
||||
"down": w(hidden, intermediate),
|
||||
}
|
||||
)
|
||||
self.final_ln = (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32)
|
||||
self.lm_head_w = w(vocab, hidden)
|
||||
|
||||
inv_freq = 1.0 / (
|
||||
rope_theta ** (np.arange(0, self.head_dim, 2, dtype=np.float32) / self.head_dim)
|
||||
)
|
||||
self.inv_freq = inv_freq.astype(np.float32)
|
||||
|
||||
# -- primitive ops -----------------------------------------------------
|
||||
def _rmsnorm(self, x: np.ndarray, weight: np.ndarray) -> np.ndarray:
|
||||
variance = np.mean(x.astype(np.float32) ** 2, axis=-1, keepdims=True)
|
||||
normed = x / np.sqrt(variance + self.rms_eps)
|
||||
return (normed * weight).astype(np.float32)
|
||||
|
||||
def _rope(self, positions: np.ndarray):
|
||||
# positions: (batch, seq) -> cos/sin: (batch, seq, head_dim)
|
||||
angles = positions[..., None].astype(np.float32) * self.inv_freq[None, None, :]
|
||||
emb = np.concatenate([angles, angles], axis=-1)
|
||||
return np.cos(emb).astype(np.float32), np.sin(emb).astype(np.float32)
|
||||
|
||||
@staticmethod
|
||||
def _rotate_half(x: np.ndarray) -> np.ndarray:
|
||||
half = x.shape[-1] // 2
|
||||
return np.concatenate([-x[..., half:], x[..., :half]], axis=-1)
|
||||
|
||||
def _apply_rope(self, t: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
|
||||
# t: (batch, n_heads, seq, head_dim); cos/sin: (batch, seq, head_dim)
|
||||
cos = cos[:, None, :, :]
|
||||
sin = sin[:, None, :, :]
|
||||
return t * cos + self._rotate_half(t) * sin
|
||||
|
||||
def _attention(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
|
||||
batch, seq, _ = x.shape
|
||||
q = (x @ layer["q"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
||||
k = (x @ layer["k"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
||||
v = (x @ layer["v"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
||||
q = q.transpose(0, 2, 1, 3)
|
||||
k = k.transpose(0, 2, 1, 3)
|
||||
v = v.transpose(0, 2, 1, 3)
|
||||
cos, sin = self._rope(positions)
|
||||
q = self._apply_rope(q, cos, sin)
|
||||
k = self._apply_rope(k, cos, sin)
|
||||
scores = (q @ k.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
|
||||
causal = np.triu(np.full((seq, seq), -1e30, dtype=np.float32), k=1)
|
||||
scores = scores + causal[None, None, :, :]
|
||||
scores = scores - scores.max(axis=-1, keepdims=True)
|
||||
weights = np.exp(scores)
|
||||
weights = weights / weights.sum(axis=-1, keepdims=True)
|
||||
out = weights @ v
|
||||
out = out.transpose(0, 2, 1, 3).reshape(batch, seq, self.hidden)
|
||||
return (out @ layer["o"].T).astype(np.float32)
|
||||
|
||||
def _mlp(self, x: np.ndarray, layer: dict) -> np.ndarray:
|
||||
gate = x @ layer["gate"].T
|
||||
up = x @ layer["up"].T
|
||||
silu = gate * (1.0 / (1.0 + np.exp(-gate)))
|
||||
return ((silu * up) @ layer["down"].T).astype(np.float32)
|
||||
|
||||
def _run_layer(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
|
||||
h = x + self._attention(self._rmsnorm(x, layer["in_ln"]), layer, positions)
|
||||
h = h + self._mlp(self._rmsnorm(h, layer["post_ln"]), layer)
|
||||
return h.astype(np.float32)
|
||||
|
||||
|
||||
class _ReferenceShard:
|
||||
"""A contiguous inclusive layer range of the reference model.
|
||||
|
||||
Satisfies the ``ShardComputation`` duck type used by ``BoundaryAdapter``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: _ReferenceDenseLlama,
|
||||
start_layer: int,
|
||||
end_layer: int,
|
||||
*,
|
||||
architecture_adapter: str | None = None,
|
||||
) -> None:
|
||||
self._model = model
|
||||
self.start_layer = start_layer
|
||||
self.end_layer = end_layer
|
||||
self.total_layers = model.n_layers
|
||||
self.architecture_adapter = architecture_adapter or model.architecture_adapter
|
||||
|
||||
def embed_tokens(self, token_ids: np.ndarray) -> np.ndarray:
|
||||
return self._model.embed[np.asarray(token_ids)]
|
||||
|
||||
def run_layers(self, hidden: np.ndarray, *, positions: np.ndarray) -> np.ndarray:
|
||||
h = np.asarray(hidden, dtype=np.float32)
|
||||
for idx in range(self.start_layer, self.end_layer + 1):
|
||||
h = self._model._run_layer(h, self._model.layers[idx], positions)
|
||||
return h
|
||||
|
||||
def final_norm(self, hidden: np.ndarray) -> np.ndarray:
|
||||
return self._model._rmsnorm(np.asarray(hidden, dtype=np.float32), self._model.final_ln)
|
||||
|
||||
def lm_head(self, hidden: np.ndarray) -> np.ndarray:
|
||||
return np.asarray(hidden, dtype=np.float32) @ self._model.lm_head_w.T
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Whole-model and split reference drivers.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _whole_model_next_token(model: _ReferenceDenseLlama, token_ids: list[int]) -> TailOutput:
|
||||
shard = _ReferenceShard(model, 0, model.n_layers - 1)
|
||||
adapter = BoundaryAdapter(shard)
|
||||
result = adapter.forward(token_ids=np.asarray(token_ids)[None, :])
|
||||
assert isinstance(result, TailOutput)
|
||||
return result
|
||||
|
||||
|
||||
def _split_next_token(
|
||||
model: _ReferenceDenseLlama,
|
||||
token_ids: list[int],
|
||||
cut_points: list[int],
|
||||
*,
|
||||
through_wire: bool = True,
|
||||
) -> TailOutput:
|
||||
"""Run the model as N contiguous ranges, shipping the bundle across each seam.
|
||||
|
||||
``cut_points`` are the last (inclusive) layer of each non-final range.
|
||||
"""
|
||||
bounds = _ranges_from_cuts(cut_points, model.n_layers)
|
||||
boundary: BoundaryBundle | None = None
|
||||
result: BoundaryBundle | TailOutput | None = None
|
||||
for i, (start, end) in enumerate(bounds):
|
||||
shard = _ReferenceShard(model, start, end)
|
||||
adapter = BoundaryAdapter(shard)
|
||||
if i == 0:
|
||||
result = adapter.forward(token_ids=np.asarray(token_ids)[None, :])
|
||||
else:
|
||||
assert isinstance(boundary, BoundaryBundle)
|
||||
incoming = BoundaryBundle.unpack(boundary.pack()) if through_wire else boundary
|
||||
result = adapter.forward(boundary=incoming)
|
||||
if isinstance(result, BoundaryBundle):
|
||||
boundary = result
|
||||
assert isinstance(result, TailOutput)
|
||||
return result
|
||||
|
||||
|
||||
def _ranges_from_cuts(cut_points: list[int], n_layers: int) -> list[tuple[int, int]]:
|
||||
bounds: list[tuple[int, int]] = []
|
||||
start = 0
|
||||
for cut in cut_points:
|
||||
bounds.append((start, cut))
|
||||
start = cut + 1
|
||||
bounds.append((start, n_layers - 1))
|
||||
return bounds
|
||||
|
||||
|
||||
def _greedy_generate(next_token_fn, prompt: list[int], n_new: int) -> list[int]:
|
||||
tokens = list(prompt)
|
||||
generated: list[int] = []
|
||||
for _ in range(n_new):
|
||||
out = next_token_fn(tokens)
|
||||
tokens.append(out.token_id)
|
||||
generated.append(out.token_id)
|
||||
return generated
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Certification / fail-closed.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_dense_llama_and_aliases_are_certified():
|
||||
"Dense Llama-family identifiers all resolve to the one certified adapter.\n\nTags: node, boundary"
|
||||
for name in ("dense-llama", "llama", "LlamaForCausalLM", "LlamaModel"):
|
||||
boundary = certified_architecture(name)
|
||||
assert boundary.adapter == "dense-llama"
|
||||
assert boundary.boundary_tensor_name == "residual_stream"
|
||||
assert is_certified_architecture(name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["qwen3", "qwen3-moe", "mixtral", "gpt2", "", None, 123])
|
||||
def test_uncertified_architectures_fail_closed(name):
|
||||
"Uncertified architectures raise instead of guessing a tensor layout.\n\nTags: node, boundary"
|
||||
assert not is_certified_architecture(name)
|
||||
with pytest.raises(UncertifiedArchitectureError):
|
||||
certified_architecture(name)
|
||||
|
||||
|
||||
def test_adapter_construction_fails_closed_for_uncertified_backend():
|
||||
"Building the adapter over an uncertified computation fails closed.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
shard = _ReferenceShard(model, 0, 2, architecture_adapter="qwen3-moe")
|
||||
with pytest.raises(UncertifiedArchitectureError):
|
||||
BoundaryAdapter(shard)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Roles.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_role_classification():
|
||||
"Range endpoints map to head/middle/tail/full roles.\n\nTags: node, boundary"
|
||||
assert role_for_range(0, 2, 6) is ShardRole.HEAD
|
||||
assert role_for_range(2, 3, 6) is ShardRole.MIDDLE
|
||||
assert role_for_range(4, 5, 6) is ShardRole.TAIL
|
||||
assert role_for_range(0, 5, 6) is ShardRole.FULL
|
||||
assert ShardRole.HEAD.owns_embedding and not ShardRole.HEAD.owns_final_head
|
||||
assert ShardRole.TAIL.owns_final_head and not ShardRole.TAIL.owns_embedding
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Input-side contract.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_head_accepts_token_ids_and_owns_embedding():
|
||||
"The head embeds token IDs and refuses an upstream boundary bundle.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
||||
out = head.forward(token_ids=[1, 2, 3])
|
||||
assert isinstance(out, BoundaryBundle)
|
||||
|
||||
# Head owns embedding: a residual bundle from upstream is a contract error.
|
||||
bundle = out
|
||||
with pytest.raises(BoundaryContractError, match="head owns token embedding"):
|
||||
head.forward(boundary=bundle)
|
||||
|
||||
|
||||
def test_middle_and_tail_bypass_embedding_and_require_the_bundle():
|
||||
"Middle/tail Shards reject token IDs and demand the named boundary bundle.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
tail = BoundaryAdapter(_ReferenceShard(model, 3, 5))
|
||||
with pytest.raises(BoundaryContractError, match="bypass token embedding"):
|
||||
tail.forward(token_ids=[1, 2, 3])
|
||||
with pytest.raises(BoundaryContractError, match="must receive the named boundary bundle"):
|
||||
tail.forward()
|
||||
|
||||
|
||||
def test_boundary_seam_layer_mismatch_is_rejected():
|
||||
"A bundle handed to the wrong range (seam layer mismatch) is rejected.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
||||
bundle = head.forward(token_ids=[1, 2, 3])
|
||||
assert isinstance(bundle, BoundaryBundle)
|
||||
assert bundle.next_layer == 3
|
||||
|
||||
# A range that starts at layer 4 must not accept a bundle cut at layer 3.
|
||||
wrong = BoundaryAdapter(_ReferenceShard(model, 4, 5))
|
||||
with pytest.raises(BoundaryContractError, match="starts at layer 4"):
|
||||
wrong.forward(boundary=bundle)
|
||||
|
||||
|
||||
def test_normalized_bundle_is_rejected():
|
||||
"A normalized residual is not the architecture-defined boundary.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
||||
bundle = head.forward(token_ids=[1, 2, 3])
|
||||
assert isinstance(bundle, BoundaryBundle)
|
||||
normalized = BoundaryBundle(
|
||||
architecture_adapter=bundle.architecture_adapter,
|
||||
schema_version=bundle.schema_version,
|
||||
tensor_name=bundle.tensor_name,
|
||||
residual=bundle.residual,
|
||||
positions=bundle.positions,
|
||||
next_layer=bundle.next_layer,
|
||||
normalized=True,
|
||||
)
|
||||
tail = BoundaryAdapter(_ReferenceShard(model, 3, 5))
|
||||
with pytest.raises(BoundaryContractError, match="UNNORMALIZED"):
|
||||
tail.forward(boundary=normalized)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Output-side contract.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_non_tail_emits_unnormalized_full_row_boundary():
|
||||
"A non-tail Shard emits the unnormalized residual with every position row.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
tokens = [3, 7, 1, 9, 2]
|
||||
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
||||
bundle = head.forward(token_ids=tokens)
|
||||
assert isinstance(bundle, BoundaryBundle)
|
||||
assert bundle.normalized is False
|
||||
assert bundle.tensor_name == "residual_stream"
|
||||
assert bundle.schema_version == BOUNDARY_SCHEMA_VERSION
|
||||
assert bundle.next_layer == 3
|
||||
# No tail-only row pruning: all sequence positions are forwarded.
|
||||
assert bundle.residual.shape == (1, len(tokens), model.hidden)
|
||||
assert bundle.positions.shape == (1, len(tokens))
|
||||
|
||||
# The emitted residual must be exactly the whole model's residual after layer 2
|
||||
# (i.e. before any final norm) — prove it is NOT normalized.
|
||||
positions = np.arange(len(tokens))[None, :]
|
||||
hidden = model.embed[np.asarray(tokens)][None, :]
|
||||
for idx in range(0, 3):
|
||||
hidden = model._run_layer(hidden, model.layers[idx], positions)
|
||||
assert np.allclose(bundle.residual, hidden, atol=0)
|
||||
assert not np.allclose(bundle.residual, model._rmsnorm(hidden, model.final_ln))
|
||||
|
||||
|
||||
def test_tail_emits_pruned_logits_through_the_sampling_contract():
|
||||
"The tail prunes to the final row and samples through an explicit contract.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
out = _whole_model_next_token(model, [4, 8, 15, 16, 23])
|
||||
assert isinstance(out, TailOutput)
|
||||
assert out.logits.shape == (1, model.vocab) # tail-only row pruning to last row
|
||||
assert out.sampling.mode == "greedy"
|
||||
assert 0 <= out.token_id < model.vocab
|
||||
assert out.token_id == int(np.argmax(out.logits[0]))
|
||||
|
||||
|
||||
def test_sampling_contract_rejects_uncertified_modes():
|
||||
"Only the certified greedy sampling mode is accepted.\n\nTags: node, boundary"
|
||||
with pytest.raises(BoundaryContractError):
|
||||
SamplingContract(mode="top_p")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# The core parity gate.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_two_range_prefill_parity_matches_whole_model():
|
||||
"Whole-model vs two-range prefill produce the same next-token logits and token.\n\nTags: node, boundary, parity"
|
||||
model = _ReferenceDenseLlama()
|
||||
prompt = [5, 12, 3, 41, 7, 19, 2, 33]
|
||||
|
||||
whole = _whole_model_next_token(model, prompt)
|
||||
split = _split_next_token(model, prompt, cut_points=[2])
|
||||
|
||||
assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
|
||||
assert whole.token_id == split.token_id
|
||||
|
||||
|
||||
def test_three_range_prefill_parity_exercises_the_middle_role():
|
||||
"A head/middle/tail split reproduces whole-model prefill through two seams.\n\nTags: node, boundary, parity"
|
||||
model = _ReferenceDenseLlama()
|
||||
prompt = [9, 1, 44, 6, 30, 11]
|
||||
|
||||
whole = _whole_model_next_token(model, prompt)
|
||||
split = _split_next_token(model, prompt, cut_points=[1, 3])
|
||||
|
||||
assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
|
||||
assert whole.token_id == split.token_id
|
||||
|
||||
|
||||
def test_two_range_greedy_decode_parity_matches_whole_model():
|
||||
"Whole-model vs two-range greedy decode produce identical token sequences.\n\nTags: node, boundary, parity"
|
||||
model = _ReferenceDenseLlama()
|
||||
prompt = [2, 17, 8, 25]
|
||||
n_new = 12
|
||||
|
||||
whole_tokens = _greedy_generate(
|
||||
lambda toks: _whole_model_next_token(model, toks), prompt, n_new
|
||||
)
|
||||
split_tokens = _greedy_generate(
|
||||
lambda toks: _split_next_token(model, toks, cut_points=[2]), prompt, n_new
|
||||
)
|
||||
|
||||
assert whole_tokens == split_tokens
|
||||
assert len(whole_tokens) == n_new
|
||||
|
||||
|
||||
def test_boundary_bundle_wire_round_trip_is_exact():
|
||||
"Packing and unpacking the boundary bundle reconstructs the exact arrays.\n\nTags: node, boundary"
|
||||
model = _ReferenceDenseLlama()
|
||||
head = BoundaryAdapter(_ReferenceShard(model, 0, 2))
|
||||
bundle = head.forward(token_ids=[1, 2, 3, 4])
|
||||
assert isinstance(bundle, BoundaryBundle)
|
||||
|
||||
restored = BoundaryBundle.unpack(bundle.pack())
|
||||
assert np.array_equal(restored.residual, bundle.residual)
|
||||
assert np.array_equal(restored.positions, bundle.positions)
|
||||
assert restored.next_layer == bundle.next_layer
|
||||
assert restored.architecture_adapter == bundle.architecture_adapter
|
||||
|
||||
fields = bundle.named_tensor_fields()
|
||||
assert fields["name"] == "residual_stream"
|
||||
assert fields["shape"] == [1, 4, model.hidden]
|
||||
assert fields["byte_order"] in ("little", "big")
|
||||
|
||||
|
||||
def test_alias_architecture_still_parity_matches():
|
||||
"A Shard advertised as 'llama' interoperates with the canonical adapter.\n\nTags: node, boundary, parity"
|
||||
model = _ReferenceDenseLlama()
|
||||
prompt = [7, 3, 22, 5]
|
||||
|
||||
whole = _whole_model_next_token(model, prompt)
|
||||
|
||||
# Head advertises 'LlamaForCausalLM', tail advertises 'llama'; both certify to
|
||||
# the same canonical adapter, so the seam contract still matches.
|
||||
head = BoundaryAdapter(_ReferenceShard(model, 0, 2, architecture_adapter="LlamaForCausalLM"))
|
||||
bundle = head.forward(token_ids=np.asarray(prompt)[None, :])
|
||||
assert isinstance(bundle, BoundaryBundle)
|
||||
tail = BoundaryAdapter(_ReferenceShard(model, 3, 5, architecture_adapter="llama"))
|
||||
split = tail.forward(boundary=BoundaryBundle.unpack(bundle.pack()))
|
||||
assert isinstance(split, TailOutput)
|
||||
|
||||
assert np.allclose(whole.logits, split.logits, atol=PARITY_ATOL)
|
||||
assert whole.token_id == split.token_id
|
||||
186
tests/test_gguf_backend.py
Normal file
186
tests/test_gguf_backend.py
Normal file
@@ -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
|
||||
88
tests/test_gguf_ownership.py
Normal file
88
tests/test_gguf_ownership.py
Normal file
@@ -0,0 +1,88 @@
|
||||
"""Dense-Llama GGUF ownership selection and introspection tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_node.gguf_ownership import (
|
||||
DenseLlamaShardOwnership,
|
||||
authoritative_dense_llama_ownership,
|
||||
infer_dense_llama_ownership,
|
||||
select_dense_llama_tensor_names,
|
||||
)
|
||||
|
||||
|
||||
def test_dense_llama_selection_only_picks_block_range_and_endpoints():
|
||||
"Dense-Llama selection keeps only the owned blocks plus the correct endpoints.\n\nTags: node, GGUF"
|
||||
tensor_inventory = {
|
||||
"token_embd.weight": 10_000,
|
||||
"blk.0.attn_q.weight": 1_000,
|
||||
"blk.0.ffn_down.weight": 1_000,
|
||||
"blk.1.attn_q.weight": 2_000,
|
||||
"blk.1.ffn_down.weight": 2_000,
|
||||
"blk.2.attn_q.weight": 3_000,
|
||||
"blk.2.ffn_down.weight": 3_000,
|
||||
"output_norm.weight": 256,
|
||||
"output.weight": 10_000,
|
||||
"rope.freqs": 128,
|
||||
}
|
||||
|
||||
selected = select_dense_llama_tensor_names(
|
||||
tensor_inventory,
|
||||
1,
|
||||
2,
|
||||
total_layers=3,
|
||||
)
|
||||
|
||||
assert selected == {
|
||||
"blk.1.attn_q.weight",
|
||||
"blk.1.ffn_down.weight",
|
||||
"blk.2.attn_q.weight",
|
||||
"blk.2.ffn_down.weight",
|
||||
"output_norm.weight",
|
||||
"output.weight",
|
||||
}
|
||||
|
||||
selected_bytes = sum(tensor_inventory[name] for name in selected)
|
||||
full_bytes = sum(tensor_inventory.values())
|
||||
assert selected_bytes == 20_256
|
||||
assert selected_bytes < full_bytes
|
||||
|
||||
|
||||
def test_dense_llama_loaded_range_is_authoritative_from_tensor_inventory():
|
||||
"The backend's loaded tensor inventory is the source of truth for range and ownership.\n\nTags: node, GGUF"
|
||||
|
||||
class Backend:
|
||||
loaded_tensor_names = (
|
||||
"token_embd.weight",
|
||||
"blk.4.attn_q.weight",
|
||||
"blk.5.ffn_down.weight",
|
||||
"output_norm.weight",
|
||||
"output.weight",
|
||||
)
|
||||
|
||||
ownership = authoritative_dense_llama_ownership(Backend(), selection=None)
|
||||
|
||||
assert isinstance(ownership, DenseLlamaShardOwnership)
|
||||
assert ownership.range == (4, 5)
|
||||
assert ownership.owns_embedding is True
|
||||
assert ownership.owns_final_head is True
|
||||
|
||||
|
||||
def test_derivative_slice_requires_source_and_slice_hashes():
|
||||
"Temporary derivative GGUF slices must carry hashes and cannot claim final semantics.\n\nTags: node, GGUF"
|
||||
with pytest.raises(ValueError, match="source and slice hashes"):
|
||||
infer_dense_llama_ownership(
|
||||
["blk.1.attn_q.weight"],
|
||||
derivative_slice=True,
|
||||
final_artifact_semantics=False,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="final artifacts"):
|
||||
infer_dense_llama_ownership(
|
||||
["blk.1.attn_q.weight"],
|
||||
source_artifact_hash="sha256:source",
|
||||
slice_artifact_hash="sha256:slice",
|
||||
derivative_slice=True,
|
||||
final_artifact_semantics=True,
|
||||
)
|
||||
769
tests/test_hot_kv_state.py
Normal file
769
tests/test_hot_kv_state.py
Normal file
@@ -0,0 +1,769 @@
|
||||
"""Isolated concurrent local Hot KV State (DGR-007).
|
||||
|
||||
These tests prove the KV/session manager with a *pure-numpy* KV-cached dense-Llama
|
||||
reference: no download, no GPU, no torch, no API credit. The reference implements
|
||||
the DGR-006 ``ShardComputation`` duck type plus ``run_layers_cached`` so cached
|
||||
prefill/decode over a per-session KV context reproduces the stateless whole-model
|
||||
tokens bit-for-bit. On top of that correctness core, the tests exercise the
|
||||
manager's lifecycle: owned-layer allocation, prefill/decode append, truncate,
|
||||
release, TTL/LRU eviction, explicit cache-miss responses, stale-epoch and
|
||||
incompatible-recipe rejection, four concurrent cross-talk-free sessions, and
|
||||
budget-bounded cancellation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from meshnet_node.boundary_adapter import BoundaryBundle, TailOutput
|
||||
from meshnet_node.hot_kv_state import (
|
||||
CacheMiss,
|
||||
CacheMissReason,
|
||||
HotKvStateConfig,
|
||||
HotKvStateManager,
|
||||
IncompatibleCacheRecipeError,
|
||||
KvBoundaryAdapter,
|
||||
KvBudgetExceededError,
|
||||
KvCacheMissError,
|
||||
KvCacheRecipe,
|
||||
LayerKvCache,
|
||||
StaleRouteEpochError,
|
||||
kv_recipe_for,
|
||||
)
|
||||
|
||||
PARITY_ATOL = 1e-6
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Pure-numpy KV-cached dense-Llama reference (test fixture, not production).
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class _KvDenseLlama:
|
||||
"""A tiny deterministic dense-Llama with both stateless and cached runners."""
|
||||
|
||||
architecture_adapter = "dense-llama"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
vocab: int = 48,
|
||||
hidden: int = 32,
|
||||
n_layers: int = 6,
|
||||
n_heads: int = 4,
|
||||
intermediate: int = 64,
|
||||
rms_eps: float = 1e-6,
|
||||
rope_theta: float = 10000.0,
|
||||
seed: int = 20260716,
|
||||
) -> None:
|
||||
assert hidden % n_heads == 0
|
||||
self.vocab = vocab
|
||||
self.hidden = hidden
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.head_dim = hidden // n_heads
|
||||
assert self.head_dim % 2 == 0
|
||||
self.rms_eps = rms_eps
|
||||
self.rope_theta = rope_theta
|
||||
|
||||
rng = np.random.default_rng(seed)
|
||||
|
||||
def w(*shape: int) -> np.ndarray:
|
||||
return (rng.standard_normal(shape) * 0.08).astype(np.float32)
|
||||
|
||||
self.embed = w(vocab, hidden)
|
||||
self.layers = []
|
||||
for _ in range(n_layers):
|
||||
self.layers.append(
|
||||
{
|
||||
"in_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
||||
"q": w(hidden, hidden),
|
||||
"k": w(hidden, hidden),
|
||||
"v": w(hidden, hidden),
|
||||
"o": w(hidden, hidden),
|
||||
"post_ln": (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32),
|
||||
"gate": w(intermediate, hidden),
|
||||
"up": w(intermediate, hidden),
|
||||
"down": w(hidden, intermediate),
|
||||
}
|
||||
)
|
||||
self.final_ln = (1.0 + rng.standard_normal(hidden) * 0.02).astype(np.float32)
|
||||
self.lm_head_w = w(vocab, hidden)
|
||||
|
||||
inv_freq = 1.0 / (
|
||||
rope_theta ** (np.arange(0, self.head_dim, 2, dtype=np.float32) / self.head_dim)
|
||||
)
|
||||
self.inv_freq = inv_freq.astype(np.float32)
|
||||
|
||||
# -- primitive ops -----------------------------------------------------
|
||||
def _rmsnorm(self, x: np.ndarray, weight: np.ndarray) -> np.ndarray:
|
||||
variance = np.mean(x.astype(np.float32) ** 2, axis=-1, keepdims=True)
|
||||
normed = x / np.sqrt(variance + self.rms_eps)
|
||||
return (normed * weight).astype(np.float32)
|
||||
|
||||
def _rope(self, positions: np.ndarray):
|
||||
angles = positions[..., None].astype(np.float32) * self.inv_freq[None, None, :]
|
||||
emb = np.concatenate([angles, angles], axis=-1)
|
||||
return np.cos(emb).astype(np.float32), np.sin(emb).astype(np.float32)
|
||||
|
||||
@staticmethod
|
||||
def _rotate_half(x: np.ndarray) -> np.ndarray:
|
||||
half = x.shape[-1] // 2
|
||||
return np.concatenate([-x[..., half:], x[..., :half]], axis=-1)
|
||||
|
||||
def _apply_rope(self, t: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
|
||||
cos = cos[:, None, :, :]
|
||||
sin = sin[:, None, :, :]
|
||||
return t * cos + self._rotate_half(t) * sin
|
||||
|
||||
def _project_qkv(self, normed: np.ndarray, layer: dict, positions: np.ndarray):
|
||||
batch, seq, _ = normed.shape
|
||||
q = (normed @ layer["q"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
||||
k = (normed @ layer["k"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
||||
v = (normed @ layer["v"].T).reshape(batch, seq, self.n_heads, self.head_dim)
|
||||
q = q.transpose(0, 2, 1, 3)
|
||||
k = k.transpose(0, 2, 1, 3)
|
||||
v = v.transpose(0, 2, 1, 3)
|
||||
cos, sin = self._rope(positions)
|
||||
q = self._apply_rope(q, cos, sin)
|
||||
k = self._apply_rope(k, cos, sin)
|
||||
return q, k, v
|
||||
|
||||
def _attend(
|
||||
self,
|
||||
q: np.ndarray,
|
||||
k_all: np.ndarray,
|
||||
v_all: np.ndarray,
|
||||
layer: dict,
|
||||
q_positions: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
batch, _, seq_new, _ = q.shape
|
||||
total = k_all.shape[2]
|
||||
scores = (q @ k_all.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
|
||||
# Causal mask by absolute position: keys are stored in absolute order
|
||||
# 0..total-1; query row i lives at absolute position q_positions[i].
|
||||
key_abs = np.arange(total, dtype=np.int64)
|
||||
q_abs = np.asarray(q_positions).reshape(seq_new).astype(np.int64)
|
||||
mask = np.where(key_abs[None, :] <= q_abs[:, None], 0.0, -1e30).astype(np.float32)
|
||||
scores = scores + mask[None, None, :, :]
|
||||
scores = scores - scores.max(axis=-1, keepdims=True)
|
||||
weights = np.exp(scores)
|
||||
weights = weights / weights.sum(axis=-1, keepdims=True)
|
||||
out = weights @ v_all
|
||||
out = out.transpose(0, 2, 1, 3).reshape(batch, seq_new, self.hidden)
|
||||
return (out @ layer["o"].T).astype(np.float32)
|
||||
|
||||
def _mlp(self, x: np.ndarray, layer: dict) -> np.ndarray:
|
||||
gate = x @ layer["gate"].T
|
||||
up = x @ layer["up"].T
|
||||
silu = gate * (1.0 / (1.0 + np.exp(-gate)))
|
||||
return ((silu * up) @ layer["down"].T).astype(np.float32)
|
||||
|
||||
# -- stateless whole-sequence layer (ground truth) ---------------------
|
||||
def _run_layer_stateless(self, x: np.ndarray, layer: dict, positions: np.ndarray) -> np.ndarray:
|
||||
normed = self._rmsnorm(x, layer["in_ln"])
|
||||
q, k, v = self._project_qkv(normed, layer, positions)
|
||||
attn = self._attend(q, k, v, layer, positions[0])
|
||||
h = x + attn
|
||||
h = h + self._mlp(self._rmsnorm(h, layer["post_ln"]), layer)
|
||||
return h.astype(np.float32)
|
||||
|
||||
def whole_model_next_token(self, token_ids: list[int]) -> int:
|
||||
positions = np.arange(len(token_ids))[None, :]
|
||||
h = self.embed[np.asarray(token_ids)][None, :]
|
||||
for idx in range(self.n_layers):
|
||||
h = self._run_layer_stateless(h, self.layers[idx], positions)
|
||||
h = self._rmsnorm(h[:, -1:, :], self.final_ln)
|
||||
logits = h @ self.lm_head_w.T
|
||||
return int(np.argmax(logits[0, -1]))
|
||||
|
||||
def stateless_greedy(self, prompt: list[int], n_new: int) -> list[int]:
|
||||
tokens = list(prompt)
|
||||
out: list[int] = []
|
||||
for _ in range(n_new):
|
||||
tok = self.whole_model_next_token(tokens)
|
||||
tokens.append(tok)
|
||||
out.append(tok)
|
||||
return out
|
||||
|
||||
|
||||
class _KvReferenceShard:
|
||||
"""A contiguous inclusive layer range with a KV-cached runner.
|
||||
|
||||
Satisfies the KV-aware ``ShardComputation`` duck type used by
|
||||
``KvBoundaryAdapter``: DGR-006 methods plus ``run_layers_cached`` and the KV
|
||||
geometry (``n_kv_heads`` / ``head_dim`` / ``kv_dtype``).
|
||||
"""
|
||||
|
||||
kv_dtype = "float32"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: _KvDenseLlama,
|
||||
start_layer: int,
|
||||
end_layer: int,
|
||||
*,
|
||||
architecture_adapter: str | None = None,
|
||||
) -> None:
|
||||
self._model = model
|
||||
self.start_layer = start_layer
|
||||
self.end_layer = end_layer
|
||||
self.total_layers = model.n_layers
|
||||
self.n_kv_heads = model.n_heads
|
||||
self.head_dim = model.head_dim
|
||||
self.architecture_adapter = architecture_adapter or model.architecture_adapter
|
||||
|
||||
def embed_tokens(self, token_ids: np.ndarray) -> np.ndarray:
|
||||
return self._model.embed[np.asarray(token_ids)]
|
||||
|
||||
def final_norm(self, hidden: np.ndarray) -> np.ndarray:
|
||||
return self._model._rmsnorm(np.asarray(hidden, dtype=np.float32), self._model.final_ln)
|
||||
|
||||
def lm_head(self, hidden: np.ndarray) -> np.ndarray:
|
||||
return np.asarray(hidden, dtype=np.float32) @ self._model.lm_head_w.T
|
||||
|
||||
def run_layers_cached(self, hidden, *, positions, past_kv):
|
||||
m = self._model
|
||||
x = np.asarray(hidden, dtype=np.float32)
|
||||
positions = np.asarray(positions)
|
||||
new_kv: dict[int, tuple[np.ndarray, np.ndarray]] = {}
|
||||
for idx in range(self.start_layer, self.end_layer + 1):
|
||||
layer = m.layers[idx]
|
||||
normed = m._rmsnorm(x, layer["in_ln"])
|
||||
q, k, v = m._project_qkv(normed, layer, positions)
|
||||
# Post-RoPE new K/V stored as (seq_new, n_heads, head_dim).
|
||||
new_k = k[0].transpose(1, 0, 2).copy()
|
||||
new_v = v[0].transpose(1, 0, 2).copy()
|
||||
cache = past_kv.get(idx)
|
||||
if cache is not None and cache.length > 0:
|
||||
past_k = cache.keys[None].transpose(0, 2, 1, 3)
|
||||
past_v = cache.values[None].transpose(0, 2, 1, 3)
|
||||
k_all = np.concatenate([past_k, k], axis=2)
|
||||
v_all = np.concatenate([past_v, v], axis=2)
|
||||
else:
|
||||
k_all, v_all = k, v
|
||||
attn = m._attend(q, k_all, v_all, layer, positions[0])
|
||||
h = x + attn
|
||||
x = h + m._mlp(m._rmsnorm(h, layer["post_ln"]), layer)
|
||||
x = x.astype(np.float32)
|
||||
new_kv[idx] = (new_k, new_v)
|
||||
return x, new_kv
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Helpers.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class _FakeClock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, delta: float) -> None:
|
||||
self.now += delta
|
||||
|
||||
|
||||
def _full_shard(model: _KvDenseLlama):
|
||||
return _KvReferenceShard(model, 0, model.n_layers - 1)
|
||||
|
||||
|
||||
def _manager_for(shard, config: HotKvStateConfig | None = None, clock=None) -> HotKvStateManager:
|
||||
return HotKvStateManager(kv_recipe_for(shard), config=config, clock=clock)
|
||||
|
||||
|
||||
def _cached_greedy(
|
||||
adapter: KvBoundaryAdapter,
|
||||
manager: HotKvStateManager,
|
||||
session_id: str,
|
||||
epoch: int,
|
||||
prompt: list[int],
|
||||
n_new: int,
|
||||
) -> list[int]:
|
||||
"""Greedy decode one full-model session through the KV manager."""
|
||||
out = adapter.prefill(session_id, epoch, token_ids=np.asarray(prompt))
|
||||
assert isinstance(out, TailOutput)
|
||||
tokens = [out.token_id]
|
||||
for _ in range(n_new - 1):
|
||||
step = adapter.decode(session_id, epoch, token_ids=[out.token_id])
|
||||
assert isinstance(step, TailOutput)
|
||||
out = step
|
||||
tokens.append(out.token_id)
|
||||
return tokens
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Recipe identity.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_recipe_owned_layers_and_fingerprint_aliasing():
|
||||
"The KV recipe covers only owned layers and canonicalizes architecture aliases.\n\nTags: node, kv"
|
||||
recipe = KvCacheRecipe(
|
||||
architecture_adapter="LlamaForCausalLM",
|
||||
kv_dtype="float32",
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
total_layers=6,
|
||||
start_layer=2,
|
||||
end_layer=3,
|
||||
)
|
||||
assert recipe.owned_layers == (2, 3)
|
||||
alias = KvCacheRecipe(
|
||||
architecture_adapter="llama",
|
||||
kv_dtype="float32",
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
total_layers=6,
|
||||
start_layer=2,
|
||||
end_layer=3,
|
||||
)
|
||||
assert recipe.is_compatible(alias)
|
||||
# A different owned range is not compatible.
|
||||
other = KvCacheRecipe(
|
||||
architecture_adapter="llama",
|
||||
kv_dtype="float32",
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
total_layers=6,
|
||||
start_layer=0,
|
||||
end_layer=1,
|
||||
)
|
||||
assert not recipe.is_compatible(other)
|
||||
|
||||
|
||||
def test_recipe_bytes_per_token_scales_with_owned_layers():
|
||||
"KV bytes-per-token counts keys+values across owned layers only.\n\nTags: node, kv"
|
||||
base = dict(
|
||||
architecture_adapter="dense-llama",
|
||||
kv_dtype="float32",
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
total_layers=6,
|
||||
)
|
||||
one = KvCacheRecipe(**base, start_layer=0, end_layer=0)
|
||||
two = KvCacheRecipe(**base, start_layer=0, end_layer=1)
|
||||
# 2 (k+v) * heads * dim * 4 bytes per layer.
|
||||
assert one.bytes_per_token() == 2 * 4 * 8 * 4
|
||||
assert two.bytes_per_token() == 2 * one.bytes_per_token()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Owned-layer allocation.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_manager_allocates_kv_only_for_owned_layers():
|
||||
"A middle shard allocates KV state only for its owned layer range.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _KvReferenceShard(model, 2, 3)
|
||||
manager = _manager_for(shard)
|
||||
session = manager.open("sess-mid", 0)
|
||||
assert session.owned_layers == (2, 3)
|
||||
assert set(session.layers) == {2, 3}
|
||||
with pytest.raises(KeyError):
|
||||
session.layer(0)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Prefill / decode / truncate.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_prefill_then_decode_append_grows_owned_layers():
|
||||
"Prefill and decode append advance every owned layer in lockstep.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
prompt = [5, 12, 3, 41]
|
||||
out = adapter.prefill("s", 0, token_ids=np.asarray(prompt))
|
||||
assert isinstance(out, TailOutput)
|
||||
session = manager.get("s", 0)
|
||||
assert session.seq_len == len(prompt)
|
||||
for cache in session.layers.values():
|
||||
assert cache.length == len(prompt)
|
||||
|
||||
step = adapter.decode("s", 0, token_ids=[out.token_id])
|
||||
assert isinstance(step, TailOutput)
|
||||
assert manager.get("s", 0).seq_len == len(prompt) + 1
|
||||
|
||||
|
||||
def test_truncate_rolls_back_all_owned_layers():
|
||||
"Truncate drops cached positions beyond a length across owned layers.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
adapter.prefill("s", 0, token_ids=np.asarray([1, 2, 3, 4, 5]))
|
||||
assert manager.get("s", 0).seq_len == 5
|
||||
manager.truncate("s", 0, 2)
|
||||
session = manager.get("s", 0)
|
||||
assert session.seq_len == 2
|
||||
for cache in session.layers.values():
|
||||
assert cache.length == 2
|
||||
|
||||
|
||||
def test_layer_kv_cache_rejects_wrong_shape():
|
||||
"LayerKvCache rejects K/V that do not match its head geometry.\n\nTags: node, kv"
|
||||
cache = LayerKvCache(0, n_kv_heads=4, head_dim=8, dtype="float32")
|
||||
with pytest.raises(ValueError):
|
||||
cache.append(np.zeros((1, 3, 8), dtype=np.float32), np.zeros((1, 3, 8), dtype=np.float32))
|
||||
cache.append(np.zeros((2, 4, 8), dtype=np.float32), np.zeros((2, 4, 8), dtype=np.float32))
|
||||
assert cache.length == 2
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Cached vs stateless parity (correctness core).
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_cached_full_shard_decode_matches_stateless_whole_model():
|
||||
"Cached full-model greedy decode reproduces stateless whole-model tokens.\n\nTags: node, kv, parity"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
prompt = [2, 17, 8, 25, 6]
|
||||
n_new = 12
|
||||
reference = model.stateless_greedy(prompt, n_new)
|
||||
cached = _cached_greedy(adapter, manager, "s", 0, prompt, n_new)
|
||||
assert cached == reference
|
||||
assert len(cached) == n_new
|
||||
|
||||
|
||||
def test_cached_prefill_next_token_matches_whole_model_logits():
|
||||
"Cached prefill produces the same next-token logits as the whole model.\n\nTags: node, kv, parity"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
prompt = [9, 1, 44, 6, 30, 11]
|
||||
out = adapter.prefill("s", 0, token_ids=np.asarray(prompt))
|
||||
assert isinstance(out, TailOutput)
|
||||
assert out.token_id == model.whole_model_next_token(prompt)
|
||||
|
||||
|
||||
def test_multi_range_cached_decode_parity_across_a_seam():
|
||||
"A head/tail split with independent per-range KV reproduces whole-model decode.\n\nTags: node, kv, parity"
|
||||
model = _KvDenseLlama()
|
||||
head_shard = _KvReferenceShard(model, 0, 2)
|
||||
tail_shard = _KvReferenceShard(model, 3, 5)
|
||||
head_mgr = _manager_for(head_shard)
|
||||
tail_mgr = _manager_for(tail_shard)
|
||||
head = KvBoundaryAdapter(head_shard, head_mgr)
|
||||
tail = KvBoundaryAdapter(tail_shard, tail_mgr)
|
||||
|
||||
prompt = [7, 3, 22, 5, 9]
|
||||
n_new = 8
|
||||
|
||||
# Each range only allocates its owned layers.
|
||||
def step(token_ids, is_prefill):
|
||||
if is_prefill:
|
||||
bundle = head.prefill("s", 0, token_ids=np.asarray(token_ids))
|
||||
out = tail.prefill("s", 0, boundary=bundle)
|
||||
else:
|
||||
bundle = head.decode("s", 0, token_ids=[token_ids])
|
||||
assert isinstance(bundle, BoundaryBundle)
|
||||
out = tail.decode("s", 0, boundary=bundle)
|
||||
assert isinstance(out, TailOutput)
|
||||
return out.token_id
|
||||
|
||||
tokens = [step(prompt, True)]
|
||||
for _ in range(n_new - 1):
|
||||
tokens.append(step(tokens[-1], False))
|
||||
|
||||
assert head_mgr.get("s", 0).owned_layers == (0, 1, 2)
|
||||
assert tail_mgr.get("s", 0).owned_layers == (3, 4, 5)
|
||||
assert tokens == model.stateless_greedy(prompt, n_new)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Four concurrent sessions with no cross-talk.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_four_interleaved_sessions_have_no_kv_cross_talk():
|
||||
"Four interleaved sessions each decode their own tokens without cross-talk.\n\nTags: node, kv, concurrency"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
prompts = {
|
||||
"alpha": [1, 2, 3, 4],
|
||||
"bravo": [40, 39, 2, 15],
|
||||
"charlie": [7, 7, 7, 7],
|
||||
"delta": [31, 5, 18, 22],
|
||||
}
|
||||
n_new = 10
|
||||
references = {sid: model.stateless_greedy(p, n_new) for sid, p in prompts.items()}
|
||||
# The four prompts must actually diverge, else "no cross-talk" is vacuous.
|
||||
assert len({tuple(v) for v in references.values()}) == 4
|
||||
|
||||
generated: dict[str, list[int]] = {}
|
||||
for sid, prompt in prompts.items():
|
||||
out = adapter.prefill(sid, 0, token_ids=np.asarray(prompt))
|
||||
assert isinstance(out, TailOutput)
|
||||
generated[sid] = [out.token_id]
|
||||
|
||||
# Round-robin decode: every session takes one step per round, interleaved.
|
||||
for _ in range(n_new - 1):
|
||||
for sid in prompts:
|
||||
step = adapter.decode(sid, 0, token_ids=[generated[sid][-1]])
|
||||
assert isinstance(step, TailOutput)
|
||||
generated[sid].append(step.token_id)
|
||||
|
||||
for sid in prompts:
|
||||
assert generated[sid] == references[sid], sid
|
||||
assert manager.session_count == 4
|
||||
|
||||
|
||||
def test_four_sessions_on_real_threads_stay_isolated():
|
||||
"Four sessions decoding on real threads produce their own reference tokens.\n\nTags: node, kv, concurrency"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard, HotKvStateConfig(max_sessions=8))
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
prompts = {
|
||||
"t-alpha": [3, 14, 1, 5],
|
||||
"t-bravo": [2, 27, 18, 4],
|
||||
"t-charlie": [9, 9, 1, 2],
|
||||
"t-delta": [44, 6, 30, 11],
|
||||
}
|
||||
n_new = 8
|
||||
references = {sid: model.stateless_greedy(p, n_new) for sid, p in prompts.items()}
|
||||
results: dict[str, list[int]] = {}
|
||||
errors: list[Exception] = []
|
||||
|
||||
def run(sid: str, prompt: list[int]) -> None:
|
||||
try:
|
||||
results[sid] = _cached_greedy(adapter, manager, sid, 0, prompt, n_new)
|
||||
except Exception as exc: # pragma: no cover - surfaced via assert below
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=run, args=(sid, p)) for sid, p in prompts.items()]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert not errors
|
||||
for sid in prompts:
|
||||
assert results[sid] == references[sid], sid
|
||||
|
||||
|
||||
def test_release_one_session_leaves_others_intact_and_returns_memory():
|
||||
"Releasing one session frees its budget and does not disturb the others.\n\nTags: node, kv, concurrency"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
prompts = {"keep-1": [1, 2, 3], "drop": [10, 11, 12, 13], "keep-2": [5, 6, 7]}
|
||||
n_new = 6
|
||||
references = {sid: model.stateless_greedy(p, n_new) for sid, p in prompts.items()}
|
||||
|
||||
gen: dict[str, list[int]] = {}
|
||||
for sid, prompt in prompts.items():
|
||||
out = adapter.prefill(sid, 0, token_ids=np.asarray(prompt))
|
||||
gen[sid] = [out.token_id]
|
||||
|
||||
bytes_before = manager.total_bytes
|
||||
assert manager.release("drop", 0) is True
|
||||
assert manager.total_bytes < bytes_before
|
||||
|
||||
# A decode on the released session is an explicit cache miss, not corruption.
|
||||
miss = adapter.decode("drop", 0, token_ids=[gen["drop"][-1]])
|
||||
assert isinstance(miss, CacheMiss)
|
||||
assert miss.reason is CacheMissReason.RELEASED
|
||||
|
||||
# The survivors keep decoding to their own references.
|
||||
for _ in range(n_new - 1):
|
||||
for sid in ("keep-1", "keep-2"):
|
||||
step = adapter.decode(sid, 0, token_ids=[gen[sid][-1]])
|
||||
assert isinstance(step, TailOutput)
|
||||
gen[sid].append(step.token_id)
|
||||
for sid in ("keep-1", "keep-2"):
|
||||
assert gen[sid] == references[sid], sid
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Stale epoch / incompatible recipe rejection.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_stale_route_epoch_is_rejected():
|
||||
"A request for an older route epoch than the current one is rejected.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
manager = _manager_for(_full_shard(model))
|
||||
manager.open("s", 5)
|
||||
with pytest.raises(StaleRouteEpochError):
|
||||
manager.open("s", 4)
|
||||
with pytest.raises(StaleRouteEpochError):
|
||||
manager.resolve("s", 4)
|
||||
with pytest.raises(StaleRouteEpochError):
|
||||
manager.append("s", 4, {})
|
||||
|
||||
|
||||
def test_new_route_epoch_supersedes_and_frees_old_epoch():
|
||||
"A newer route epoch supersedes the old one, freeing its KV and reporting a miss.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
adapter.prefill("s", 1, token_ids=np.asarray([1, 2, 3, 4]))
|
||||
bytes_epoch1 = manager.total_bytes
|
||||
assert bytes_epoch1 > 0
|
||||
|
||||
# Re-planned route: epoch 2 starts a fresh isolated context.
|
||||
adapter.prefill("s", 2, token_ids=np.asarray([9, 8]))
|
||||
assert manager.session_keys() == [("s", 2)]
|
||||
# Old epoch is gone; a lookup for it is now stale (epoch < current).
|
||||
with pytest.raises(StaleRouteEpochError):
|
||||
manager.resolve("s", 1)
|
||||
|
||||
|
||||
def test_incompatible_cache_recipe_is_rejected():
|
||||
"A request carrying a different KV recipe is rejected, not silently reused.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
manager.open("s", 0)
|
||||
|
||||
incompatible = KvCacheRecipe(
|
||||
architecture_adapter="dense-llama",
|
||||
kv_dtype="float16", # different KV dtype
|
||||
n_kv_heads=model.n_heads,
|
||||
head_dim=model.head_dim,
|
||||
total_layers=model.n_layers,
|
||||
start_layer=0,
|
||||
end_layer=model.n_layers - 1,
|
||||
)
|
||||
with pytest.raises(IncompatibleCacheRecipeError):
|
||||
manager.resolve("s", 0, recipe=incompatible)
|
||||
with pytest.raises(IncompatibleCacheRecipeError):
|
||||
manager.open("s2", 0, recipe=incompatible)
|
||||
|
||||
|
||||
def test_uncertified_architecture_recipe_fails_closed():
|
||||
"A KV recipe for an uncertified architecture fails closed at construction.\n\nTags: node, kv"
|
||||
from meshnet_node.boundary_adapter import UncertifiedArchitectureError
|
||||
|
||||
with pytest.raises(UncertifiedArchitectureError):
|
||||
KvCacheRecipe(
|
||||
architecture_adapter="qwen3-moe",
|
||||
kv_dtype="float32",
|
||||
n_kv_heads=4,
|
||||
head_dim=8,
|
||||
total_layers=6,
|
||||
start_layer=0,
|
||||
end_layer=5,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Explicit cache-miss responses.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_unknown_session_is_an_explicit_cache_miss():
|
||||
"Resolving an unknown session returns an explicit unknown-session miss.\n\nTags: node, kv"
|
||||
manager = _manager_for(_full_shard(_KvDenseLlama()))
|
||||
miss = manager.resolve("nope", 0)
|
||||
assert isinstance(miss, CacheMiss)
|
||||
assert miss.reason is CacheMissReason.UNKNOWN_SESSION
|
||||
with pytest.raises(KvCacheMissError):
|
||||
manager.get("nope", 0)
|
||||
|
||||
|
||||
def test_seq_len_mismatch_is_an_explicit_cache_miss():
|
||||
"A decode whose expected length disagrees with the cache is an explicit miss.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
out = adapter.prefill("s", 0, token_ids=np.asarray([1, 2, 3]))
|
||||
# Cache holds 3 tokens; claim it holds 99.
|
||||
miss = adapter.decode("s", 0, token_ids=[out.token_id], expected_seq_len=99)
|
||||
assert isinstance(miss, CacheMiss)
|
||||
assert miss.reason is CacheMissReason.SEQ_LEN_MISMATCH
|
||||
|
||||
|
||||
def test_ttl_eviction_yields_an_explicit_cache_miss():
|
||||
"A session idle past its TTL is evicted and reported as a TTL cache miss.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
clock = _FakeClock()
|
||||
manager = _manager_for(shard, HotKvStateConfig(ttl_seconds=10.0), clock=clock)
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
adapter.prefill("s", 0, token_ids=np.asarray([1, 2, 3]))
|
||||
clock.advance(11.0)
|
||||
miss = manager.resolve("s", 0)
|
||||
assert isinstance(miss, CacheMiss)
|
||||
assert miss.reason is CacheMissReason.EVICTED_TTL
|
||||
assert manager.total_bytes == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Eviction and budget.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_lru_eviction_by_session_cap_reports_a_miss():
|
||||
"Exceeding the session cap evicts the least-recently-used session.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
manager = _manager_for(shard, HotKvStateConfig(max_sessions=2))
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
adapter.prefill("a", 0, token_ids=np.asarray([1, 2]))
|
||||
adapter.prefill("b", 0, token_ids=np.asarray([3, 4]))
|
||||
# Touch 'a' so 'b' becomes the LRU victim.
|
||||
adapter.decode("a", 0, token_ids=[1])
|
||||
adapter.prefill("c", 0, token_ids=np.asarray([5, 6]))
|
||||
|
||||
miss = manager.resolve("b", 0)
|
||||
assert isinstance(miss, CacheMiss)
|
||||
assert miss.reason is CacheMissReason.EVICTED_LRU
|
||||
assert set(k[0] for k in manager.session_keys()) == {"a", "c"}
|
||||
|
||||
|
||||
def test_budget_eviction_keeps_total_within_budget():
|
||||
"Byte-budget pressure evicts LRU sessions so the store stays within budget.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
recipe = kv_recipe_for(shard)
|
||||
# Budget for ~5 tokens of one session; a second big session forces eviction.
|
||||
budget = recipe.bytes_per_token() * 5
|
||||
manager = _manager_for(shard, HotKvStateConfig(budget_bytes=budget, max_sessions=8))
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
|
||||
adapter.prefill("a", 0, token_ids=np.asarray([1, 2, 3]))
|
||||
adapter.prefill("b", 0, token_ids=np.asarray([4, 5, 6, 7]))
|
||||
assert manager.total_bytes <= budget
|
||||
# 'a' (older, LRU) was evicted to make room for 'b'.
|
||||
miss = manager.resolve("a", 0)
|
||||
assert isinstance(miss, CacheMiss)
|
||||
assert miss.reason is CacheMissReason.EVICTED_LRU
|
||||
assert manager.get("b", 0).seq_len == 4
|
||||
|
||||
|
||||
def test_single_session_exceeding_budget_raises():
|
||||
"A single session that cannot fit the budget raises instead of evicting itself.\n\nTags: node, kv"
|
||||
model = _KvDenseLlama()
|
||||
shard = _full_shard(model)
|
||||
recipe = kv_recipe_for(shard)
|
||||
budget = recipe.bytes_per_token() * 2 # only 2 tokens fit
|
||||
manager = _manager_for(shard, HotKvStateConfig(budget_bytes=budget))
|
||||
adapter = KvBoundaryAdapter(shard, manager)
|
||||
with pytest.raises(KvBudgetExceededError):
|
||||
adapter.prefill("a", 0, token_ids=np.asarray([1, 2, 3, 4, 5]))
|
||||
78
tests/test_llama_worker_build.py
Normal file
78
tests/test_llama_worker_build.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
SCRIPT = ROOT / "packages" / "node" / "native" / "scripts" / "build_llama_worker.sh"
|
||||
PIN_FILE = ROOT / "packages" / "node" / "native" / "llama" / "UPSTREAM_COMMIT"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not SCRIPT.exists(), reason="llama worker build script is missing")
|
||||
def test_llama_worker_build_smoke_rebuild(tmp_path: Path) -> None:
|
||||
if not shutil_which("git"):
|
||||
pytest.skip("git is unavailable")
|
||||
if not (shutil_which("g++") or shutil_which("c++") or shutil_which("clang++")):
|
||||
pytest.skip("no C++ compiler is unavailable")
|
||||
|
||||
source_dir = tmp_path / "llama.cpp"
|
||||
build_one = tmp_path / "build-1"
|
||||
build_two = tmp_path / "build-2"
|
||||
pin = PIN_FILE.read_text(encoding="utf-8").strip()
|
||||
|
||||
source_dir.mkdir()
|
||||
_write_fake_upstream_tree(source_dir, pin)
|
||||
_git_init(source_dir)
|
||||
|
||||
_run_build(source_dir, build_one)
|
||||
_run_build(source_dir, build_two)
|
||||
|
||||
binary = build_two / "meshnet_worker"
|
||||
assert binary.exists()
|
||||
output = subprocess.run(
|
||||
[str(binary), "--smoke"],
|
||||
cwd=ROOT,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert "meshnet worker scaffold ok" in output.stdout
|
||||
assert pin in output.stdout
|
||||
|
||||
|
||||
def _run_build(source_dir: Path, build_dir: Path) -> None:
|
||||
env = os.environ.copy()
|
||||
env.setdefault("PATH", os.environ.get("PATH", ""))
|
||||
subprocess.run(
|
||||
[str(SCRIPT), "--source-dir", str(source_dir), "--build-dir", str(build_dir)],
|
||||
cwd=ROOT,
|
||||
check=True,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
|
||||
def _write_fake_upstream_tree(source_dir: Path, pin: str) -> None:
|
||||
(source_dir / "LICENSE").write_text("MIT License placeholder\n", encoding="utf-8")
|
||||
(source_dir / "AUTHORS").write_text("Georgi Gerganov\nMeshnet maintainers\n", encoding="utf-8")
|
||||
(source_dir / "CMakeLists.txt").write_text("# upstream placeholder\n", encoding="utf-8")
|
||||
(source_dir / ".meshnet-upstream-commit").write_text(f"{pin}\n", encoding="utf-8")
|
||||
(source_dir / ".meshnet-upstream-repository").write_text(
|
||||
"https://github.com/ggml-org/llama.cpp.git\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _git_init(source_dir: Path) -> None:
|
||||
subprocess.run(["git", "init", "-q"], cwd=source_dir, check=True)
|
||||
|
||||
|
||||
def shutil_which(name: str) -> str | None:
|
||||
from shutil import which
|
||||
|
||||
return which(name)
|
||||
508
tests/test_native_shard_protocol.py
Normal file
508
tests/test_native_shard_protocol.py
Normal file
@@ -0,0 +1,508 @@
|
||||
"""DGR-002: generated-schema round-trip and compatibility tests.
|
||||
|
||||
Covers the versioned gRPC Shard protocol (``packages/node/native/proto``):
|
||||
* Python round-trip across the full envelope, tensor bundle, and every service.
|
||||
* Proto3 forward/backward compatibility (unknown-field preservation, defaults).
|
||||
* Bounded-fragment tensor bundle framing + checksums.
|
||||
* Cross-language Python<->C++ round-trip when the C++ toolchain is available;
|
||||
otherwise the C++ test skips with an explicit reason (deterministic, GPU-free,
|
||||
model-download-free, API-credit-free by construction).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
import pytest
|
||||
|
||||
# grpc_tools (grpcio-tools) is required to generate the stubs. It is present in
|
||||
# the project .venv; skip cleanly elsewhere rather than error.
|
||||
native_protocol = pytest.importorskip(
|
||||
"meshnet_node.native_protocol",
|
||||
reason="meshnet_node.native_protocol import failed",
|
||||
)
|
||||
|
||||
try:
|
||||
native_protocol.generate()
|
||||
_GEN_ERROR = None
|
||||
except native_protocol.ProtocGenerationError as exc: # pragma: no cover
|
||||
_GEN_ERROR = str(exc)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
_GEN_ERROR is not None,
|
||||
reason=f"protobuf stubs unavailable: {_GEN_ERROR}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pb2():
|
||||
return native_protocol.load()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Envelope / header round-trip and field coverage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _full_header(pb2):
|
||||
return pb2.MessageHeader(
|
||||
schema_version=pb2.SCHEMA_VERSION_1,
|
||||
work_id="work-42",
|
||||
route_session_id="rs-7",
|
||||
route_epoch=9,
|
||||
fingerprint=pb2.ArtifactFingerprint(
|
||||
model_id="meta-llama/Llama-3.1-8B",
|
||||
revision="main",
|
||||
artifact_hash="sha256:deadbeef",
|
||||
quantization="Q4_K_M",
|
||||
runtime_recipe_fingerprint="recipe-123",
|
||||
),
|
||||
shard_range=pb2.ShardRange(
|
||||
start_layer=8,
|
||||
end_layer=16,
|
||||
effective_start_layer=9,
|
||||
owns_embedding=False,
|
||||
owns_final_head=False,
|
||||
),
|
||||
phase=pb2.PHASE_PREFILL,
|
||||
position=pb2.Position(start_position=0, token_count=12, sequence_length=12),
|
||||
idempotency_step=3,
|
||||
cache_expectation=pb2.CACHE_REUSE,
|
||||
compression=pb2.COMPRESSION_ZSTD,
|
||||
checksum=pb2.Checksum(algorithm=pb2.CHECKSUM_CRC32C, value=b"\x00\x01\x02\x03"),
|
||||
)
|
||||
|
||||
|
||||
def test_message_header_carries_every_required_field(pb2):
|
||||
"""The header carries every identifier the transport contract demands.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
header = _full_header(pb2)
|
||||
raw = header.SerializeToString()
|
||||
back = pb2.MessageHeader()
|
||||
back.ParseFromString(raw)
|
||||
|
||||
assert back.schema_version == pb2.SCHEMA_VERSION_1
|
||||
assert back.work_id == "work-42"
|
||||
assert back.route_session_id == "rs-7"
|
||||
assert back.route_epoch == 9
|
||||
assert back.fingerprint.artifact_hash == "sha256:deadbeef"
|
||||
assert back.fingerprint.runtime_recipe_fingerprint == "recipe-123"
|
||||
assert back.shard_range.effective_start_layer == 9
|
||||
assert back.phase == pb2.PHASE_PREFILL
|
||||
assert back.position.token_count == 12
|
||||
assert back.idempotency_step == 3
|
||||
assert back.cache_expectation == pb2.CACHE_REUSE
|
||||
assert back.compression == pb2.COMPRESSION_ZSTD
|
||||
assert back.checksum.algorithm == pb2.CHECKSUM_CRC32C
|
||||
assert back.checksum.value == b"\x00\x01\x02\x03"
|
||||
|
||||
|
||||
def test_named_tensor_bundle_describes_shape_dtype_byteorder_and_fragments(pb2):
|
||||
"""A tensor bundle round-trips name, shape, dtype, byte order and fragments.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
bundle = pb2.TensorBundle(
|
||||
bundle_version=1,
|
||||
tensors=[
|
||||
pb2.NamedTensor(
|
||||
name="hidden_states",
|
||||
shape=[2, 3, 4096],
|
||||
dtype=pb2.DTYPE_BF16,
|
||||
byte_order=pb2.BYTE_ORDER_LITTLE_ENDIAN,
|
||||
total_byte_length=16,
|
||||
compression=pb2.COMPRESSION_NONE,
|
||||
fragments=[
|
||||
pb2.TensorFragment(
|
||||
fragment_index=0,
|
||||
fragment_count=2,
|
||||
byte_offset=0,
|
||||
data=b"\x00" * 8,
|
||||
),
|
||||
pb2.TensorFragment(
|
||||
fragment_index=1,
|
||||
fragment_count=2,
|
||||
byte_offset=8,
|
||||
data=b"\x01" * 8,
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
back = pb2.TensorBundle()
|
||||
back.ParseFromString(bundle.SerializeToString())
|
||||
tensor = back.tensors[0]
|
||||
assert tensor.name == "hidden_states"
|
||||
assert list(tensor.shape) == [2, 3, 4096]
|
||||
assert tensor.dtype == pb2.DTYPE_BF16
|
||||
assert tensor.byte_order == pb2.BYTE_ORDER_LITTLE_ENDIAN
|
||||
assert [f.byte_offset for f in tensor.fragments] == [0, 8]
|
||||
|
||||
|
||||
def test_session_stream_carries_open_prefill_decode_release_cancel(pb2):
|
||||
"""The bidi stream oneof expresses every seam operation.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
header = _full_header(pb2)
|
||||
frames = {
|
||||
"open": pb2.SessionActivation(
|
||||
open=pb2.SessionOpen(
|
||||
header=header,
|
||||
deadline_unix_nanos=1_000_000,
|
||||
max_prefill_tokens_per_chunk=256,
|
||||
max_fragment_bytes=1 << 20,
|
||||
initial_credit=pb2.FlowControl(credits=8, max_in_flight_bytes=1 << 24),
|
||||
)
|
||||
),
|
||||
"prefill": pb2.SessionActivation(
|
||||
prefill=pb2.PrefillChunk(
|
||||
header=header, chunk_index=0, chunk_count=2, final_chunk=False
|
||||
)
|
||||
),
|
||||
"decode": pb2.SessionActivation(decode=pb2.DecodeStep(header=header)),
|
||||
"release": pb2.SessionActivation(
|
||||
release=pb2.ReleaseRequest(header=header, reason="done")
|
||||
),
|
||||
"cancel": pb2.SessionActivation(
|
||||
cancel=pb2.CancelRequest(header=header, reason="client abort")
|
||||
),
|
||||
"flow_control": pb2.SessionActivation(
|
||||
flow_control=pb2.FlowControl(credits=4)
|
||||
),
|
||||
}
|
||||
for name, frame in frames.items():
|
||||
back = pb2.SessionActivation()
|
||||
back.ParseFromString(frame.SerializeToString())
|
||||
assert back.WhichOneof("payload") == name
|
||||
|
||||
|
||||
def test_session_response_carries_structured_status_and_results(pb2):
|
||||
"""Server frames carry accepted/result/status/acks with structured Status.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
status = pb2.Status(
|
||||
code=8,
|
||||
message="resource exhausted",
|
||||
retry_class=pb2.RETRY_CLASS_RETRYABLE,
|
||||
details={"queue_depth": "128"},
|
||||
)
|
||||
resp = pb2.SessionResponse(
|
||||
result=pb2.ActivationResult(
|
||||
header=_full_header(pb2),
|
||||
outputs=pb2.TensorBundle(bundle_version=1),
|
||||
cache_result=pb2.CACHE_WRITTEN,
|
||||
status=status,
|
||||
)
|
||||
)
|
||||
back = pb2.SessionResponse()
|
||||
back.ParseFromString(resp.SerializeToString())
|
||||
assert back.WhichOneof("payload") == "result"
|
||||
assert back.result.cache_result == pb2.CACHE_WRITTEN
|
||||
assert back.result.status.retry_class == pb2.RETRY_CLASS_RETRYABLE
|
||||
assert back.result.status.details["queue_depth"] == "128"
|
||||
|
||||
|
||||
def test_capability_and_health_round_trip(pb2):
|
||||
"""Capability and health messages round-trip their admission fields.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
cap = pb2.CapabilityResponse(
|
||||
schema_version=pb2.SCHEMA_VERSION_1,
|
||||
supported_schema_versions=[pb2.SCHEMA_VERSION_1],
|
||||
supported_architectures=["llama"],
|
||||
supported_quantizations=["Q4_K_M", "F16"],
|
||||
servable_range=pb2.ShardRange(start_layer=0, end_layer=16),
|
||||
budget=pb2.ResourceBudget(
|
||||
weight_bytes=1 << 32, kv_bytes=1 << 30, max_concurrent_sessions=4
|
||||
),
|
||||
supported_compression=[pb2.COMPRESSION_NONE, pb2.COMPRESSION_ZSTD],
|
||||
supported_checksums=[pb2.CHECKSUM_CRC32C, pb2.CHECKSUM_SHA256],
|
||||
)
|
||||
cap_back = pb2.CapabilityResponse()
|
||||
cap_back.ParseFromString(cap.SerializeToString())
|
||||
assert cap_back.budget.max_concurrent_sessions == 4
|
||||
assert list(cap_back.supported_quantizations) == ["Q4_K_M", "F16"]
|
||||
|
||||
health = pb2.HealthResponse(
|
||||
status=pb2.SERVING, active_sessions=2, queued_requests=1, kv_pressure=0.5
|
||||
)
|
||||
health_back = pb2.HealthResponse()
|
||||
health_back.ParseFromString(health.SerializeToString())
|
||||
assert health_back.status == pb2.SERVING
|
||||
assert health_back.kv_pressure == pytest.approx(0.5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Compatibility
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_unknown_fields_are_preserved_for_forward_compatibility(pb2):
|
||||
"""An older reader tolerates and preserves fields it does not know.
|
||||
|
||||
A newer sender may add a field; parsing into the current schema must not
|
||||
fail and must round-trip the unknown bytes.
|
||||
|
||||
Tags: protocol, compatibility
|
||||
"""
|
||||
header = _full_header(pb2)
|
||||
raw = bytearray(header.SerializeToString())
|
||||
# Append an unknown field: number 5000, wire type 2 (length-delimited).
|
||||
tag = (5000 << 3) | 2
|
||||
raw += _encode_varint(tag)
|
||||
payload = b"future-field"
|
||||
raw += _encode_varint(len(payload))
|
||||
raw += payload
|
||||
|
||||
parsed = pb2.MessageHeader()
|
||||
# Parsing must not raise on the unknown field.
|
||||
parsed.ParseFromString(bytes(raw))
|
||||
# Known fields survive intact.
|
||||
assert parsed.work_id == "work-42"
|
||||
assert parsed.route_epoch == 9
|
||||
# The unknown bytes are preserved and re-emitted on re-serialization. This is
|
||||
# the behavioural compatibility guarantee; the introspection accessor
|
||||
# (UnknownFields()) is not implemented by the upb backend, so we assert the
|
||||
# observable outcome rather than the accessor.
|
||||
reserialized = parsed.SerializeToString()
|
||||
assert payload in reserialized
|
||||
assert _encode_varint(tag) in reserialized
|
||||
|
||||
|
||||
def test_defaults_are_stable_for_backward_compatibility(pb2):
|
||||
"""A message from an older sender (missing new fields) reads as enum zero.
|
||||
|
||||
Tags: protocol, compatibility
|
||||
"""
|
||||
empty = pb2.MessageHeader()
|
||||
back = pb2.MessageHeader()
|
||||
back.ParseFromString(empty.SerializeToString())
|
||||
assert back.schema_version == pb2.SCHEMA_VERSION_UNSPECIFIED
|
||||
assert back.phase == pb2.PHASE_UNSPECIFIED
|
||||
assert back.cache_expectation == pb2.CACHE_EXPECTATION_UNSPECIFIED
|
||||
assert back.work_id == ""
|
||||
assert back.route_epoch == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bounded-fragment helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fragment_and_reassemble_round_trip_with_checksums(pb2):
|
||||
"""Bounded fragmentation reassembles exactly and validates checksums.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
payload = bytes((i * 7) % 256 for i in range(10_000))
|
||||
tensor = native_protocol.fragment_tensor(
|
||||
name="hidden",
|
||||
shape=[1, 4096],
|
||||
dtype=pb2.DTYPE_F16,
|
||||
payload=payload,
|
||||
max_fragment_bytes=4096,
|
||||
checksum_algorithm=pb2.CHECKSUM_CRC32C,
|
||||
)
|
||||
assert len(tensor.fragments) == 3
|
||||
assert all(len(f.data) <= 4096 for f in tensor.fragments)
|
||||
# Survives a serialization round-trip before reassembly.
|
||||
back = pb2.NamedTensor()
|
||||
back.ParseFromString(tensor.SerializeToString())
|
||||
assert native_protocol.reassemble_tensor(back) == payload
|
||||
|
||||
|
||||
def test_reassemble_detects_fragment_corruption(pb2):
|
||||
"""A flipped fragment byte fails checksum verification.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
payload = b"abcdefabcdef" * 100
|
||||
tensor = native_protocol.fragment_tensor(
|
||||
name="t",
|
||||
shape=[len(payload)],
|
||||
dtype=pb2.DTYPE_U8,
|
||||
payload=payload,
|
||||
max_fragment_bytes=256,
|
||||
checksum_algorithm=pb2.CHECKSUM_SHA256,
|
||||
)
|
||||
tensor.fragments[1].data = tensor.fragments[1].data[:-1] + b"\x00"
|
||||
with pytest.raises(ValueError, match="checksum mismatch"):
|
||||
native_protocol.reassemble_tensor(tensor)
|
||||
|
||||
|
||||
def test_checksum_algorithms_verify(pb2):
|
||||
"""CRC32C, CRC32 and SHA256 all verify their own payloads.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
data = b"the quick brown fox"
|
||||
for algo in (pb2.CHECKSUM_CRC32C, pb2.CHECKSUM_CRC32, pb2.CHECKSUM_SHA256):
|
||||
checksum = native_protocol.compute_checksum(algo, data)
|
||||
assert native_protocol.verify_checksum(checksum, data)
|
||||
assert not native_protocol.verify_checksum(checksum, data + b"!")
|
||||
|
||||
|
||||
def test_service_descriptor_exposes_all_operations(pb2):
|
||||
"""The generated service defines capability/health/session/release/cancel.
|
||||
|
||||
Requires the grpc runtime; skips cleanly without it.
|
||||
|
||||
Tags: protocol
|
||||
"""
|
||||
grpc = pytest.importorskip("grpc", reason="grpc runtime not installed")
|
||||
assert grpc is not None
|
||||
grpc_mod = native_protocol.load_grpc()
|
||||
assert hasattr(grpc_mod, "ShardRuntimeStub")
|
||||
assert hasattr(grpc_mod, "ShardRuntimeServicer")
|
||||
# Confirm the streaming seam and unary ops exist on the servicer.
|
||||
servicer = grpc_mod.ShardRuntimeServicer
|
||||
for op in ("GetCapability", "Health", "ActivateSession", "Release", "Cancel"):
|
||||
assert hasattr(servicer, op), op
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-language Python <-> C++ compatibility
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _cpp_toolchain_reason() -> str | None:
|
||||
"""Return a skip reason if the C++ build toolchain is unavailable."""
|
||||
for tool in ("cmake", "protoc"):
|
||||
if shutil.which(tool) is None:
|
||||
return f"{tool} not found on PATH"
|
||||
return None
|
||||
|
||||
|
||||
def _build_cpp_compatible_sample(pb2):
|
||||
"""Python message matching what roundtrip_test.cpp CheckSample expects."""
|
||||
header = pb2.MessageHeader(
|
||||
schema_version=pb2.SCHEMA_VERSION_1,
|
||||
work_id="w1",
|
||||
route_session_id="s1",
|
||||
route_epoch=3,
|
||||
phase=pb2.PHASE_PREFILL,
|
||||
idempotency_step=7,
|
||||
cache_expectation=pb2.CACHE_FRESH,
|
||||
compression=pb2.COMPRESSION_NONE,
|
||||
fingerprint=pb2.ArtifactFingerprint(
|
||||
model_id="meta-llama/Llama-3.1-8B",
|
||||
quantization="Q4_K_M",
|
||||
runtime_recipe_fingerprint="recipe-abc",
|
||||
),
|
||||
shard_range=pb2.ShardRange(
|
||||
start_layer=0, end_layer=16, effective_start_layer=0, owns_embedding=True
|
||||
),
|
||||
position=pb2.Position(start_position=0, token_count=5, sequence_length=5),
|
||||
)
|
||||
return pb2.SessionActivation(
|
||||
prefill=pb2.PrefillChunk(
|
||||
header=header,
|
||||
chunk_index=0,
|
||||
chunk_count=1,
|
||||
final_chunk=True,
|
||||
activations=pb2.TensorBundle(
|
||||
bundle_version=1,
|
||||
tensors=[
|
||||
pb2.NamedTensor(
|
||||
name="hidden",
|
||||
shape=[1, 4096],
|
||||
dtype=pb2.DTYPE_F16,
|
||||
byte_order=pb2.BYTE_ORDER_LITTLE_ENDIAN,
|
||||
total_byte_length=8,
|
||||
compression=pb2.COMPRESSION_NONE,
|
||||
fragments=[
|
||||
pb2.TensorFragment(
|
||||
fragment_index=0,
|
||||
fragment_count=1,
|
||||
byte_offset=0,
|
||||
data=bytes(range(1, 9)),
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_cross_language_roundtrip_python_and_cpp(pb2, tmp_path):
|
||||
"""Python and C++ parse each other's serialized frames (both directions).
|
||||
|
||||
Builds the C++ round-trip binary via CMake, feeds it a Python-serialized
|
||||
fixture (C++ must parse it), and parses the C++-serialized output back in
|
||||
Python. Skips with an explicit reason when the C++ toolchain is absent.
|
||||
|
||||
Tags: protocol, compatibility, cpp
|
||||
"""
|
||||
reason = _cpp_toolchain_reason()
|
||||
if reason is not None:
|
||||
pytest.skip(f"C++ toolchain unavailable: {reason}")
|
||||
|
||||
native_root = native_protocol.PROTO_DIR.parent
|
||||
build_dir = tmp_path / "cpp-build"
|
||||
|
||||
configure = subprocess.run(
|
||||
["cmake", "-S", str(native_root), "-B", str(build_dir)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if configure.returncode != 0:
|
||||
pytest.skip(
|
||||
"cmake configure failed (protobuf C++ dev likely missing):\n"
|
||||
+ configure.stderr[-2000:]
|
||||
)
|
||||
|
||||
build = subprocess.run(
|
||||
["cmake", "--build", str(build_dir), "--target",
|
||||
"shard_protocol_roundtrip_test"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert build.returncode == 0, f"C++ build failed:\n{build.stderr[-2000:]}"
|
||||
|
||||
binary = build_dir / "shard_protocol_roundtrip_test"
|
||||
assert binary.exists(), "C++ test binary not produced"
|
||||
|
||||
py_fixture = tmp_path / "py_sample.bin"
|
||||
cpp_out = tmp_path / "cpp_sample.bin"
|
||||
py_fixture.write_bytes(_build_cpp_compatible_sample(pb2).SerializeToString())
|
||||
|
||||
run = subprocess.run(
|
||||
[str(binary), "--selftest", "--read", str(py_fixture),
|
||||
"--write", str(cpp_out)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert run.returncode == 0, f"C++ round-trip failed:\n{run.stdout}\n{run.stderr}"
|
||||
|
||||
# C++ parsed our bytes; now Python parses C++'s bytes and checks known fields.
|
||||
parsed = pb2.SessionActivation()
|
||||
parsed.ParseFromString(cpp_out.read_bytes())
|
||||
assert parsed.WhichOneof("payload") == "prefill"
|
||||
assert parsed.prefill.header.work_id == "w1"
|
||||
assert parsed.prefill.header.route_epoch == 3
|
||||
assert parsed.prefill.activations.tensors[0].name == "hidden"
|
||||
assert parsed.prefill.activations.tensors[0].dtype == pb2.DTYPE_F16
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Local helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _encode_varint(value: int) -> bytes:
|
||||
out = bytearray()
|
||||
while True:
|
||||
byte = value & 0x7F
|
||||
value >>= 7
|
||||
if value:
|
||||
out.append(byte | 0x80)
|
||||
else:
|
||||
out.append(byte)
|
||||
return bytes(out)
|
||||
@@ -22,6 +22,7 @@ import pytest
|
||||
|
||||
from meshnet_node.admission import (
|
||||
REASON_BACKEND_MISMATCH,
|
||||
REASON_COMPATIBILITY_MISMATCH,
|
||||
REASON_MODEL_MISMATCH,
|
||||
REASON_NO_REPORT,
|
||||
REASON_NOT_PASSED,
|
||||
@@ -68,11 +69,26 @@ class _FakeBackend:
|
||||
total_layers = 24
|
||||
hidden_size = 8
|
||||
|
||||
def __init__(self, *, shard_start=0, shard_end=23, device="cpu", forward_error=None):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
shard_start=0,
|
||||
shard_end=23,
|
||||
device="cpu",
|
||||
forward_error=None,
|
||||
loaded_shard_start=None,
|
||||
loaded_shard_end=None,
|
||||
owns_embedding=None,
|
||||
owns_final_head=None,
|
||||
):
|
||||
self.shard_start = shard_start
|
||||
self.shard_end = shard_end
|
||||
self.is_head = shard_start == 0
|
||||
self.is_tail = shard_end == self.total_layers - 1
|
||||
self.loaded_shard_start = shard_start if loaded_shard_start is None else loaded_shard_start
|
||||
self.loaded_shard_end = shard_end if loaded_shard_end is None else loaded_shard_end
|
||||
self.owns_embedding = self.is_head if owns_embedding is None else owns_embedding
|
||||
self.owns_final_head = self.is_tail if owns_final_head is None else owns_final_head
|
||||
self.device = _FakeDevice(device)
|
||||
self.model_id = MODEL
|
||||
self._forward_error = forward_error
|
||||
@@ -192,6 +208,17 @@ def test_a_passing_report_from_another_backend_or_device_is_refused():
|
||||
assert exc.value.reason == REASON_BACKEND_MISMATCH
|
||||
|
||||
|
||||
def test_a_passing_report_with_the_wrong_cache_layout_is_refused():
|
||||
"The compatibility fingerprint fails closed when cache layout changes.\n\nTags: node, admission"
|
||||
ctx = _context()
|
||||
report = capability_report_for(ctx, cache_layout="local-hot-kv")
|
||||
|
||||
with pytest.raises(CapabilityAdmissionError) as exc:
|
||||
admit(AdmissionRequirement.for_context(ctx), report)
|
||||
|
||||
assert exc.value.reason == REASON_COMPATIBILITY_MISMATCH
|
||||
|
||||
|
||||
def test_a_report_older_than_the_freshness_window_is_refused():
|
||||
"Hardware, drivers and weights move; an old proof is not a current one.\n\nTags: node, admission"
|
||||
ctx = _context()
|
||||
@@ -371,10 +398,31 @@ def test_a_matching_passing_report_registers_and_travels_with_the_payload(startu
|
||||
assert report["status"] == "passed"
|
||||
assert report["model"]["model_id"] == MODEL
|
||||
assert (report["shard"]["start"], report["shard"]["end"]) == (0, 23)
|
||||
assert report["shard"]["owns_embedding"] is True
|
||||
assert report["shard"]["owns_final_head"] is True
|
||||
assert report["recipe"]["recipe_id"] == DEFAULT_RECIPE_ID
|
||||
assert report["backend"]["device"] == "cpu"
|
||||
|
||||
|
||||
def test_capability_report_prefers_backend_loaded_range_over_cli_claims():
|
||||
"The proof reports the model's loaded range, not the CLI's requested range.\n\nTags: node, admission"
|
||||
backend = _FakeBackend(
|
||||
shard_start=0,
|
||||
shard_end=23,
|
||||
loaded_shard_start=8,
|
||||
loaded_shard_end=15,
|
||||
owns_embedding=False,
|
||||
owns_final_head=True,
|
||||
)
|
||||
report = capability_report_for(
|
||||
_context(backend=backend, shard_start=0, shard_end=23),
|
||||
)
|
||||
|
||||
assert (report.shard.start, report.shard.end) == (8, 15)
|
||||
assert report.shard.owns_embedding is False
|
||||
assert report.shard.owns_final_head is True
|
||||
|
||||
|
||||
def test_the_served_backend_is_loaded_with_the_recipe_that_was_validated(startup_env):
|
||||
"The recipe named in the report is the one the serving backend actually ran.\n\nTags: node, admission, startup"
|
||||
node = _start(recipe_id="eager-attention")
|
||||
|
||||
@@ -42,9 +42,12 @@ def _report(**overrides):
|
||||
status="passed",
|
||||
duration_ms=142,
|
||||
validated_at=1_760_000_000.0,
|
||||
owns_embedding=True,
|
||||
owns_final_head=False,
|
||||
)
|
||||
kwargs.update(overrides)
|
||||
return build_capability_report(**kwargs)
|
||||
report = build_capability_report(**kwargs)
|
||||
return report
|
||||
|
||||
|
||||
# --- model-agnostic identity ------------------------------------------------
|
||||
@@ -114,6 +117,9 @@ def test_report_dict_has_the_stable_documented_key_set():
|
||||
"shard",
|
||||
"recipe",
|
||||
"backend",
|
||||
"artifact",
|
||||
"runtime_recipe",
|
||||
"compatibility_fingerprint",
|
||||
"status",
|
||||
"validated_at",
|
||||
"duration_ms",
|
||||
@@ -121,12 +127,38 @@ def test_report_dict_has_the_stable_documented_key_set():
|
||||
}
|
||||
assert payload["schema_version"] == CAPABILITY_SCHEMA_VERSION
|
||||
assert set(payload["model"]) == {"model_id", "revision", "config_fingerprint"}
|
||||
assert set(payload["shard"]) == {"start", "end"}
|
||||
assert set(payload["shard"]) == {
|
||||
"start",
|
||||
"end",
|
||||
"owns_embedding",
|
||||
"owns_final_head",
|
||||
}
|
||||
assert set(payload["recipe"]) == {
|
||||
"recipe_id",
|
||||
"recipe_version",
|
||||
"catalogue_version",
|
||||
}
|
||||
assert set(payload["artifact"]) == {
|
||||
"model_id",
|
||||
"revision",
|
||||
"artifact_hash",
|
||||
"shard_start",
|
||||
"shard_end",
|
||||
}
|
||||
assert set(payload["runtime_recipe"]) == {
|
||||
"weight_quantization",
|
||||
"activation_dtype",
|
||||
"compute_dtype",
|
||||
"kv_dtype",
|
||||
"kv_layout",
|
||||
"tokenizer_revision",
|
||||
"architecture_adapter",
|
||||
"backend_id",
|
||||
"runtime_version",
|
||||
"boundary_schema_version",
|
||||
"cache_layout",
|
||||
"fingerprint",
|
||||
}
|
||||
assert set(payload["backend"]) == {
|
||||
"backend_id",
|
||||
"device",
|
||||
@@ -134,10 +166,19 @@ def test_report_dict_has_the_stable_documented_key_set():
|
||||
"quantization",
|
||||
"runtime",
|
||||
}
|
||||
assert payload["compatibility_fingerprint"].startswith("sha256:")
|
||||
# JSON-serializable end to end.
|
||||
assert json.loads(json.dumps(payload)) == payload
|
||||
|
||||
|
||||
def test_report_carries_endpoint_ownership():
|
||||
"Endpoint ownership is recorded alongside the shard range.\n\nTags: node, startup"
|
||||
payload = _report().to_dict()
|
||||
|
||||
assert payload["shard"]["owns_embedding"] is True
|
||||
assert payload["shard"]["owns_final_head"] is False
|
||||
|
||||
|
||||
def test_identity_key_pins_model_shard_recipe_and_backend():
|
||||
"Identity key pins model shard recipe and backend\n\nTags: node, startup"
|
||||
base = _report()
|
||||
@@ -156,6 +197,15 @@ def test_identity_key_pins_model_shard_recipe_and_backend():
|
||||
assert _report(device="other-device").identity_key() != base.identity_key()
|
||||
|
||||
|
||||
def test_compatibility_fingerprint_changes_when_the_runtime_recipe_changes():
|
||||
"The compatibility fingerprint changes when the runtime recipe changes.\n\nTags: node, startup"
|
||||
base = _report()
|
||||
altered = _report(cache_layout="stateless")
|
||||
|
||||
assert base.compatibility_fingerprint != altered.compatibility_fingerprint
|
||||
assert base.runtime_recipe.fingerprint != altered.runtime_recipe.fingerprint
|
||||
|
||||
|
||||
def test_config_fingerprint_is_stable_under_key_order_and_detects_change():
|
||||
"Config fingerprint is stable under key order and detects change\n\nTags: node, startup"
|
||||
a = config_fingerprint({"num_hidden_layers": 8, "hidden_size": 512})
|
||||
|
||||
@@ -5,6 +5,7 @@ Model ids here are arbitrary and made up on purpose: nothing in the admission or
|
||||
routing path may branch on a vendor, model or kernel name.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
import urllib.error
|
||||
@@ -18,6 +19,7 @@ from meshnet_tracker.capability import (
|
||||
POLICY_ENFORCE,
|
||||
STATE_ABSENT,
|
||||
STATE_ADMITTED,
|
||||
STATE_COMPATIBILITY_MISMATCH,
|
||||
STATE_CATALOGUE_INCOMPATIBLE,
|
||||
STATE_FAILED,
|
||||
STATE_INVALID,
|
||||
@@ -41,6 +43,14 @@ SHORT = "oracle-9b"
|
||||
LAYERS = 32
|
||||
|
||||
|
||||
def _stable_json(data: dict) -> str:
|
||||
return json.dumps(data, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
|
||||
|
||||
|
||||
def _sha256_text(data: dict) -> str:
|
||||
return "sha256:" + hashlib.sha256(_stable_json(data).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _post_json(url: str, payload: dict) -> dict:
|
||||
data = json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
@@ -60,6 +70,8 @@ def _report(
|
||||
model_id: str = MODEL,
|
||||
start: int = 0,
|
||||
end: int = 15,
|
||||
owns_embedding: bool | None = None,
|
||||
owns_final_head: bool | None = None,
|
||||
status: str = "passed",
|
||||
validated_at: float | None = None,
|
||||
recipe_id: str = "baseline",
|
||||
@@ -70,10 +82,48 @@ def _report(
|
||||
diagnostics: list | None = None,
|
||||
) -> dict:
|
||||
"""A capability report shaped exactly as `meshnet_node.capability` emits it."""
|
||||
return {
|
||||
if owns_embedding is None:
|
||||
owns_embedding = start == 0
|
||||
if owns_final_head is None:
|
||||
owns_final_head = end >= LAYERS - 1
|
||||
artifact = {
|
||||
"model_id": model_id,
|
||||
"revision": None,
|
||||
"artifact_hash": _sha256_text(
|
||||
{
|
||||
"model_id": model_id,
|
||||
"shard_start": start,
|
||||
"shard_end": end,
|
||||
"recipe_id": recipe_id,
|
||||
"recipe_version": recipe_version,
|
||||
}
|
||||
),
|
||||
"shard_start": start,
|
||||
"shard_end": end,
|
||||
}
|
||||
runtime_recipe = {
|
||||
"weight_quantization": "bfloat16",
|
||||
"activation_dtype": "bfloat16",
|
||||
"compute_dtype": "bfloat16",
|
||||
"kv_dtype": "bfloat16",
|
||||
"kv_layout": "session-cache",
|
||||
"tokenizer_revision": model_id,
|
||||
"architecture_adapter": "unknown",
|
||||
"backend_id": "torch-transformers",
|
||||
"runtime_version": "0.1.0",
|
||||
"boundary_schema_version": 1,
|
||||
"cache_layout": "local-hot-kv",
|
||||
}
|
||||
runtime_recipe["fingerprint"] = _sha256_text(runtime_recipe)
|
||||
payload = {
|
||||
"schema_version": schema_version,
|
||||
"model": {"model_id": model_id, "revision": None, "config_fingerprint": None},
|
||||
"shard": {"start": start, "end": end},
|
||||
"shard": {
|
||||
"start": start,
|
||||
"end": end,
|
||||
"owns_embedding": owns_embedding,
|
||||
"owns_final_head": owns_final_head,
|
||||
},
|
||||
"recipe": {
|
||||
"recipe_id": recipe_id,
|
||||
"recipe_version": recipe_version,
|
||||
@@ -86,11 +136,24 @@ def _report(
|
||||
"quantization": "bfloat16",
|
||||
"runtime": {},
|
||||
},
|
||||
"artifact": artifact,
|
||||
"runtime_recipe": runtime_recipe,
|
||||
"status": status,
|
||||
"validated_at": time.time() if validated_at is None else validated_at,
|
||||
"duration_ms": 42,
|
||||
"diagnostics": list(diagnostics or []),
|
||||
}
|
||||
payload["compatibility_fingerprint"] = _sha256_text(
|
||||
{
|
||||
"model": payload["model"],
|
||||
"shard": payload["shard"],
|
||||
"recipe": payload["recipe"],
|
||||
"backend": payload["backend"],
|
||||
"artifact": payload["artifact"],
|
||||
"runtime_recipe": payload["runtime_recipe"],
|
||||
}
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
def _registration(
|
||||
@@ -119,6 +182,7 @@ def _registration(
|
||||
report = _report(start=start, end=end)
|
||||
if report is not None:
|
||||
payload["capability_report"] = report
|
||||
payload["compatibility_fingerprint"] = report["compatibility_fingerprint"]
|
||||
if recipe_id is not None:
|
||||
payload["recipe_id"] = recipe_id
|
||||
if recipe_version is not None:
|
||||
@@ -196,6 +260,15 @@ def test_a_report_for_a_different_recipe_than_the_node_declares_is_a_recipe_mism
|
||||
assert versioned.state == STATE_RECIPE_MISMATCH
|
||||
|
||||
|
||||
def test_a_report_for_a_different_compatibility_fingerprint_is_a_compatibility_mismatch():
|
||||
"The exact artifact/runtime recipe fingerprint gates admission.\n\nTags: routing, tracker"
|
||||
state = _evaluate(
|
||||
_report(),
|
||||
declared_compatibility_fingerprint="sha256:deadbeef",
|
||||
)
|
||||
assert state.state == STATE_COMPATIBILITY_MISMATCH
|
||||
|
||||
|
||||
def test_an_older_recipe_catalogue_is_incompatible():
|
||||
"Recipe ids from a catalogue older than the tracker's minimum cannot be matched.\n\nTags: routing, tracker"
|
||||
state = _evaluate(_report(catalogue_version="2025.01.1"))
|
||||
|
||||
Reference in New Issue
Block a user