303 lines
12 KiB
Python
303 lines
12 KiB
Python
"""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,
|
|
)
|