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

tests/shard_engine_contract.py

@@ -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