story: DGR-031 Introduce the project-owned ShardEngine interface
This commit is contained in:
Dobromir Popov
241
tests/test_shard_engine.py
Normal file
241
tests/test_shard_engine.py
Normal file
@@ -0,0 +1,241 @@
|
||||
"""DGR-031 ``ShardEngine`` contract tests.
|
||||
|
||||
``_ReferenceEngine`` below is a minimal, in-memory ``ShardEngine`` that exists
|
||||
only to prove :func:`assert_shard_engine_contract` is non-vacuous and to pin
|
||||
the abstract contract's own validation rules. It is deliberately not the
|
||||
DGR-032 deterministic fixture (delay/memory-pressure/malformed/crash
|
||||
injection, full session/epoch modeling for the fake worker) — that is a
|
||||
separate, larger story. DGR-032 and DGR-037 are expected to import
|
||||
``assert_shard_engine_contract`` from ``tests/shard_engine_contract.py``
|
||||
against their own engines.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_node.shard_engine import (
|
||||
ArchitectureAuxStateHook,
|
||||
BoundaryBundle,
|
||||
DecodeRequest,
|
||||
EngineCapabilities,
|
||||
EngineTensor,
|
||||
HealthResult,
|
||||
LoadRequest,
|
||||
LoadResult,
|
||||
MetricsResult,
|
||||
MtpHook,
|
||||
PrefillRequest,
|
||||
ShardEngine,
|
||||
StepResult,
|
||||
TokenOutput,
|
||||
)
|
||||
from meshnet_node.shard_lifecycle import CacheResult, StatusCode, StructuredStatus
|
||||
|
||||
from shard_engine_contract import assert_shard_engine_contract
|
||||
|
||||
|
||||
class _ReferenceEngine(ShardEngine):
|
||||
"""Minimal in-memory engine used only to exercise the shared contract."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._loaded: LoadRequest | None = None
|
||||
self._sessions: dict[str, dict] = {}
|
||||
self._cancelled_total = 0
|
||||
|
||||
def load(self, request: LoadRequest) -> LoadResult:
|
||||
self._loaded = request
|
||||
return LoadResult(
|
||||
status=StructuredStatus(StatusCode.OK, "loaded"),
|
||||
effective_start=request.shard_start,
|
||||
architecture="dense",
|
||||
)
|
||||
|
||||
def capabilities(self) -> EngineCapabilities:
|
||||
if self._loaded is None:
|
||||
return EngineCapabilities(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "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="dense",
|
||||
max_concurrent_sessions=8,
|
||||
max_context_tokens=131072,
|
||||
supports_mtp=False,
|
||||
)
|
||||
|
||||
def prefill(self, request: PrefillRequest) -> StepResult:
|
||||
if self._loaded is None:
|
||||
return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "engine not loaded"))
|
||||
self._sessions[request.session_id] = {"epoch": request.route_epoch, "cancelled": False}
|
||||
output = self._transform(self._seed_bytes(request.token_ids, request.input), request.idempotency_step)
|
||||
return StepResult(status=StructuredStatus(StatusCode.OK, "prefilled"), cache_result=CacheResult.STORED, output=output)
|
||||
|
||||
def decode(self, request: DecodeRequest) -> StepResult:
|
||||
session = self._sessions.get(request.session_id)
|
||||
if session is None:
|
||||
return StepResult(
|
||||
status=StructuredStatus(StatusCode.NOT_FOUND, "no cached session state"),
|
||||
cache_result=CacheResult.MISS,
|
||||
)
|
||||
if session["cancelled"]:
|
||||
return StepResult(status=StructuredStatus(StatusCode.CANCELLED, "session cancelled"))
|
||||
if request.route_epoch < session["epoch"]:
|
||||
return StepResult(status=StructuredStatus(StatusCode.FAILED_PRECONDITION, "stale route epoch"))
|
||||
session["epoch"] = request.route_epoch
|
||||
token_ids = (request.token_id,) if request.token_id is not None else None
|
||||
output = self._transform(self._seed_bytes(token_ids, request.input), request.idempotency_step)
|
||||
return StepResult(status=StructuredStatus(StatusCode.OK, "decoded"), cache_result=CacheResult.HIT, output=output)
|
||||
|
||||
def cancel(self, session_id: str, *, work_id: str = "", reason: str = "") -> StructuredStatus:
|
||||
session = self._sessions.setdefault(session_id, {"epoch": 0, "cancelled": False})
|
||||
if not session["cancelled"]:
|
||||
self._cancelled_total += 1
|
||||
session["cancelled"] = True
|
||||
return StructuredStatus(StatusCode.CANCELLED, reason or "cancelled")
|
||||
|
||||
def release(self, session_id: str) -> StructuredStatus:
|
||||
self._sessions.pop(session_id, None)
|
||||
return StructuredStatus(StatusCode.OK, "released")
|
||||
|
||||
def health(self) -> HealthResult:
|
||||
return HealthResult(
|
||||
status=StructuredStatus(StatusCode.OK, "ok"),
|
||||
serving=self._loaded is not None,
|
||||
state="SERVING" if self._loaded is not None else "NOT_LOADED",
|
||||
active_sessions=len(self._sessions),
|
||||
)
|
||||
|
||||
def metrics(self) -> MetricsResult:
|
||||
return MetricsResult(
|
||||
status=StructuredStatus(StatusCode.OK, "ok"),
|
||||
active_sessions=len(self._sessions),
|
||||
cancelled_sessions=self._cancelled_total,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _seed_bytes(token_ids, bundle: BoundaryBundle | None) -> bytes:
|
||||
if token_ids:
|
||||
return b"".join(int(t).to_bytes(4, "big") for t in token_ids)
|
||||
if bundle is not None:
|
||||
return b"".join(tensor.data for tensor in bundle.tensors)
|
||||
return b""
|
||||
|
||||
def _transform(self, seed: bytes, idempotency_step: int) -> BoundaryBundle | TokenOutput:
|
||||
digest = hashlib.sha256(seed + idempotency_step.to_bytes(4, "big")).digest()
|
||||
assert self._loaded is not None
|
||||
if self._loaded.shard_end >= self._loaded.total_layers - 1:
|
||||
token_id = int.from_bytes(digest[:4], "big") % 50_000
|
||||
return TokenOutput(token_id=token_id)
|
||||
tensor = EngineTensor(name="hidden_states", shape=(1, max(len(seed) // 4, 1)), dtype="bfloat16", data=digest)
|
||||
return BoundaryBundle(tensors=(tensor,), architecture="dense", boundary_point="pre_tail_residual")
|
||||
|
||||
|
||||
def test_reference_engine_obeys_the_shared_shard_engine_contract():
|
||||
assert_shard_engine_contract(_ReferenceEngine)
|
||||
|
||||
|
||||
def test_shard_engine_is_abstract_and_cannot_be_instantiated_directly():
|
||||
with pytest.raises(TypeError):
|
||||
ShardEngine() # type: ignore[abstract]
|
||||
|
||||
|
||||
def test_engine_tensor_rejects_empty_name_shape_or_dtype():
|
||||
with pytest.raises(ValueError):
|
||||
EngineTensor(name="", shape=(1,), dtype="bfloat16", data=b"x")
|
||||
with pytest.raises(ValueError):
|
||||
EngineTensor(name="t", shape=(), dtype="bfloat16", data=b"x")
|
||||
with pytest.raises(ValueError):
|
||||
EngineTensor(name="t", shape=(0,), dtype="bfloat16", data=b"x")
|
||||
with pytest.raises(ValueError):
|
||||
EngineTensor(name="t", shape=(1,), dtype="", data=b"x")
|
||||
|
||||
|
||||
def test_boundary_bundle_requires_at_least_one_tensor():
|
||||
with pytest.raises(ValueError):
|
||||
BoundaryBundle(tensors=(), architecture="dense", boundary_point="pre_tail_residual")
|
||||
|
||||
|
||||
def test_boundary_bundle_tensor_lookup_by_name():
|
||||
tensor = EngineTensor(name="hidden_states", shape=(1, 1), dtype="bfloat16", data=b"\x00\x00")
|
||||
bundle = BoundaryBundle(tensors=(tensor,), architecture="dense", boundary_point="pre_tail_residual")
|
||||
assert bundle.tensor("hidden_states") is tensor
|
||||
with pytest.raises(KeyError):
|
||||
bundle.tensor("router_logits")
|
||||
|
||||
|
||||
def test_token_output_rejects_negative_token_id():
|
||||
with pytest.raises(ValueError):
|
||||
TokenOutput(token_id=-1)
|
||||
|
||||
|
||||
def test_mtp_hook_is_reserved_and_refuses_to_enable():
|
||||
MtpHook() # disabled is fine
|
||||
with pytest.raises(ValueError):
|
||||
MtpHook(enabled=True)
|
||||
with pytest.raises(ValueError):
|
||||
MtpHook(draft_token_count=-1)
|
||||
|
||||
|
||||
def test_architecture_aux_state_hook_carries_opaque_shard_local_state():
|
||||
hook = ArchitectureAuxStateHook(kind="csa", state={"window": 128})
|
||||
assert hook.kind == "csa"
|
||||
assert hook.state == {"window": 128}
|
||||
|
||||
|
||||
def test_prefill_and_decode_requests_require_exactly_one_input_kind():
|
||||
with pytest.raises(ValueError):
|
||||
PrefillRequest(session_id="s", route_epoch=0, position=0, idempotency_step=0)
|
||||
with pytest.raises(ValueError):
|
||||
PrefillRequest(
|
||||
session_id="s",
|
||||
route_epoch=0,
|
||||
position=0,
|
||||
idempotency_step=0,
|
||||
token_ids=(1,),
|
||||
input=BoundaryBundle(
|
||||
tensors=(EngineTensor(name="hidden_states", shape=(1,), dtype="bfloat16", data=b"x"),),
|
||||
architecture="dense",
|
||||
boundary_point="pre_tail_residual",
|
||||
),
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
DecodeRequest(session_id="s", route_epoch=0, position=0, idempotency_step=0)
|
||||
|
||||
|
||||
def test_load_request_validates_shard_range_against_total_layers():
|
||||
LoadRequest(artifact_path="a", shard_start=0, shard_end=3, total_layers=4)
|
||||
with pytest.raises(ValueError):
|
||||
LoadRequest(artifact_path="a", shard_start=0, shard_end=4, total_layers=4)
|
||||
with pytest.raises(ValueError):
|
||||
LoadRequest(artifact_path="", shard_start=0, shard_end=0, total_layers=1)
|
||||
with pytest.raises(ValueError):
|
||||
LoadRequest(artifact_path="a", shard_start=3, shard_end=1, total_layers=4)
|
||||
|
||||
|
||||
def test_step_result_requires_an_output_when_status_is_ok():
|
||||
with pytest.raises(ValueError):
|
||||
StepResult(status=StructuredStatus(StatusCode.OK, "ok"), output=None)
|
||||
# A non-OK status is allowed to carry no output.
|
||||
StepResult(status=StructuredStatus(StatusCode.NOT_FOUND, "missing"), output=None)
|
||||
|
||||
|
||||
def test_shard_engine_module_imports_no_native_or_grpc_or_wire_abi_types():
|
||||
import meshnet_node.shard_engine as shard_engine_module
|
||||
|
||||
# The boundary module must not *import* anything that would let a
|
||||
# ggml_tensor, llama context/scheduler handle, ctypes native handle, or a
|
||||
# generated-protobuf (ABI) message leak into a project-owned dataclass
|
||||
# field. Checking bound globals (not docstring prose) proves this
|
||||
# structurally rather than by convention.
|
||||
forbidden_modules = {"ctypes", "grpc", "meshnet_node.native_protocol"}
|
||||
for name, value in vars(shard_engine_module).items():
|
||||
module_name = getattr(value, "__name__", None)
|
||||
assert module_name not in forbidden_modules, (
|
||||
f"shard_engine.{name} binds forbidden module {module_name!r}"
|
||||
)
|
||||
Reference in New Issue
Block a user