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

story: DGR-031 Introduce the project-owned ShardEngine interface

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-23 11:00:33 +03:00
parent fd742d35c0
commit c34ab059cc
8 changed files with 2636 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

273
tests/shard_engine_contract.py Normal file
View File

@@ -0,0 +1,273 @@
"""Reusable ``ShardEngine`` lifecycle contract (DGR-031).
Any :class:`~meshnet_node.shard_engine.ShardEngine` implementation — the
DGR-032 deterministic fixture, the DGR-037 llama.cpp binding, or a throwaway
test double — can be checked against this contract by calling
:func:`assert_shard_engine_contract` with a zero-argument factory that
returns a fresh, unloaded engine instance. It proves the *lifecycle
semantics* (load/capabilities gating, cache-miss/stale-epoch/cancel/release
behavior, head vs. middle boundary-vs-token output) are identical across
implementations. It says nothing about whether the numbers an implementation
produces are numerically correct — that is DGR-036's job.
This module is not itself collected as a test file (it does not match
``test_*.py``); import ``assert_shard_engine_contract`` from a real test file
that supplies the engine factory, as ``test_shard_engine.py`` does here.
"""
from __future__ import annotations
from typing import Callable
from meshnet_node.shard_engine import (
BoundaryBundle,
DecodeRequest,
EngineTensor,
LoadRequest,
PrefillRequest,
ShardEngine,
TokenOutput,
)
from meshnet_node.shard_lifecycle import CacheResult, StatusCode
def assert_shard_engine_contract(make_engine: Callable[[], ShardEngine]) -> None:
"""Run every lifecycle check against a fresh engine instance per check.
Each check gets its own ``make_engine()`` instance so one check's session
state can never leak into another's.
"""
_assert_health_before_load_is_not_serving(make_engine())
_assert_load_then_capabilities_matches_range(make_engine())
_assert_prefill_then_decode_succeeds_and_is_deterministic(make_engine())
_assert_middle_shard_accepts_boundary_bundle_not_token_ids(make_engine())
_assert_decode_without_prefill_is_a_deterministic_cache_miss(make_engine())
_assert_stale_epoch_is_rejected(make_engine())
_assert_cancel_then_decode_is_rejected_and_cancel_is_idempotent(make_engine())
_assert_release_then_decode_is_rejected_and_release_is_idempotent(make_engine())
_assert_metrics_reports_cancelled_sessions(make_engine())
def _load(
engine: ShardEngine, *, shard_start: int = 0, shard_end: int = 3, total_layers: int = 4
):
result = engine.load(
LoadRequest(
artifact_path="fixture://contract-test",
shard_start=shard_start,
shard_end=shard_end,
total_layers=total_layers,
)
)
assert result.status.code is StatusCode.OK, result.status
return result
def _output_bytes(output: BoundaryBundle | TokenOutput | None) -> bytes:
assert output is not None
if isinstance(output, TokenOutput):
return output.token_id.to_bytes(8, "big")
return b"".join(tensor.data for tensor in output.tensors)
def _assert_health_before_load_is_not_serving(engine: ShardEngine) -> None:
health = engine.health()
assert health.status.code is StatusCode.OK
assert health.serving is False
def _assert_load_then_capabilities_matches_range(engine: ShardEngine) -> None:
_load(engine, shard_start=0, shard_end=3, total_layers=4)
caps = engine.capabilities()
assert caps.status.code is StatusCode.OK
assert caps.shard_start == 0
assert caps.shard_end == 3
assert caps.total_layers == 4
assert caps.is_head is True
assert caps.is_tail is True
assert caps.supports_mtp is False, "MTP must stay reserved-off until DGR-066"
assert engine.health().serving is True
def _assert_prefill_then_decode_succeeds_and_is_deterministic(engine: ShardEngine) -> None:
_load(engine)
prefill = engine.prefill(
PrefillRequest(
session_id="session-a",
route_epoch=1,
position=0,
idempotency_step=0,
token_ids=(1, 2, 3),
)
)
assert prefill.status.code is StatusCode.OK
assert isinstance(prefill.output, (BoundaryBundle, TokenOutput))
decode = engine.decode(
DecodeRequest(
session_id="session-a",
route_epoch=1,
position=3,
idempotency_step=1,
token_id=4,
)
)
assert decode.status.code is StatusCode.OK
assert decode.cache_result is CacheResult.HIT
assert isinstance(decode.output, (BoundaryBundle, TokenOutput))
# Determinism: the identical prefill replayed on a brand-new session
# produces byte-identical output. The transform is a pure function of
# its inputs, not of hidden randomness or cross-session state.
replay = engine.prefill(
PrefillRequest(
session_id="session-b",
route_epoch=1,
position=0,
idempotency_step=0,
token_ids=(1, 2, 3),
)
)
assert _output_bytes(replay.output) == _output_bytes(prefill.output)
def _assert_middle_shard_accepts_boundary_bundle_not_token_ids(engine: ShardEngine) -> None:
_load(engine, shard_start=1, shard_end=2, total_layers=8)
caps = engine.capabilities()
assert caps.is_head is False
assert caps.is_tail is False
input_bundle = BoundaryBundle(
tensors=(
EngineTensor(name="hidden_states", shape=(1, 3), dtype="bfloat16", data=b"\x00" * 8),
),
architecture="dense",
boundary_point="pre_tail_residual",
)
result = engine.prefill(
PrefillRequest(
session_id="session-middle",
route_epoch=1,
position=0,
idempotency_step=0,
input=input_bundle,
)
)
assert result.status.code is StatusCode.OK
assert isinstance(result.output, BoundaryBundle), "a non-tail shard must hand off a boundary bundle, never a sampled token"
def _assert_decode_without_prefill_is_a_deterministic_cache_miss(engine: ShardEngine) -> None:
_load(engine)
result = engine.decode(
DecodeRequest(
session_id="never-opened",
route_epoch=1,
position=0,
idempotency_step=0,
token_id=9,
)
)
assert result.status.code is not StatusCode.OK
assert result.cache_result is CacheResult.MISS
assert result.output is None
def _assert_stale_epoch_is_rejected(engine: ShardEngine) -> None:
_load(engine)
engine.prefill(
PrefillRequest(
session_id="session-epoch",
route_epoch=5,
position=0,
idempotency_step=0,
token_ids=(1,),
)
)
stale = engine.decode(
DecodeRequest(
session_id="session-epoch",
route_epoch=4,
position=1,
idempotency_step=1,
token_id=2,
)
)
assert stale.status.code is not StatusCode.OK
assert stale.output is None
def _assert_cancel_then_decode_is_rejected_and_cancel_is_idempotent(engine: ShardEngine) -> None:
_load(engine)
engine.prefill(
PrefillRequest(
session_id="session-cancel",
route_epoch=1,
position=0,
idempotency_step=0,
token_ids=(1,),
)
)
cancelled = engine.cancel("session-cancel")
assert cancelled.code is StatusCode.CANCELLED
after = engine.decode(
DecodeRequest(
session_id="session-cancel",
route_epoch=1,
position=1,
idempotency_step=1,
token_id=2,
)
)
assert after.status.code is StatusCode.CANCELLED
assert after.output is None
again = engine.cancel("session-cancel")
assert again.code is StatusCode.CANCELLED
def _assert_release_then_decode_is_rejected_and_release_is_idempotent(engine: ShardEngine) -> None:
_load(engine)
engine.prefill(
PrefillRequest(
session_id="session-release",
route_epoch=1,
position=0,
idempotency_step=0,
token_ids=(1,),
)
)
released = engine.release("session-release")
assert released.code is StatusCode.OK
after = engine.decode(
DecodeRequest(
session_id="session-release",
route_epoch=1,
position=1,
idempotency_step=1,
token_id=2,
)
)
assert after.status.code is not StatusCode.OK
again = engine.release("session-release")
assert again.code is StatusCode.OK
def _assert_metrics_reports_cancelled_sessions(engine: ShardEngine) -> None:
_load(engine)
engine.prefill(
PrefillRequest(
session_id="session-metrics",
route_epoch=1,
position=0,
idempotency_step=0,
token_ids=(1,),
)
)
engine.cancel("session-metrics")
metrics = engine.metrics()
assert metrics.status.code is StatusCode.OK
assert metrics.cancelled_sessions >= 1

241
tests/test_shard_engine.py Normal file
View File

@@ -0,0 +1,241 @@
"""DGR-031 ``ShardEngine`` contract tests.
``_ReferenceEngine`` below is a minimal, in-memory ``ShardEngine`` that exists
only to prove :func:`assert_shard_engine_contract` is non-vacuous and to pin
the abstract contract's own validation rules. It is deliberately not the
DGR-032 deterministic fixture (delay/memory-pressure/malformed/crash
injection, full session/epoch modeling for the fake worker) — that is a
separate, larger story. DGR-032 and DGR-037 are expected to import
``assert_shard_engine_contract`` from ``tests/shard_engine_contract.py``
against their own engines.
"""
from __future__ import annotations
import hashlib
import pytest
from meshnet_node.shard_engine import (
ArchitectureAuxStateHook,
BoundaryBundle,
DecodeRequest,
EngineCapabilities,
EngineTensor,
HealthResult,
LoadRequest,
LoadResult,
MetricsResult,
MtpHook,
PrefillRequest,
ShardEngine,
StepResult,
TokenOutput,
)
from meshnet_node.shard_lifecycle import CacheResult, StatusCode, StructuredStatus
from shard_engine_contract import assert_shard_engine_contract
class _ReferenceEngine(ShardEngine):
"""Minimal in-memory engine used only to exercise the shared contract."""
def __init__(self) -> None:
self._loaded: LoadRequest | None = None
self._sessions: dict[str, dict] = {}
self._cancelled_total = 0
def load(self, request: LoadRequest) -> LoadResult:
self._loaded = request
return LoadResult(
status=StructuredStatus(StatusCode.OK, "loaded"),
effective_start=request.shard_start,
architecture="dense",
)
def capabilities(self) -> EngineCapabilities:
if self._loaded is None:
return EngineCapabilities(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "not loaded"))
request = self._loaded
return EngineCapabilities(
status=StructuredStatus(StatusCode.OK, "ready"),
shard_start=request.shard_start,
shard_end=request.shard_end,
effective_start=request.shard_start,
total_layers=request.total_layers,
architecture="dense",
max_concurrent_sessions=8,
max_context_tokens=131072,
supports_mtp=False,
)
def prefill(self, request: PrefillRequest) -> StepResult:
if self._loaded is None:
return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine not loaded"))
self._sessions[request.session_id] = {"epoch": request.route_epoch, "cancelled": False}
output = self._transform(self._seed_bytes(request.token_ids, request.input), request.idempotency_step)
return StepResult(status=StructuredStatus(StatusCode.OK, "prefilled"), cache_result=CacheResult.STORED, output=output)
def decode(self, request: DecodeRequest) -> StepResult:
session = self._sessions.get(request.session_id)
if session is None:
return StepResult(
status=StructuredStatus(StatusCode.NOT_FOUND, "no cached session state"),
cache_result=CacheResult.MISS,
)
if session["cancelled"]:
return StepResult(status=StructuredStatus(StatusCode.CANCELLED, "session cancelled"))
if request.route_epoch < session["epoch"]:
return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "stale route epoch"))
session["epoch"] = request.route_epoch
token_ids = (request.token_id,) if request.token_id is not None else None
output = self._transform(self._seed_bytes(token_ids, request.input), request.idempotency_step)
return StepResult(status=StructuredStatus(StatusCode.OK, "decoded"), cache_result=CacheResult.HIT, output=output)
def cancel(self, session_id: str, *, work_id: str = "", reason: str = "") -> StructuredStatus:
session = self._sessions.setdefault(session_id, {"epoch": 0, "cancelled": False})
if not session["cancelled"]:
self._cancelled_total += 1
session["cancelled"] = True
return StructuredStatus(StatusCode.CANCELLED, reason or "cancelled")
def release(self, session_id: str) -> StructuredStatus:
self._sessions.pop(session_id, None)
return StructuredStatus(StatusCode.OK, "released")
def health(self) -> HealthResult:
return HealthResult(
status=StructuredStatus(StatusCode.OK, "ok"),
serving=self._loaded is not None,
state="SERVING" if self._loaded is not None else "NOT_LOADED",
active_sessions=len(self._sessions),
)
def metrics(self) -> MetricsResult:
return MetricsResult(
status=StructuredStatus(StatusCode.OK, "ok"),
active_sessions=len(self._sessions),
cancelled_sessions=self._cancelled_total,
)
@staticmethod
def _seed_bytes(token_ids, bundle: BoundaryBundle | None) -> bytes:
if token_ids:
return b"".join(int(t).to_bytes(4, "big") for t in token_ids)
if bundle is not None:
return b"".join(tensor.data for tensor in bundle.tensors)
return b""
def _transform(self, seed: bytes, idempotency_step: int) -> BoundaryBundle | TokenOutput:
digest = hashlib.sha256(seed + idempotency_step.to_bytes(4, "big")).digest()
assert self._loaded is not None
if self._loaded.shard_end >= self._loaded.total_layers - 1:
token_id = int.from_bytes(digest[:4], "big") % 50_000
return TokenOutput(token_id=token_id)
tensor = EngineTensor(name="hidden_states", shape=(1, max(len(seed) // 4, 1)), dtype="bfloat16", data=digest)
return BoundaryBundle(tensors=(tensor,), architecture="dense", boundary_point="pre_tail_residual")
def test_reference_engine_obeys_the_shared_shard_engine_contract():
assert_shard_engine_contract(_ReferenceEngine)
def test_shard_engine_is_abstract_and_cannot_be_instantiated_directly():
with pytest.raises(TypeError):
ShardEngine() # type: ignore[abstract]
def test_engine_tensor_rejects_empty_name_shape_or_dtype():
with pytest.raises(ValueError):
EngineTensor(name="", shape=(1,), dtype="bfloat16", data=b"x")
with pytest.raises(ValueError):
EngineTensor(name="t", shape=(), dtype="bfloat16", data=b"x")
with pytest.raises(ValueError):
EngineTensor(name="t", shape=(0,), dtype="bfloat16", data=b"x")
with pytest.raises(ValueError):
EngineTensor(name="t", shape=(1,), dtype="", data=b"x")
def test_boundary_bundle_requires_at_least_one_tensor():
with pytest.raises(ValueError):
BoundaryBundle(tensors=(), architecture="dense", boundary_point="pre_tail_residual")
def test_boundary_bundle_tensor_lookup_by_name():
tensor = EngineTensor(name="hidden_states", shape=(1, 1), dtype="bfloat16", data=b"\x00\x00")
bundle = BoundaryBundle(tensors=(tensor,), architecture="dense", boundary_point="pre_tail_residual")
assert bundle.tensor("hidden_states") is tensor
with pytest.raises(KeyError):
bundle.tensor("router_logits")
def test_token_output_rejects_negative_token_id():
with pytest.raises(ValueError):
TokenOutput(token_id=-1)
def test_mtp_hook_is_reserved_and_refuses_to_enable():
MtpHook() # disabled is fine
with pytest.raises(ValueError):
MtpHook(enabled=True)
with pytest.raises(ValueError):
MtpHook(draft_token_count=-1)
def test_architecture_aux_state_hook_carries_opaque_shard_local_state():
hook = ArchitectureAuxStateHook(kind="csa", state={"window": 128})
assert hook.kind == "csa"
assert hook.state == {"window": 128}
def test_prefill_and_decode_requests_require_exactly_one_input_kind():
with pytest.raises(ValueError):
PrefillRequest(session_id="s", route_epoch=0, position=0, idempotency_step=0)
with pytest.raises(ValueError):
PrefillRequest(
session_id="s",
route_epoch=0,
position=0,
idempotency_step=0,
token_ids=(1,),
input=BoundaryBundle(
tensors=(EngineTensor(name="hidden_states", shape=(1,), dtype="bfloat16", data=b"x"),),
architecture="dense",
boundary_point="pre_tail_residual",
),
)
with pytest.raises(ValueError):
DecodeRequest(session_id="s", route_epoch=0, position=0, idempotency_step=0)
def test_load_request_validates_shard_range_against_total_layers():
LoadRequest(artifact_path="a", shard_start=0, shard_end=3, total_layers=4)
with pytest.raises(ValueError):
LoadRequest(artifact_path="a", shard_start=0, shard_end=4, total_layers=4)
with pytest.raises(ValueError):
LoadRequest(artifact_path="", shard_start=0, shard_end=0, total_layers=1)
with pytest.raises(ValueError):
LoadRequest(artifact_path="a", shard_start=3, shard_end=1, total_layers=4)
def test_step_result_requires_an_output_when_status_is_ok():
with pytest.raises(ValueError):
StepResult(status=StructuredStatus(StatusCode.OK, "ok"), output=None)
# A non-OK status is allowed to carry no output.
StepResult(status=StructuredStatus(StatusCode.NOT_FOUND, "missing"), output=None)
def test_shard_engine_module_imports_no_native_or_grpc_or_wire_abi_types():
import meshnet_node.shard_engine as shard_engine_module
# The boundary module must not *import* anything that would let a
# ggml_tensor, llama context/scheduler handle, ctypes native handle, or a
# generated-protobuf (ABI) message leak into a project-owned dataclass
# field. Checking bound globals (not docstring prose) proves this
# structurally rather than by convention.
forbidden_modules = {"ctypes", "grpc", "meshnet_node.native_protocol"}
for name, value in vars(shard_engine_module).items():
module_name = getattr(value, "__name__", None)
assert module_name not in forbidden_modules, (
f"shard_engine.{name} binds forbidden module {module_name!r}"
)