neuron-tai/tests/shard_engine_contract.py

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