"""Deterministic fake ``ShardEngine`` fixture (DGR-032). ``FakeShardEngine`` is a pure-Python, allocation-cheap subclass of :class:`~meshnet_node.shard_engine.ShardEngine`: no llama.cpp, no native buffers, no GPU, no filesystem or network I/O. Every prefill/decode output is a deterministic pure function of ``(loaded range, request inputs, idempotency_step)`` — hashed with SHA-256 — so replaying identical inputs on a fresh session always yields byte-identical output. It exists so worker wiring, gRPC harnesses (DGR-033), and lifecycle/session logic can be exercised end-to-end before a real llama.cpp-backed engine (DGR-037) exists. This is FIXTURE evidence only. ``EVIDENCE_CLASS`` is set to ``"fixture"`` (as opposed to ``"real"``) precisely so a later story comparing engines programmatically — DGR-036's fixture-vs-real-model parity check — can assert it is actually comparing a fixture against a real engine rather than two fixtures. This module proves lifecycle/session/epoch/fault-injection semantics; it says nothing about numerical parity with a real model. Real- model certification is DGR-036 onward (DGR-053/DGR-054 for V4 alpha). Fault injection (delay, memory pressure, malformed output, crash) is deterministic and opt-in via :class:`FakeShardEngineConfig`. Every knob defaults to off, so a bare ``FakeShardEngine()`` reproduces plain deterministic fixture behavior and passes :func:`tests.shard_engine_contract.assert_shard_engine_contract` unmodified. """ from __future__ import annotations import hashlib import time from dataclasses import dataclass from typing import Callable from .shard_engine import ( BoundaryBundle, DecodeRequest, EngineCapabilities, EngineTensor, HealthResult, LoadRequest, LoadResult, MetricsResult, PrefillRequest, ShardEngine, StepResult, TokenOutput, ) from .shard_lifecycle import CacheResult, StatusCode, StructuredStatus __all__ = ["FakeShardEngineConfig", "FakeShardEngine", "TOKEN_ID_VOCAB_SIZE", "MALFORMED_TOKEN_ID_FLOOR"] TOKEN_ID_VOCAB_SIZE = 50_000 # A malformed tail output is deterministically pushed past the fixture's own # advertised vocabulary range, so a downstream consumer checking "is this # token_id within the vocab this fixture promises" can detect it without any # extra signalling from the engine. MALFORMED_TOKEN_ID_FLOOR = 100_000_000 def _default_crash_exception() -> BaseException: return RuntimeError( "FakeShardEngine: injected crash (simulated process failure, not a StructuredStatus)" ) @dataclass(frozen=True) class FakeShardEngineConfig: """Deterministic fault-injection knobs. Every knob is off (``0``/``None``/``False``) by default. ``sleep`` is injectable so tests can assert a delay was requested without an actual process sleep; ``crash_exception_factory`` is injectable so tests can assert on a specific exception type/instance. """ step_delay_seconds: float = 0.0 sleep: Callable[[float], None] = time.sleep memory_budget_bytes: int | None = None malformed_output: bool = False crash_after_calls: int | None = None crash_exception_factory: Callable[[], BaseException] = _default_crash_exception def __post_init__(self) -> None: if self.step_delay_seconds < 0: raise ValueError("step_delay_seconds must be non-negative") if self.memory_budget_bytes is not None and self.memory_budget_bytes < 0: raise ValueError("memory_budget_bytes must be non-negative") if self.crash_after_calls is not None and self.crash_after_calls <= 0: raise ValueError("crash_after_calls must be positive when set") @dataclass class _SessionState: epoch: int cancelled: bool = False class FakeShardEngine(ShardEngine): """Deterministic fixture ``ShardEngine``. See module docstring.""" EVIDENCE_CLASS = "fixture" def __init__(self, config: FakeShardEngineConfig | None = None) -> None: self._config = config or FakeShardEngineConfig() self._loaded: LoadRequest | None = None self._sessions: dict[str, _SessionState] = {} self._cancelled_total = 0 self._generated_tokens = 0 self._call_count = 0 self._bytes_used = 0 # -- lifecycle ----------------------------------------------------- def load(self, request: LoadRequest) -> LoadResult: self._loaded = request return LoadResult( status=StructuredStatus(StatusCode.OK, "fake engine loaded"), effective_start=request.shard_start, architecture=str(request.recipe.get("architecture", "fake")), ) def capabilities(self) -> EngineCapabilities: if self._loaded is None: return EngineCapabilities( status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine 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=str(request.recipe.get("architecture", "fake")), max_concurrent_sessions=64, max_context_tokens=131072, supports_mtp=False, ) def prefill(self, request: PrefillRequest) -> StepResult: return self._step( session_id=request.session_id, route_epoch=request.route_epoch, idempotency_step=request.idempotency_step, token_ids=request.token_ids, input_bundle=request.input, cache_result_on_success=CacheResult.STORED, opens_session=True, ) def decode(self, request: DecodeRequest) -> StepResult: token_ids = (request.token_id,) if request.token_id is not None else None return self._step( session_id=request.session_id, route_epoch=request.route_epoch, idempotency_step=request.idempotency_step, token_ids=token_ids, input_bundle=request.input, cache_result_on_success=CacheResult.HIT, opens_session=False, ) def cancel(self, session_id: str, *, work_id: str = "", reason: str = "") -> StructuredStatus: session = self._sessions.get(session_id) if session is None: session = _SessionState(epoch=0) self._sessions[session_id] = session if not session.cancelled: self._cancelled_total += 1 session.cancelled = True return StructuredStatus(StatusCode.CANCELLED, reason or "fake engine: session cancelled") def release(self, session_id: str) -> StructuredStatus: self._sessions.pop(session_id, None) return StructuredStatus(StatusCode.OK, "fake engine: session released") def health(self) -> HealthResult: loaded = self._loaded is not None return HealthResult( status=StructuredStatus(StatusCode.OK, "ok"), serving=loaded, state="SERVING" if loaded else "NOT_LOADED", active_sessions=len(self._sessions), ) def metrics(self) -> MetricsResult: return MetricsResult( status=StructuredStatus(StatusCode.OK, "ok"), active_sessions=len(self._sessions), queued_frames=0, inflight_bytes=0, kv_entries=len(self._sessions), generated_tokens=self._generated_tokens, cancelled_sessions=self._cancelled_total, ) # -- shared step machinery ------------------------------------------ def _step( self, *, session_id: str, route_epoch: int, idempotency_step: int, token_ids: tuple[int, ...] | None, input_bundle: BoundaryBundle | None, cache_result_on_success: CacheResult, opens_session: bool, ) -> StepResult: if self._loaded is None: return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine not loaded")) self._call_count += 1 if self._config.crash_after_calls is not None and self._call_count == self._config.crash_after_calls: raise self._config.crash_exception_factory() session = self._sessions.get(session_id) if session is None: if not opens_session: return StepResult( status=StructuredStatus(StatusCode.NOT_FOUND, "no cached session state for decode"), cache_result=CacheResult.MISS, ) session = _SessionState(epoch=route_epoch) self._sessions[session_id] = session if session.cancelled: return StepResult(status=StructuredStatus(StatusCode.CANCELLED, "session cancelled")) if route_epoch < session.epoch: return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "stale route epoch")) session.epoch = route_epoch if self._config.step_delay_seconds: self._config.sleep(self._config.step_delay_seconds) seed = self._seed_bytes(token_ids, input_bundle) self._bytes_used += len(seed) budget = self._config.memory_budget_bytes if budget is not None and self._bytes_used > budget: return StepResult( status=StructuredStatus( StatusCode.RESOURCE_EXHAUSTED, "fake engine memory pressure budget exceeded", retryable=True, details={"memory_budget_bytes": str(budget), "bytes_used": str(self._bytes_used)}, ) ) output = self._transform(seed, idempotency_step, input_bundle) if isinstance(output, TokenOutput): self._generated_tokens += 1 return StepResult(status=StructuredStatus(StatusCode.OK, "ok"), cache_result=cache_result_on_success, output=output) @staticmethod def _seed_bytes(token_ids: tuple[int, ...] | None, bundle: BoundaryBundle | None) -> bytes: if token_ids: seed = b"".join(int(t).to_bytes(8, "big") for t in token_ids) elif bundle is not None: seed = b"".join(tensor.data for tensor in bundle.tensors) if bundle.token_id_sideband: seed += b"".join(int(t).to_bytes(8, "big") for t in bundle.token_id_sideband) else: seed = b"" return seed def _transform( self, seed: bytes, idempotency_step: int, input_bundle: BoundaryBundle | None ) -> BoundaryBundle | TokenOutput: assert self._loaded is not None digest = hashlib.sha256(seed + idempotency_step.to_bytes(8, "big")).digest() loaded = self._loaded is_tail = loaded.shard_end >= loaded.total_layers - 1 is_head = loaded.shard_start == 0 if is_tail: token_id = int.from_bytes(digest[:4], "big") % TOKEN_ID_VOCAB_SIZE if self._config.malformed_output: token_id = MALFORMED_TOKEN_ID_FLOOR + token_id return TokenOutput(token_id=token_id) boundary_point = "post_head_residual" if is_head else "post_middle_residual" architecture = ( input_bundle.architecture if input_bundle is not None else str(loaded.recipe.get("architecture", "fake")) ) data = digest if self._config.malformed_output: architecture = f"malformed:{architecture}" data = digest[:1] tensor = EngineTensor( name="hidden_states", shape=(1, max(len(seed) // 8, 1)), dtype="bfloat16", data=data, ) token_id_sideband = input_bundle.token_id_sideband if input_bundle is not None else None return BoundaryBundle( tensors=(tensor,), architecture=architecture, boundary_point=boundary_point, token_id_sideband=token_id_sideband, )