popov/neuron-tai
1
0
Fork 0
Code Issues 52 Pull Requests Actions Packages Projects Releases Wiki Activity

chore: clean superseded GGUF scaffolding

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-16 22:32:37 +03:00
parent 81b1fa6074
commit 369b2072cc
102 changed files with 219 additions and 15274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

472
tests/test_batch_scheduler.py
View File

@@ -1,472 +0,0 @@
"""Continuous batching and bounded admission (DGR-012).
These tests drive the node-local continuous-batching scheduler with the *same*
pure-numpy KV-cached dense-Llama reference the Hot KV State manager uses
(DGR-007), imported from ``test_hot_kv_state``. That keeps the whole gate
deterministic, download-free, GPU-free, and API-credit-free while exercising the
real KV isolation path (``KvBoundaryAdapter`` + ``HotKvStateManager``) rather than
a mock.
Coverage maps to the story's acceptance criteria:
* bounded admission against weight/KV/scratch/queue budgets,
* compatible decode steps batched with per-session positions/outputs preserved,
* prefill never starving in-flight decode (explicit decode-first policy),
* backpressure when the bounded queue is full,
* capability telemetry reporting every required signal,
* a deterministic 1/2/4/8 concurrency sweep showing saturation and no
cross-session corruption.
"""
from __future__ import annotations
import numpy as np
import pytest
from meshnet_node.hot_kv_state import (
HotKvStateConfig,
HotKvStateManager,
KvBoundaryAdapter,
kv_recipe_for,
)
from meshnet_node.batch_scheduler import (
AdmissionReason,
ContinuousBatchScheduler,
GenerationRequest,
KvBatchEngine,
NodeBudget,
Phase,
run_concurrency_sweep,
)
# Reuse the certified numpy dense-Llama reference and shard from the DGR-007 gate.
from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard
# --------------------------------------------------------------------------- #
# 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 _make_engine(
model: _KvDenseLlama | None = None,
*,
config: HotKvStateConfig | None = None,
) -> KvBatchEngine:
"""A full-shard KV batch engine over the deterministic numpy dense-Llama."""
model = model or _KvDenseLlama()
shard = _KvReferenceShard(model, 0, model.n_layers - 1)
manager = HotKvStateManager(kv_recipe_for(shard), config=config)
adapter = KvBoundaryAdapter(shard, manager)
return KvBatchEngine(adapter)
def _reference_tokens(model: _KvDenseLlama, prompt, n_new: int) -> list[int]:
return model.stateless_greedy(list(prompt), n_new)
def _generation(session_id: str, prompt, n_new: int, epoch: int = 0) -> GenerationRequest:
return GenerationRequest(
session_id=session_id,
route_epoch=epoch,
prompt_token_ids=tuple(prompt),
max_new_tokens=n_new,
)
# --------------------------------------------------------------------------- #
# Bounded admission (weight / KV / scratch / queue budgets).
# --------------------------------------------------------------------------- #
def test_admission_respects_active_scratch_and_queue_budgets():
"Admission fills active slots, queues the overflow, then rejects a full queue.\n\nTags: node, scheduler, admission"
engine = _make_engine()
budget = NodeBudget(
max_active_sessions=2,
scratch_bytes_per_session=1,
scratch_budget_bytes=2, # scratch also caps at 2 concurrent
max_queue_depth=1,
max_batch_size=2,
)
scheduler = ContinuousBatchScheduler(engine, budget)
a = scheduler.submit(_generation("a", [1, 2, 3], 4))
b = scheduler.submit(_generation("b", [4, 5, 6], 4))
assert a.reason is AdmissionReason.ADMITTED
assert b.reason is AdmissionReason.ADMITTED
# Two active slots full -> the next goes to the bounded queue.
c = scheduler.submit(_generation("c", [7, 8, 9], 4))
assert c.reason is AdmissionReason.QUEUED
# Queue depth 1 is now full -> backpressure rejection.
d = scheduler.submit(_generation("d", [1, 1, 1], 4))
assert d.reason is AdmissionReason.REJECTED_QUEUE_FULL
assert d.rejected
telem = scheduler.telemetry()
assert telem.active_sessions == 2
assert telem.queue_depth == 1
assert telem.rejected_admissions_total == 1
assert telem.rejected_by_reason[AdmissionReason.REJECTED_QUEUE_FULL.value] == 1
def test_admission_rejects_a_session_that_cannot_fit_the_kv_budget():
"A generation whose whole KV cannot fit the node budget is rejected up front.\n\nTags: node, scheduler, admission"
engine = _make_engine()
per_token = engine._manager.recipe.bytes_per_token()
# Budget holds only 3 positions; a prompt(4)+7 new = 10 final positions cannot fit.
budget = NodeBudget(kv_budget_bytes=per_token * 3)
scheduler = ContinuousBatchScheduler(engine, budget)
decision = scheduler.submit(_generation("big", [1, 2, 3, 4], 7))
assert decision.reason is AdmissionReason.REJECTED_KV_BUDGET
assert scheduler.telemetry().rejected_admissions_total == 1
def test_admission_rejects_when_per_session_scratch_exceeds_budget():
"A per-session scratch larger than the whole scratch envelope is rejected.\n\nTags: node, scheduler, admission"
engine = _make_engine()
budget = NodeBudget(scratch_bytes_per_session=1024, scratch_budget_bytes=512)
scheduler = ContinuousBatchScheduler(engine, budget)
decision = scheduler.submit(_generation("s", [1, 2], 2))
assert decision.reason is AdmissionReason.REJECTED_SCRATCH_BUDGET
def test_duplicate_submission_is_rejected():
"Submitting a session id that is already scheduled is rejected as a duplicate.\n\nTags: node, scheduler, admission"
engine = _make_engine()
scheduler = ContinuousBatchScheduler(engine, NodeBudget(max_active_sessions=4))
assert scheduler.submit(_generation("dup", [1, 2], 3)).reason is AdmissionReason.ADMITTED
assert scheduler.submit(_generation("dup", [3, 4], 3)).reason is AdmissionReason.REJECTED_DUPLICATE
def test_weight_budget_is_reported_in_telemetry():
"The resident weight footprint is surfaced as a capability signal.\n\nTags: node, scheduler, telemetry"
engine = _make_engine()
budget = NodeBudget(weight_bytes=123_456)
scheduler = ContinuousBatchScheduler(engine, budget)
assert scheduler.telemetry().weight_bytes == 123_456
# --------------------------------------------------------------------------- #
# Continuous batching preserves per-session positions and outputs.
# --------------------------------------------------------------------------- #
def test_batched_decode_preserves_per_session_positions_and_outputs():
"Four sessions batched together each reproduce their own stateless tokens.\n\nTags: node, scheduler, batching"
model = _KvDenseLlama()
engine = _make_engine(model)
budget = NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4)
scheduler = ContinuousBatchScheduler(engine, budget)
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: _reference_tokens(model, p, n_new) for sid, p in prompts.items()}
# The four references must diverge, else "no cross-talk" would be vacuous.
assert len({tuple(v) for v in references.values()}) == 4
for sid, prompt in prompts.items():
assert scheduler.submit(_generation(sid, prompt, n_new)).running
outputs = scheduler.run_to_completion()
for sid in prompts:
assert outputs[sid] == references[sid], sid
telem = scheduler.telemetry()
# A genuine batch formed: at least one decode tick carried all four sessions.
assert telem.batch_occupancy_max == 4
assert telem.completed_sessions == 4
assert telem.active_sessions == 0
def test_positions_are_isolated_across_different_prompt_lengths():
"Sessions with different prompt lengths keep independent positions when batched.\n\nTags: node, scheduler, batching"
model = _KvDenseLlama()
engine = _make_engine(model)
scheduler = ContinuousBatchScheduler(
engine, NodeBudget(max_active_sessions=3, max_batch_size=3, max_queue_depth=3)
)
jobs = {
"short": ([5], 6),
"medium": ([2, 9, 14], 6),
"long": ([1, 2, 3, 4, 5, 6, 7], 6),
}
refs = {sid: _reference_tokens(model, p, n) for sid, (p, n) in jobs.items()}
for sid, (prompt, n) in jobs.items():
scheduler.submit(_generation(sid, prompt, n))
outputs = scheduler.run_to_completion()
for sid in jobs:
assert outputs[sid] == refs[sid], sid
# --------------------------------------------------------------------------- #
# Prefill does not starve decode.
# --------------------------------------------------------------------------- #
def test_prefill_does_not_starve_in_flight_decode():
"A burst of new prefills never stalls an already-decoding session.\n\nTags: node, scheduler, fairness"
model = _KvDenseLlama()
engine = _make_engine(model)
# One prefill per tick (budget == a single prompt) so prefill is throttled and
# we can observe that decode still advances every tick.
budget = NodeBudget(
max_active_sessions=8,
max_batch_size=8,
max_queue_depth=8,
scratch_bytes_per_session=1,
scratch_budget_bytes=8,
max_prefill_tokens_per_tick=4,
)
scheduler = ContinuousBatchScheduler(engine, budget)
# Session A starts and prefills on tick 1.
scheduler.submit(_generation("A", [3, 14, 1, 5], 12))
scheduler.run_tick()
a_state = scheduler.session_result("A")
assert a_state.phase is Phase.DECODING
a_len = len(a_state.generated)
assert a_len == 1
# Burst of new work arrives while A is decoding.
for sid in ("B", "C", "D", "E"):
scheduler.submit(_generation(sid, [2, 27, 18, 4], 12))
# Over the next few ticks A must decode on *every* tick (never starved),
# while at most one new session prefills per tick (prefill is bounded).
prefill_counts = []
for _ in range(4):
report = scheduler.run_tick()
new_a_len = len(scheduler.session_result("A").generated)
assert new_a_len == a_len + 1, "decode of A stalled while prefills were pending"
a_len = new_a_len
assert "A" in report.decoded
prefill_counts.append(len(report.prefilled))
assert max(prefill_counts) <= 1, "prefill was not bounded per tick"
def test_decode_first_policy_is_explicit_in_a_single_tick():
"In one tick decode of active sessions precedes prefill of new ones.\n\nTags: node, scheduler, fairness"
model = _KvDenseLlama()
engine = _make_engine(model)
scheduler = ContinuousBatchScheduler(
engine,
NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4,
scratch_bytes_per_session=1, scratch_budget_bytes=4),
)
scheduler.submit(_generation("live", [1, 2, 3], 8))
scheduler.run_tick() # 'live' prefills, now decoding
scheduler.submit(_generation("fresh", [9, 8, 7], 8))
report = scheduler.run_tick()
assert "live" in report.decoded
assert "fresh" in report.prefilled
# --------------------------------------------------------------------------- #
# Backpressure and bounded memory.
# --------------------------------------------------------------------------- #
def test_backpressure_signals_when_queue_full_then_recovers():
"A full queue rejects new work; a completed session frees a slot for the queue.\n\nTags: node, scheduler, backpressure"
engine = _make_engine()
budget = NodeBudget(
max_active_sessions=1,
max_batch_size=1,
max_queue_depth=1,
scratch_bytes_per_session=1,
scratch_budget_bytes=1,
)
scheduler = ContinuousBatchScheduler(engine, budget)
assert scheduler.submit(_generation("first", [1, 2], 2)).running
assert scheduler.submit(_generation("second", [3, 4], 2)).reason is AdmissionReason.QUEUED
# Both a slot and the queue are full now.
assert scheduler.submit(_generation("third", [5, 6], 2)).reason is AdmissionReason.REJECTED_QUEUE_FULL
# Drain 'first'; the queued 'second' must be pulled into the freed slot.
scheduler.run_to_completion()
outputs = scheduler.outputs()
assert set(outputs) == {"first", "second"}
def test_completed_sessions_release_kv_so_growth_is_bounded():
"Finished sessions release their KV, so total KV returns to zero.\n\nTags: node, scheduler, backpressure"
engine = _make_engine()
scheduler = ContinuousBatchScheduler(
engine, NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=8)
)
for sid in ("a", "b", "c", "d"):
scheduler.submit(_generation(sid, [1, 2, 3], 4))
scheduler.run_to_completion()
telem = scheduler.telemetry()
assert telem.kv_total_bytes == 0, "KV not released after completion"
assert telem.active_sessions == 0
assert telem.completed_sessions == 4
# --------------------------------------------------------------------------- #
# Telemetry.
# --------------------------------------------------------------------------- #
def test_telemetry_reports_every_required_signal():
"The capability snapshot reports sessions, queue, batch, KV, rates, rejections.\n\nTags: node, scheduler, telemetry"
model = _KvDenseLlama()
engine = _make_engine(model)
clock = _FakeClock()
budget = NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=1)
scheduler = ContinuousBatchScheduler(engine, budget, clock=clock)
scheduler.submit(_generation("x", [1, 2, 3], 4))
scheduler.submit(_generation("y", [4, 5, 6], 4))
scheduler.submit(_generation("z", [7, 8, 9], 4)) # queued
rejected = scheduler.submit(_generation("w", [1, 1, 1], 4)) # queue full
assert rejected.rejected
clock.advance(1.0)
scheduler.run_tick() # both prefill
clock.advance(1.0)
scheduler.run_tick() # both decode as a batch of 2
clock.advance(2.0)
telem = scheduler.telemetry()
snap = telem.to_dict()
for key in (
"active_sessions", "queue_depth", "batch_occupancy_last",
"batch_occupancy_avg", "batch_occupancy_max", "weight_bytes",
"kv_total_bytes", "kv_budget_bytes", "kv_pressure",
"scratch_used_bytes", "scratch_budget_bytes", "scratch_pressure",
"prefill_tokens_total", "decode_tokens_total",
"prefill_tokens_per_sec", "decode_tokens_per_sec",
"rejected_admissions_total", "rejected_by_reason",
"completed_sessions", "ticks",
):
assert key in snap, key
assert telem.batch_occupancy_max == 2
assert telem.prefill_tokens_total == 6 # two prompts of length 3
assert telem.decode_tokens_total == 2 # one batched decode step, two sessions
assert telem.rejected_admissions_total == 1
# Rates are deterministic under the injected clock: 4 seconds elapsed.
assert telem.decode_tokens_per_sec == pytest.approx(2 / 4.0)
assert telem.prefill_tokens_per_sec == pytest.approx(6 / 4.0)
assert 0.0 < telem.kv_pressure <= 1.0
# --------------------------------------------------------------------------- #
# Concurrency 1/2/4/8 sweep: saturation and no corruption.
# --------------------------------------------------------------------------- #
def test_concurrency_sweep_identifies_saturation_without_corruption():
"A 1/2/4/8 sweep raises batch occupancy, cuts ticks, and never corrupts output.\n\nTags: node, scheduler, benchmark"
model = _KvDenseLlama()
prompts = {
"s0": [1, 2, 3, 4], "s1": [5, 6, 7, 8], "s2": [9, 10, 11, 12],
"s3": [13, 14, 15, 16], "s4": [17, 18, 19, 20], "s5": [21, 22, 23, 24],
"s6": [25, 26, 27, 28], "s7": [29, 30, 31, 32],
}
n_new = 8
requests = [_generation(sid, p, n_new) for sid, p in prompts.items()]
sweep = run_concurrency_sweep(
lambda: _make_engine(model),
requests,
concurrency_levels=(1, 2, 4, 8),
)
assert sweep.corruption_free
assert [r.concurrency for r in sweep.results] == [1, 2, 4, 8]
# No session hit a cache miss (budgets are sized to never evict here).
assert all(r.cache_misses == 0 for r in sweep.results)
assert all(r.rejected_admissions == 0 for r in sweep.results)
# Each per-session stream matches the serialized (concurrency-1) reference.
for sid, prompt in prompts.items():
assert list(sweep.reference_outputs[sid]) == _reference_tokens(model, prompt, n_new)
occupancies = [r.avg_batch_occupancy for r in sweep.results]
ticks = [r.ticks for r in sweep.results]
tokens_per_tick = [r.tokens_per_tick for r in sweep.results]
# Batching packs more sessions per decode step as concurrency rises, so
# average occupancy strictly increases and total ticks strictly decrease.
assert occupancies == sorted(occupancies) and len(set(occupancies)) == 4
assert ticks == sorted(ticks, reverse=True) and len(set(ticks)) == 4
# Aggregate work per tick rises with concurrency (the throughput win).
assert tokens_per_tick == sorted(tokens_per_tick)
# For eight equal-length jobs the node keeps saturating up to the top level.
assert sweep.saturation_concurrency == 8
# The report is JSON-safe for durable evidence.
import json
json.dumps(sweep.to_dict())
def test_concurrency_sweep_saturates_below_max_when_load_is_small():
"With fewer concurrent jobs than slots, saturation is found below the top level.\n\nTags: node, scheduler, benchmark"
model = _KvDenseLlama()
# Only three jobs: at concurrency 4 and 8 the batch can never exceed 3, so
# occupancy stops rising past the load and saturation is detected early.
requests = [
_generation("j0", [1, 2, 3], 6),
_generation("j1", [4, 5, 6], 6),
_generation("j2", [7, 8, 9], 6),
]
sweep = run_concurrency_sweep(
lambda: _make_engine(model), requests, concurrency_levels=(1, 2, 4, 8)
)
assert sweep.corruption_free
assert sweep.saturation_concurrency <= 4
# Levels at or above the load size share the same occupancy/tick profile.
top = [r for r in sweep.results if r.concurrency >= 4]
assert len({r.ticks for r in top}) == 1
# --------------------------------------------------------------------------- #
# Engine contract guards.
# --------------------------------------------------------------------------- #
def test_kv_batch_engine_requires_a_full_shard():
"The batch engine rejects a partial (non head+tail) shard.\n\nTags: node, scheduler"
model = _KvDenseLlama()
head = _KvReferenceShard(model, 0, 2) # head only, not tail
manager = HotKvStateManager(kv_recipe_for(head))
adapter = KvBoundaryAdapter(head, manager)
with pytest.raises(Exception):
KvBatchEngine(adapter)
def test_run_to_completion_is_bounded_against_misconfiguration():
"run_to_completion raises rather than looping forever when work cannot drain.\n\nTags: node, scheduler"
engine = _make_engine()
scheduler = ContinuousBatchScheduler(
engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4)
)
scheduler.submit(_generation("only", [1, 2], 3))
# A tiny explicit tick ceiling is exceeded deterministically.
with pytest.raises(Exception):
scheduler.run_to_completion(max_ticks=1)

488
tests/test_boundary_adapter.py
View File

@@ -1,488 +0,0 @@
"""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

611
tests/test_failure_semantics.py
View File

@@ -1,611 +0,0 @@
"""Bounded failure, cancellation, and restart semantics (DGR-013).
These tests drive the hardened per-session decode stream with the *same*
pure-numpy KV-cached dense-Llama reference the Hot KV State manager (DGR-007) and
the continuous-batch scheduler (DGR-012) use, imported from ``test_hot_kv_state``.
The whole matrix stays deterministic, download-free, GPU-free, and API-credit-free
while exercising the real KV isolation path (``KvBoundaryAdapter`` +
``HotKvStateManager``) rather than a mock.
Coverage maps to the story's acceptance criteria:
* deadlines and heartbeat/health loss terminate blocked stream operations,
* cancellation propagates across every Shard and releases KV + queued buffers,
* duplicate steps are idempotent; uncertain mutations are never replayed silently,
* alpha failover restarts from token zero rather than importing unverified KV,
* worker death / stream reset / malformed bundle / stale epoch / cache miss,
* billing/work records distinguish completed, cancelled, failed, and unverified.
"""
from __future__ import annotations
import json
import numpy as np
import pytest
from meshnet_node.batch_scheduler import (
ContinuousBatchScheduler,
DoneReason,
GenerationRequest,
KvBatchEngine,
NodeBudget,
)
from meshnet_node.boundary_adapter import BoundaryBundle, BoundaryContractError
from meshnet_node.hot_kv_state import (
CacheMiss,
CacheMissReason,
HotKvStateConfig,
HotKvStateManager,
KvBoundaryAdapter,
StaleRouteEpochError,
kv_recipe_for,
)
from meshnet_node.failure_semantics import (
CancellationToken,
DeadlineGuard,
FailureKind,
HardenedSessionRunner,
IdempotencyLedger,
OperationCancelled,
RestartController,
ShardCancellationGroup,
StepKey,
StreamTerminated,
UncertainMutationError,
WorkLedger,
WorkRecord,
WorkStatus,
classify_exception,
work_status_for,
)
# Reuse the certified numpy dense-Llama reference and shard from the DGR-007 gate.
from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard
# --------------------------------------------------------------------------- #
# 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
class _FaultyShard(_KvReferenceShard):
"""A full-shard reference that raises on the Nth ``run_layers_cached`` call.
``run_layers_cached`` is invoked once per stream step, so ``fail_at_call=k``
simulates a worker dying at step ``k-1`` (calls are 1-indexed). The call
counter persists across attempts, so a restart on a fresh epoch keeps counting
and does not re-trip the same fault.
"""
def __init__(self, model, start, end, *, fail_at_call=None, error=None):
super().__init__(model, start, end)
self._fail_at_call = fail_at_call
self._error = error or RuntimeError("worker died mid-step")
self.calls = 0
def run_layers_cached(self, hidden, *, positions, past_kv):
self.calls += 1
if self._fail_at_call is not None and self.calls == self._fail_at_call:
raise self._error
return super().run_layers_cached(hidden, positions=positions, past_kv=past_kv)
def _make_adapter(model=None, *, config=None, shard=None):
"""A full-shard KV boundary adapter over the deterministic numpy dense-Llama."""
model = model or _KvDenseLlama()
shard = shard or _KvReferenceShard(model, 0, model.n_layers - 1)
manager = HotKvStateManager(kv_recipe_for(shard), config=config)
adapter = KvBoundaryAdapter(shard, manager)
return adapter
def _generation(session_id, prompt, n_new, epoch=0):
return GenerationRequest(
session_id=session_id,
route_epoch=epoch,
prompt_token_ids=tuple(prompt),
max_new_tokens=n_new,
)
# --------------------------------------------------------------------------- #
# Happy path (the baseline the failure paths deviate from).
# --------------------------------------------------------------------------- #
def test_clean_run_matches_stateless_reference_and_is_billable():
"A clean stream reproduces the stateless tokens and records completed work.\n\nTags: node, failure, billing"
model = _KvDenseLlama()
adapter = _make_adapter(model)
runner = HardenedSessionRunner(adapter)
prompt = [1, 2, 3, 4]
n_new = 8
outcome = runner.run(_generation("clean", prompt, n_new))
assert outcome.status is WorkStatus.COMPLETED
assert list(outcome.tokens) == model.stateless_greedy(prompt, n_new)
record = runner.work_ledger.records_for("clean")[0]
assert record.billable
assert record.tokens == n_new
assert runner.work_ledger.billable_tokens() == n_new
# --------------------------------------------------------------------------- #
# Deadlines and heartbeat/health loss terminate blocked operations.
# --------------------------------------------------------------------------- #
def test_deadline_terminates_a_blocked_stream_and_releases_kv():
"A deadline reached mid-stream terminates the run and frees its KV.\n\nTags: node, failure, deadline"
clock = _FakeClock()
adapter = _make_adapter()
manager = adapter.manager
runner = HardenedSessionRunner(adapter, clock=clock)
# Each step advances the clock by 1.0; the deadline fires at t=3.
def before_step(_step):
clock.advance(1.0)
outcome = runner.run(
_generation("slow", [5, 6, 7], 20),
deadline=3.0,
before_step=before_step,
)
assert outcome.status is WorkStatus.FAILED
assert outcome.failure_kind is FailureKind.DEADLINE_EXCEEDED
# The stream did not hang and did not finish: only the steps before the
# deadline committed, and the session's KV was released.
assert outcome.token_count < 20
assert isinstance(manager.resolve("slow", 0), CacheMiss)
def test_heartbeat_loss_terminates_a_blocked_stream():
"Losing the peer heartbeat past the timeout terminates the stream.\n\nTags: node, failure, heartbeat"
clock = _FakeClock()
adapter = _make_adapter()
runner = HardenedSessionRunner(adapter, clock=clock)
def before_step(_step):
clock.advance(1.0)
# Heartbeats stop arriving after step 2; with a timeout of 1.5 the gap grows
# past the bound and the stream is terminated (health loss).
def heartbeat(step):
return step < 2
outcome = runner.run(
_generation("hb", [9, 8, 7], 20),
heartbeat_timeout=1.5,
heartbeat=heartbeat,
before_step=before_step,
)
assert outcome.status is WorkStatus.FAILED
assert outcome.failure_kind is FailureKind.HEARTBEAT_LOST
assert outcome.token_count < 20
def test_deadline_guard_reports_remaining_and_resets_on_heartbeat():
"The guard exposes remaining time and a heartbeat resets the health timer.\n\nTags: node, failure, deadline"
clock = _FakeClock()
guard = DeadlineGuard(deadline=10.0, heartbeat_timeout=2.0, clock=clock)
guard.start()
guard.check()
assert guard.remaining() == 10.0
clock.advance(1.5)
guard.heartbeat() # health refreshed at t=1.5
clock.advance(1.0) # gap since heartbeat is 1.0 < 2.0
guard.check()
clock.advance(2.5) # gap since heartbeat is now 3.5 > 2.0
with pytest.raises(StreamTerminated) as exc:
guard.check()
assert exc.value.kind is FailureKind.HEARTBEAT_LOST
# --------------------------------------------------------------------------- #
# Cancellation propagates across shards and releases KV + queued buffers.
# --------------------------------------------------------------------------- #
def test_cancellation_token_terminates_stream_and_releases_kv():
"A client cancel mid-stream stops the run and releases the session KV.\n\nTags: node, failure, cancel"
adapter = _make_adapter()
manager = adapter.manager
token = CancellationToken()
runner = HardenedSessionRunner(adapter)
# Cancel after two steps have run.
def before_step(step):
if step == 2:
token.cancel("client-hangup")
outcome = runner.run(
_generation("cancelme", [1, 2, 3], 20),
cancel_token=token,
before_step=before_step,
)
assert outcome.status is WorkStatus.CANCELLED
assert outcome.failure_kind is FailureKind.CANCELLED
assert outcome.token_count == 2 # steps 0 and 1 committed before the cancel
assert isinstance(manager.resolve("cancelme", 0), CacheMiss)
def test_shard_cancellation_group_releases_every_shard_and_queued_buffers():
"One cancel frees KV on every node-local shard and releases queued buffers.\n\nTags: node, failure, cancel"
model = _KvDenseLlama()
# Three node-local shards of the same route, each with its own KV manager.
managers = []
for start, end in ((0, 1), (2, 3), (4, 5)):
shard = _KvReferenceShard(model, start, end)
mgr = HotKvStateManager(kv_recipe_for(shard))
mgr.open("route", 0) # each holds live state for the session
managers.append(mgr)
released_buffers = []
group = ShardCancellationGroup("route", 0)
for mgr in managers:
group.add_shard(mgr)
group.add_queued_buffer(lambda: released_buffers.append("bundle-a"))
group.add_queued_buffer(lambda: released_buffers.append("bundle-b"))
outcome = group.cancel()
assert outcome.shards_released == 3
assert outcome.buffers_released == 2
assert released_buffers == ["bundle-a", "bundle-b"]
# Every shard's KV is gone: a lookup now yields an explicit released miss.
for mgr in managers:
miss = mgr.resolve("route", 0)
assert isinstance(miss, CacheMiss)
assert miss.reason is CacheMissReason.RELEASED
# Cancellation is idempotent.
again = group.cancel()
assert again.shards_released == 0
assert again.buffers_released == 0
def test_scheduler_cancel_drains_queue_and_releases_active_kv():
"The scheduler cancel drops queued work and frees an active session's KV.\n\nTags: node, scheduler, cancel"
model = _KvDenseLlama()
shard = _KvReferenceShard(model, 0, model.n_layers - 1)
manager = HotKvStateManager(kv_recipe_for(shard))
engine = KvBatchEngine(KvBoundaryAdapter(shard, manager))
scheduler = ContinuousBatchScheduler(
engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4)
)
assert scheduler.submit(_generation("active", [1, 2, 3], 8)).running
assert scheduler.submit(_generation("waiting", [4, 5, 6], 8)).reason.value == "queued"
scheduler.run_tick() # 'active' prefills and starts decoding, holding KV
# Cancel the queued one: it leaves the queue without ever taking a slot.
assert scheduler.cancel("waiting") is True
# Cancel the active one: its KV is released and it is recorded as cancelled.
assert scheduler.cancel("active") is True
assert manager.total_bytes == 0
telem = scheduler.telemetry()
assert telem.cancelled_sessions == 2
assert telem.completed_sessions == 0
assert telem.active_sessions == 0
assert telem.queue_depth == 0
# Cancelling an unknown / already-finished session is a no-op.
assert scheduler.cancel("active") is False
assert scheduler.cancel("never-seen") is False
def test_scheduler_cancel_rejects_a_completed_reason():
"cancel() refuses a non-terminal reason so completed work is never faked.\n\nTags: node, scheduler, cancel"
model = _KvDenseLlama()
shard = _KvReferenceShard(model, 0, model.n_layers - 1)
manager = HotKvStateManager(kv_recipe_for(shard))
engine = KvBatchEngine(KvBoundaryAdapter(shard, manager))
scheduler = ContinuousBatchScheduler(engine)
scheduler.submit(_generation("x", [1, 2], 4))
with pytest.raises(Exception):
scheduler.cancel("x", reason=DoneReason.COMPLETED)
# --------------------------------------------------------------------------- #
# Idempotency: duplicate steps are no-ops; uncertain mutations never replay.
# --------------------------------------------------------------------------- #
def test_duplicate_step_delivery_is_idempotent_no_remutation():
"Replaying a committed step returns the recorded token without re-mutating KV.\n\nTags: node, failure, idempotency"
ledger = IdempotencyLedger()
key = StepKey("s", 0, 5)
disposition = ledger.begin(key)
assert disposition.fresh
ledger.commit(key, 42)
# A duplicate delivery of the same step returns the recorded token and is a
# no-op — the caller must not re-run the mutation.
replay = ledger.begin(key)
assert replay.duplicate
assert replay.token == 42
def test_idempotent_run_replays_tokens_without_advancing_kv():
"Re-running a completed stream on the same ledger/epoch re-mutates nothing.\n\nTags: node, failure, idempotency"
model = _KvDenseLlama()
adapter = _make_adapter(model)
ledger = IdempotencyLedger()
runner = HardenedSessionRunner(adapter, idempotency=ledger)
request = _generation("idem", [3, 1, 4], 6)
first = runner.run(request)
assert first.status is WorkStatus.COMPLETED
kv_len_after_first = adapter.manager.get("idem", 0).seq_len
# A duplicate delivery of the entire stream: every step is a committed
# duplicate, so the runner replays the identical tokens and the KV length is
# unchanged (no double-append).
second = runner.run(request)
assert second.status is WorkStatus.COMPLETED
assert list(second.tokens) == list(first.tokens)
assert adapter.manager.get("idem", 0).seq_len == kv_len_after_first
def test_uncertain_mutation_is_never_replayed_silently():
"A step marked uncertain refuses a silent replay; it must be verified/restarted.\n\nTags: node, failure, idempotency"
ledger = IdempotencyLedger()
key = StepKey("s", 0, 3)
ledger.begin(key)
ledger.mark_uncertain(key, "worker died before ack")
# Replaying an uncertain mutation is refused rather than silently re-applied.
with pytest.raises(UncertainMutationError):
ledger.begin(key)
assert ledger.has_uncertain()
def test_in_flight_duplicate_is_treated_as_uncertain():
"A second begin before commit is refused (concurrent duplicate is unverified).\n\nTags: node, failure, idempotency"
ledger = IdempotencyLedger()
key = StepKey("s", 0, 1)
ledger.begin(key) # in-flight, not yet committed
with pytest.raises(UncertainMutationError):
ledger.begin(key)
# --------------------------------------------------------------------------- #
# Worker death, stream reset, malformed bundle, stale epoch, cache miss.
# --------------------------------------------------------------------------- #
def test_worker_death_midstream_is_unverified_and_marks_step_uncertain():
"A worker dying mid-step yields unverified work and an unreplayable step.\n\nTags: node, failure, worker-death"
model = _KvDenseLlama()
# Fail on the 3rd step call (step index 2), after two tokens committed.
shard = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3)
adapter = _make_adapter(model, shard=shard)
ledger = IdempotencyLedger()
runner = HardenedSessionRunner(adapter, idempotency=ledger)
outcome = runner.run(_generation("dead", [1, 2, 3], 8))
assert outcome.status is WorkStatus.UNVERIFIED
assert outcome.failure_kind is FailureKind.WORKER_DEATH
assert outcome.token_count == 2 # the two committed steps
assert not outcome.completed
# The failed step is uncertain and can never be silently replayed.
assert ledger.has_uncertain()
with pytest.raises(UncertainMutationError):
ledger.begin(StepKey("dead", 0, 2))
# KV was released on failure.
assert isinstance(adapter.manager.resolve("dead", 0), CacheMiss)
def test_stream_reset_is_restartable_failure():
"A stream reset injected mid-stream fails the run as a restartable transport loss.\n\nTags: node, failure, stream-reset"
adapter = _make_adapter()
runner = HardenedSessionRunner(adapter)
def before_step(step):
if step == 2:
raise StreamTerminated(FailureKind.STREAM_RESET, "peer reset the stream")
outcome = runner.run(_generation("reset", [1, 2, 3], 8), before_step=before_step)
assert outcome.status is WorkStatus.FAILED
assert outcome.failure_kind is FailureKind.STREAM_RESET
assert outcome.restartable
def test_malformed_bundle_is_classified_and_does_not_corrupt_kv():
"A malformed activation bundle is rejected and leaves the KV context empty.\n\nTags: node, failure, malformed-bundle"
model = _KvDenseLlama()
mid = _KvReferenceShard(model, 2, 3) # middle range: not head, not tail
manager = HotKvStateManager(kv_recipe_for(mid))
adapter = KvBoundaryAdapter(mid, manager)
assert not adapter.is_head and not adapter.is_tail
# A bundle that hands over at the wrong layer is malformed.
bad = BoundaryBundle(
architecture_adapter=adapter.architecture.adapter,
schema_version=adapter.architecture.boundary_schema_version,
tensor_name=adapter.architecture.boundary_tensor_name,
residual=np.zeros((1, 3, model.hidden), dtype=np.float32),
positions=np.arange(3, dtype=np.int64)[None, :],
next_layer=adapter.start_layer + 5, # wrong handover layer
normalized=False,
)
with pytest.raises(BoundaryContractError) as exc:
adapter.prefill("mal", 0, boundary=bad)
assert classify_exception(exc.value) is FailureKind.MALFORMED_BUNDLE
# The malformed step never appended KV: the context is empty, not corrupted.
assert manager.get("mal", 0).seq_len == 0
def test_stale_epoch_reference_is_rejected_and_classified():
"A reference to a superseded epoch is rejected as stale, never silently reused.\n\nTags: node, failure, stale-epoch"
model = _KvDenseLlama()
adapter = _make_adapter(model)
manager = adapter.manager
manager.open("sess", 5) # current epoch is now 5
with pytest.raises(StaleRouteEpochError) as exc:
manager.resolve("sess", 4) # epoch 4 is stale
assert classify_exception(exc.value) is FailureKind.STALE_EPOCH
# Driving the hardened runner on the stale epoch fails closed as STALE_EPOCH.
runner = HardenedSessionRunner(adapter)
outcome = runner.run(_generation("sess", [1, 2, 3], 4, epoch=3))
assert outcome.status is WorkStatus.FAILED
assert outcome.failure_kind is FailureKind.STALE_EPOCH
def test_cache_miss_midstream_is_restartable():
"A KV eviction mid-stream surfaces an explicit cache miss the head can restart.\n\nTags: node, failure, cache-miss"
adapter = _make_adapter()
manager = adapter.manager
runner = HardenedSessionRunner(adapter)
# Evict the session's KV just before step 3's decode.
def before_step(step):
if step == 3:
manager.release("evict", 0)
outcome = runner.run(_generation("evict", [1, 2, 3], 10), before_step=before_step)
assert outcome.failure_kind is FailureKind.CACHE_MISS
assert outcome.restartable
assert outcome.token_count == 3 # steps 0..2 committed before the eviction
# --------------------------------------------------------------------------- #
# Alpha failover: restart from token zero, never import unverified KV.
# --------------------------------------------------------------------------- #
def test_alpha_failover_restarts_from_token_zero_and_completes():
"A transient worker death fails over to a fresh epoch and reproduces the tokens.\n\nTags: node, failure, failover"
model = _KvDenseLlama()
# Die on the 3rd step of the first attempt; the persistent call counter means
# the restart (which keeps counting) does not re-trip the fault.
shard = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3)
adapter = _make_adapter(model, shard=shard)
manager = adapter.manager
runner = HardenedSessionRunner(adapter)
controller = RestartController([manager])
prompt = [7, 3, 9, 1]
n_new = 6
result = runner.run_with_failover(
_generation("alpha", prompt, n_new, epoch=0), controller, max_restarts=2
)
assert result.completed
assert result.restarts == 1
# The restart began on a fresh epoch and reproduced the full stateless stream
# from token zero — no half-computed KV was imported.
assert result.outcome.route_epoch == 1
assert list(result.outcome.tokens) == model.stateless_greedy(prompt, n_new)
# The failed epoch's KV is gone and the epoch is now stale.
with pytest.raises(StaleRouteEpochError):
manager.resolve("alpha", 0)
# First attempt was unverified, the restart completed: only the restart bills.
statuses = [a.status for a in result.attempts]
assert statuses == [WorkStatus.UNVERIFIED, WorkStatus.COMPLETED]
assert runner.work_ledger.billable_tokens() == n_new
def test_failover_refuses_to_import_unverified_kv():
"assert_fresh_start fails closed if any shard still holds new-epoch KV.\n\nTags: node, failure, failover"
model = _KvDenseLlama()
adapter = _make_adapter(model)
manager = adapter.manager
controller = RestartController([manager])
new_epoch = controller.failover("s", 0)
assert new_epoch == 1
# A clean fresh start passes.
controller.assert_fresh_start("s", new_epoch)
# If unverified KV were present under the new epoch, the guard refuses it.
manager.open("s", new_epoch)
manager.append(
"s",
new_epoch,
{i: (np.zeros((1, model.n_heads, model.head_dim), dtype=np.float32),
np.zeros((1, model.n_heads, model.head_dim), dtype=np.float32))
for i in range(model.n_layers)},
)
with pytest.raises(Exception):
controller.assert_fresh_start("s", new_epoch)
def test_non_restartable_failure_is_not_retried():
"A deterministic failure (deadline) returns immediately without a restart.\n\nTags: node, failure, failover"
clock = _FakeClock()
adapter = _make_adapter()
runner = HardenedSessionRunner(adapter, clock=clock)
controller = RestartController([adapter.manager])
def before_step(_step):
clock.advance(1.0)
result = runner.run_with_failover(
_generation("bounded", [1, 2, 3], 20),
controller,
max_restarts=3,
deadline=2.0,
before_step=before_step,
)
assert not result.completed
assert result.restarts == 0
assert result.outcome.failure_kind is FailureKind.DEADLINE_EXCEEDED
# --------------------------------------------------------------------------- #
# Billing / work records distinguish completed, cancelled, failed, unverified.
# --------------------------------------------------------------------------- #
def test_work_ledger_distinguishes_all_four_statuses():
"The work ledger keeps completed/cancelled/failed/unverified distinct.\n\nTags: node, failure, billing"
ledger = WorkLedger()
ledger.record(WorkRecord("a", 0, WorkStatus.COMPLETED, tokens=8))
ledger.record(WorkRecord("b", 0, WorkStatus.CANCELLED, tokens=3,
failure_kind=FailureKind.CANCELLED))
ledger.record(WorkRecord("c", 0, WorkStatus.FAILED, tokens=1,
failure_kind=FailureKind.DEADLINE_EXCEEDED))
ledger.record(WorkRecord("d", 0, WorkStatus.UNVERIFIED, tokens=2,
failure_kind=FailureKind.WORKER_DEATH))
counts = ledger.counts_by_status()
assert counts == {
"completed": 1, "cancelled": 1, "failed": 1, "unverified": 1,
}
# Only completed work is billable — cancelled/failed/unverified tokens are
# recorded for observability but never charged.
assert ledger.billable_tokens() == 8
assert [r.session_id for r in ledger.billable_records()] == ["a"]
# JSON-safe for durable evidence.
payload = ledger.to_dict()
assert payload["billable_tokens"] == 8
assert payload["counts_by_status"]["unverified"] == 1
json.dumps(payload)
def test_work_status_and_classification_mapping():
"Failure kinds map to the right billing status and exception classes.\n\nTags: node, failure, billing"
assert work_status_for(FailureKind.CANCELLED) is WorkStatus.CANCELLED
assert work_status_for(FailureKind.WORKER_DEATH) is WorkStatus.UNVERIFIED
# A stream reset detected at a step boundary is a certain failure (nothing
# committed for that step) — only an unexpected mid-step error is unverified.
assert work_status_for(FailureKind.STREAM_RESET) is WorkStatus.FAILED
assert work_status_for(FailureKind.DEADLINE_EXCEEDED) is WorkStatus.FAILED
assert work_status_for(FailureKind.MALFORMED_BUNDLE) is WorkStatus.FAILED
assert work_status_for(FailureKind.STALE_EPOCH) is WorkStatus.FAILED
assert work_status_for(FailureKind.CACHE_MISS) is WorkStatus.FAILED
assert classify_exception(OperationCancelled()) is FailureKind.CANCELLED
assert classify_exception(StaleRouteEpochError("x")) is FailureKind.STALE_EPOCH
assert classify_exception(BoundaryContractError("x")) is FailureKind.MALFORMED_BUNDLE
assert classify_exception(RuntimeError("boom")) is FailureKind.WORKER_DEATH
assert (
classify_exception(StreamTerminated(FailureKind.HEARTBEAT_LOST))
is FailureKind.HEARTBEAT_LOST
)

186
tests/test_gguf_backend.py
View File

@@ -1,186 +0,0 @@
"""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
View File

@@ -1,88 +0,0 @@
"""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
View File

@@ -1,769 +0,0 @@
"""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
View File

@@ -1,78 +0,0 @@
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
View File

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

50
tests/test_node_admission.py
View File

@@ -22,7 +22,6 @@ import pytest
from meshnet_node.admission import (
REASON_BACKEND_MISMATCH,
REASON_COMPATIBILITY_MISMATCH,
REASON_MODEL_MISMATCH,
REASON_NO_REPORT,
REASON_NOT_PASSED,
@@ -69,26 +68,11 @@ class _FakeBackend:
total_layers = 24
hidden_size = 8
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,
):
def __init__(self, *, shard_start=0, shard_end=23, device="cpu", forward_error=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
@@ -208,17 +192,6 @@ 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()
@@ -465,31 +438,10 @@ 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")

54
tests/test_node_capability.py
View File

@@ -42,12 +42,9 @@ 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)
report = build_capability_report(**kwargs)
return report
return build_capability_report(**kwargs)
# --- model-agnostic identity ------------------------------------------------
@@ -117,9 +114,6 @@ def test_report_dict_has_the_stable_documented_key_set():
"shard",
"recipe",
"backend",
"artifact",
"runtime_recipe",
"compatibility_fingerprint",
"status",
"validated_at",
"duration_ms",
@@ -127,38 +121,12 @@ 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",
"owns_embedding",
"owns_final_head",
}
assert set(payload["shard"]) == {"start", "end"}
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",
@@ -166,19 +134,10 @@ 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()
@@ -197,15 +156,6 @@ 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})

286
tests/test_performance_contract.py
View File

@@ -1,286 +0,0 @@
"""Tests for the DGR-001 performance contract metadata."""
from __future__ import annotations
import json
from unittest.mock import MagicMock, patch
import pytest
from meshnet_node.performance_contract import (
BENCHMARK_SCHEMA_VERSION,
DEFAULT_CONTRACT,
SCHEMA_VERSION,
main,
run_performance_benchmark,
run_real_model_endpoint_benchmark,
)
def test_default_contract_is_architecture_aligned_and_small():
"""The baseline stays on DeepSeek2 and uses the smallest DeepSeek-family GGUF.
Tags: performance, model, gguf
"""
payload = DEFAULT_CONTRACT.to_dict()
assert payload["schema_version"] == SCHEMA_VERSION
assert payload["story_id"] == "DGR-001"
assert payload["model_target"] == {
"name": "DeepSeek-V2-Lite-Chat",
"architecture": "deepseek2",
"safetensors_repo": "deepseek-ai/DeepSeek-V2-Lite-Chat",
"safetensors_precision": "bfloat16",
"gguf_repo": "second-state/DeepSeek-V2-Lite-Chat-GGUF",
"gguf_quant": "Q2_K",
"gguf_size_gb": 6.43,
"comparison_policy": (
"same model/revision, closest practical low-footprint precision pair: "
"BF16 safetensors versus Q2_K GGUF"
),
"rationale": (
"Smallest DeepSeek-family benchmark anchor that still points toward "
"DeepSeek-V4-Flash; keeps the runtime on the DeepSeek2 path instead "
"of falling back to a tiny but architecture-mismatched smoke model."
),
}
assert payload["benchmark_lanes"] == [
{
"id": "transformers-safetensors-cpu",
"runtime": "transformers",
"device": "cpu",
"recipe": "current safetensors recipe",
"concurrency_levels": [1, 4],
},
{
"id": "llama-cpp-gguf-cpu",
"runtime": "llama.cpp",
"device": "cpu",
"recipe": "whole-model GGUF recipe",
"concurrency_levels": [1, 4],
},
{
"id": "transformers-safetensors-gpu",
"runtime": "transformers",
"device": "gpu",
"recipe": "current safetensors recipe",
"concurrency_levels": [1, 4],
},
{
"id": "llama-cpp-gguf-gpu",
"runtime": "llama.cpp",
"device": "gpu",
"recipe": "whole-model GGUF recipe",
"concurrency_levels": [1, 4],
},
]
assert "ttft_ms" in payload["metrics"]
assert "output_drift" in payload["metrics"]
assert "meaningful speed or fit benefit" in payload["stop_condition"]
assert any("mounted drive" in note for note in payload["notes"])
def test_contract_cli_writes_json(tmp_path, capsys):
"""The contract can be emitted as a machine-readable artifact.
Tags: performance, artifact
"""
output = tmp_path / "performance-contract.json"
assert main(["--json-out", str(output)]) == 0
written = json.loads(output.read_text(encoding="utf-8"))
assert written == DEFAULT_CONTRACT.to_dict()
assert str(output) in capsys.readouterr().out
def test_stub_benchmark_covers_every_lane_concurrency_and_metric():
"""The runner exercises all four CPU/GPU lanes with the full metric set.
Tags: performance, benchmark, gguf
"""
report = run_performance_benchmark()
assert report["schema_version"] == BENCHMARK_SCHEMA_VERSION
assert report["story_id"] == "DGR-001"
assert report["source"] == "stub-backend"
assert report["model_target"] == DEFAULT_CONTRACT.model_target.to_dict()
assert [lane["id"] for lane in report["lanes"]] == [
lane.id for lane in DEFAULT_CONTRACT.benchmark_lanes
]
for lane in report["lanes"]:
assert [result["concurrency"] for result in lane["results"]] == [1, 4]
for result in lane["results"]:
assert set(result["metrics"]) == set(DEFAULT_CONTRACT.metrics)
assert result["metrics"]["failure_count"] == 0
assert result["metrics"]["decode_tok_per_sec"] > 0
def test_stub_benchmark_is_deterministic():
"""Two runs produce byte-identical reports; no clocks or randomness leak in.
Tags: performance, benchmark, deterministic
"""
first = run_performance_benchmark()
second = run_performance_benchmark()
assert first == second
assert json.dumps(first, sort_keys=True) == json.dumps(second, sort_keys=True)
def test_stub_benchmark_compares_gguf_against_safetensors_per_device():
"""Each device gets a GGUF-vs-safetensors comparison and a stop-condition verdict.
Tags: performance, benchmark, gguf
"""
report = run_performance_benchmark()
assert set(report["comparisons"]) == {"cpu", "gpu"}
cpu, gpu = report["comparisons"]["cpu"], report["comparisons"]["gpu"]
assert cpu["safetensors_lane"] == "transformers-safetensors-cpu"
assert cpu["gguf_lane"] == "llama-cpp-gguf-cpu"
assert cpu["memory_metric"] == "rss_bytes"
assert gpu["safetensors_lane"] == "transformers-safetensors-gpu"
assert gpu["gguf_lane"] == "llama-cpp-gguf-gpu"
assert gpu["memory_metric"] == "vram_bytes"
for comparison in (cpu, gpu):
assert comparison["decode_speedup"] > 1.0
assert comparison["artifact_bytes_ratio"] < 0.5
assert comparison["memory_bytes_ratio"] < 1.0
assert comparison["output_drift"] == 0.0
assert comparison["gguf_benefit"] is True
assert report["stop_condition"]["gguf_benefit"] is True
assert report["stop_condition"]["triggered"] is False
assert report["stop_condition"]["text"] == DEFAULT_CONTRACT.stop_condition
def test_contract_cli_writes_benchmark_report(tmp_path, capsys):
"""--benchmark-out emits the stub benchmark report next to the contract.
Tags: performance, benchmark, artifact
"""
contract_out = tmp_path / "performance-contract.json"
benchmark_out = tmp_path / "artifacts" / "stub-benchmark-report.json"
assert main(["--json-out", str(contract_out), "--benchmark-out", str(benchmark_out)]) == 0
report = json.loads(benchmark_out.read_text(encoding="utf-8"))
assert report == run_performance_benchmark()
output = capsys.readouterr().out
assert str(contract_out) in output
assert str(benchmark_out) in output
def test_real_model_endpoint_benchmark_uses_lane_specific_endpoints_and_shared_schema():
"""The live client path fans out to one endpoint per CPU/GPU lane.
Tags: performance, benchmark, live
"""
response = MagicMock()
response.read.return_value = json.dumps({"choices": [{"message": {"content": "mesh activation"}}]}).encode()
response.headers.get.return_value = "lane-session"
response.__enter__.return_value = response
endpoints = {
"transformers-safetensors-cpu": "http://cpu-safetensors",
"llama-cpp-gguf-cpu": "http://cpu-gguf",
"transformers-safetensors-gpu": "http://gpu-safetensors",
"llama-cpp-gguf-gpu": "http://gpu-gguf",
}
with patch("meshnet_node.performance_contract.urllib.request.urlopen", return_value=response) as urlopen:
report = run_real_model_endpoint_benchmark(endpoints=endpoints, model="deepseek-ai/DeepSeek-V2-Lite-Chat")
assert report["source"] == "real-model-endpoints"
assert report["model_target"] == DEFAULT_CONTRACT.model_target.to_dict()
assert set(report["comparisons"]) == {"cpu", "gpu"}
assert urlopen.call_count == len(endpoints)
called_urls = [call.args[0].full_url for call in urlopen.call_args_list]
assert called_urls == [f"{url}/v1/chat/completions" for url in endpoints.values()]
for lane in report["lanes"]:
assert lane["results"][0]["metrics"]["decode_tok_per_sec"] > 0
assert lane["results"][0]["metrics"]["ttft_ms"] > 0
assert lane["output_tokens"] == ["mesh", "activation"]
assert report["comparisons"]["cpu"]["gguf_lane"] == "llama-cpp-gguf-cpu"
assert report["comparisons"]["gpu"]["gguf_lane"] == "llama-cpp-gguf-gpu"
def test_contract_cli_runs_live_endpoint_benchmark(tmp_path, capsys):
"""--live-endpoint mappings drive the live runner and write its report.
Tags: performance, benchmark, live, artifact
"""
contract_out = tmp_path / "performance-contract.json"
live_out = tmp_path / "artifacts" / "live-benchmark-report.json"
endpoints = {
"transformers-safetensors-cpu": "http://cpu-safetensors",
"llama-cpp-gguf-cpu": "http://cpu-gguf",
"transformers-safetensors-gpu": "http://gpu-safetensors",
"llama-cpp-gguf-gpu": "http://gpu-gguf",
}
fake_report = {"schema_version": BENCHMARK_SCHEMA_VERSION, "source": "real-model-endpoints"}
argv = ["--json-out", str(contract_out), "--live-benchmark-out", str(live_out)]
for lane_id, url in endpoints.items():
argv += ["--live-endpoint", f"{lane_id}={url}"]
with patch(
"meshnet_node.performance_contract.run_real_model_endpoint_benchmark",
return_value=fake_report,
) as runner:
assert main(argv) == 0
runner.assert_called_once_with(
endpoints,
model=DEFAULT_CONTRACT.model_target.safetensors_repo,
contract=DEFAULT_CONTRACT,
)
assert json.loads(live_out.read_text(encoding="utf-8")) == fake_report
output = capsys.readouterr().out
assert str(contract_out) in output
assert str(live_out) in output
def test_contract_cli_passes_explicit_live_model(tmp_path):
"""--live-model overrides the contract's safetensors repo default.
Tags: performance, benchmark, live
"""
live_out = tmp_path / "live-benchmark-report.json"
argv = [
"--json-out", str(tmp_path / "performance-contract.json"),
"--live-benchmark-out", str(live_out),
"--live-endpoint", "transformers-safetensors-cpu=http://cpu-safetensors",
"--live-model", "local/DeepSeek-V2-Lite-Chat-Q2_K",
]
with patch(
"meshnet_node.performance_contract.run_real_model_endpoint_benchmark",
return_value={},
) as runner:
assert main(argv) == 0
assert runner.call_args.kwargs["model"] == "local/DeepSeek-V2-Lite-Chat-Q2_K"
@pytest.mark.parametrize(
"argv",
[
["--live-endpoint", "transformers-safetensors-cpu=http://cpu"],
["--live-benchmark-out", "live-report.json"],
[
"--live-endpoint", "not-a-mapping",
"--live-benchmark-out", "live-report.json",
],
],
ids=["endpoint-without-out", "out-without-endpoint", "malformed-mapping"],
)
def test_contract_cli_rejects_incomplete_live_arguments(tmp_path, argv, capsys):
"""Live flags must arrive as a consistent LANE_ID=URL + output-path set.
Tags: performance, benchmark, live, cli
"""
with pytest.raises(SystemExit) as excinfo:
main(["--json-out", str(tmp_path / "performance-contract.json"), *argv])
assert excinfo.value.code == 2
assert "--live-" in capsys.readouterr().err

77
tests/test_tracker_capability_admission.py
View File

@@ -5,7 +5,6 @@ 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
@@ -19,7 +18,6 @@ from meshnet_tracker.capability import (
POLICY_ENFORCE,
STATE_ABSENT,
STATE_ADMITTED,
STATE_COMPATIBILITY_MISMATCH,
STATE_CATALOGUE_INCOMPATIBLE,
STATE_FAILED,
STATE_INVALID,
@@ -43,14 +41,6 @@ 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(
@@ -70,8 +60,6 @@ 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",
@@ -82,48 +70,10 @@ def _report(
diagnostics: list | None = None,
) -> dict:
"""A capability report shaped exactly as `meshnet_node.capability` emits it."""
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 = {
return {
"schema_version": schema_version,
"model": {"model_id": model_id, "revision": None, "config_fingerprint": None},
"shard": {
"start": start,
"end": end,
"owns_embedding": owns_embedding,
"owns_final_head": owns_final_head,
},
"shard": {"start": start, "end": end},
"recipe": {
"recipe_id": recipe_id,
"recipe_version": recipe_version,
@@ -136,24 +86,11 @@ 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(
@@ -182,7 +119,6 @@ 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:
@@ -260,15 +196,6 @@ 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"))