story: DGR-032 Implement deterministic fake ShardEngine
This commit is contained in:
Dobromir Popov
302
packages/node/meshnet_node/fake_shard_engine.py
Normal file
302
packages/node/meshnet_node/fake_shard_engine.py
Normal file
@@ -0,0 +1,302 @@
|
||||
"""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,
|
||||
)
|
||||
Reference in New Issue
Block a user