373 lines
13 KiB
Python
373 lines
13 KiB
Python
"""The project-owned ``ShardEngine`` contract (DGR-031).
|
|
|
|
A worker process (the gRPC surface in ``shard_runtime_server.py``, or any
|
|
future transport) never talks to llama.cpp directly. It talks to a
|
|
``ShardEngine``. This module is the *only* place that boundary is defined, and
|
|
every operation on it is built from project-owned dataclasses and plain
|
|
Python values (``str``, ``int``, ``bytes``, ``Mapping``) — never a
|
|
``ggml_tensor``, a llama context/scheduler handle, or a generated-protobuf
|
|
(ABI) message. A fake fixture engine (DGR-032) and a real llama.cpp-backed
|
|
engine (DGR-037) are both, structurally, nothing more than subclasses of
|
|
:class:`ShardEngine`; the worker code that calls them does not change when one
|
|
replaces the other.
|
|
|
|
This is deliberately a fourth, distinct layer from the three that already
|
|
exist:
|
|
|
|
- ``native_protocol`` — the generated gRPC/Protobuf wire ABI (DGR-021/024).
|
|
- ``protocol.ActivationEnvelope`` — the versioned wire envelope for activation
|
|
traffic between shard *hops* over the network (DGR-021).
|
|
- ``shard_lifecycle`` — the versioned RPC/session lifecycle contract a
|
|
generated gRPC binding consumes (DGR-022).
|
|
|
|
``ShardEngine`` sits *inside* one worker process, below all three: it is the
|
|
seam between "the code that speaks Meshnet's wire protocol" and "the code
|
|
that actually runs model layers." It reuses :class:`~meshnet_node.shard_lifecycle.StructuredStatus`,
|
|
:class:`~meshnet_node.shard_lifecycle.StatusCode`, :class:`~meshnet_node.shard_lifecycle.CacheExpectation`,
|
|
and :class:`~meshnet_node.shard_lifecycle.CacheResult` rather than inventing a
|
|
parallel status vocabulary, since those are already project-owned and
|
|
version-stable.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import abc
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Mapping
|
|
|
|
from .shard_lifecycle import (
|
|
CacheExpectation,
|
|
CacheResult,
|
|
StatusCode,
|
|
StructuredStatus,
|
|
)
|
|
|
|
__all__ = [
|
|
"EngineError",
|
|
"EngineTensor",
|
|
"BoundaryBundle",
|
|
"TokenOutput",
|
|
"MtpHook",
|
|
"ArchitectureAuxStateHook",
|
|
"LoadRequest",
|
|
"LoadResult",
|
|
"EngineCapabilities",
|
|
"PrefillRequest",
|
|
"DecodeRequest",
|
|
"StepResult",
|
|
"HealthResult",
|
|
"MetricsResult",
|
|
"ShardEngine",
|
|
]
|
|
|
|
|
|
class EngineError(RuntimeError):
|
|
"""An engine-boundary failure represented by a structured status.
|
|
|
|
Mirrors :class:`~meshnet_node.shard_lifecycle.LifecycleContractError`:
|
|
callers pattern-match on ``error.status.code`` rather than on exception
|
|
subclasses, so a fake and a real engine can fail the exact same way for
|
|
the exact same reason.
|
|
"""
|
|
|
|
def __init__(self, status: StructuredStatus) -> None:
|
|
self.status = status
|
|
super().__init__(status.message)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EngineTensor:
|
|
"""One named tensor crossing the engine boundary.
|
|
|
|
Intentionally not a ``ggml_tensor`` or a framework tensor object: ``data``
|
|
is plain owned bytes, ``shape``/``dtype`` are plain metadata. An
|
|
implementation constructs this from whatever internal representation it
|
|
uses (a ``torch.Tensor``, a llama.cpp buffer, a synthetic fixture array)
|
|
without leaking that representation across the boundary.
|
|
"""
|
|
|
|
name: str
|
|
shape: tuple[int, ...]
|
|
dtype: str
|
|
data: bytes
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.name:
|
|
raise ValueError("engine tensor requires a name")
|
|
if not self.shape or any(dim <= 0 for dim in self.shape):
|
|
raise ValueError("engine tensor shape must be a non-empty tuple of positive ints")
|
|
if not self.dtype:
|
|
raise ValueError("engine tensor requires a dtype")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BoundaryBundle:
|
|
"""A named-tensor activation crossing a shard boundary (head/middle/tail-in).
|
|
|
|
``token_id_sideband`` carries token IDs alongside the activation only
|
|
where the architecture boundary requires them (V4's first three
|
|
hash-routed MoE layers); it is ``None`` everywhere else. Per-shard hot
|
|
KV/recurrent/CSA/HCA/SWA/indexer/compressor state never appears here — it
|
|
stays local to a shard via :class:`ArchitectureAuxStateHook` and is never
|
|
part of what crosses the wire.
|
|
"""
|
|
|
|
tensors: tuple[EngineTensor, ...]
|
|
architecture: str
|
|
boundary_point: str
|
|
token_id_sideband: tuple[int, ...] | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.tensors:
|
|
raise ValueError("boundary bundle requires at least one tensor")
|
|
if not self.architecture:
|
|
raise ValueError("boundary bundle requires an architecture name")
|
|
if not self.boundary_point:
|
|
raise ValueError("boundary bundle requires a boundary point name")
|
|
|
|
def tensor(self, name: str) -> EngineTensor:
|
|
for tensor in self.tensors:
|
|
if tensor.name == name:
|
|
return tensor
|
|
raise KeyError(name)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TokenOutput:
|
|
"""A tail shard's sampled decode result.
|
|
|
|
Never a raw logits tensor: the engine boundary only ever hands back the
|
|
already-sampled token (mirroring
|
|
:meth:`meshnet_node.architecture_boundary.TailOutput.sampled_token`, which
|
|
likewise refuses anything but a sampled token id).
|
|
"""
|
|
|
|
token_id: int
|
|
text: str | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.token_id < 0:
|
|
raise ValueError("sampled token id must be non-negative")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MtpHook:
|
|
"""Reserved multi-token-prediction hook — typed, but refused when enabled.
|
|
|
|
RALPH-CONTEXT is explicit that "MTP is reserved and off for alpha; its
|
|
ownership contract, implementation, and benchmark are required before
|
|
beta" (DGR-065/DGR-066). Reserving the shape now means DGR-037's real
|
|
engine and DGR-051's V4 adapter do not have to change this dataclass's
|
|
field layout later; they only flip ``enabled`` once DGR-066 lands.
|
|
"""
|
|
|
|
enabled: bool = False
|
|
draft_token_count: int = 0
|
|
aux_state: Mapping[str, Any] | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.enabled:
|
|
raise ValueError(
|
|
"MTP is reserved and must remain disabled before DGR-066; "
|
|
"this hook exists to fix its shape, not to enable it"
|
|
)
|
|
if self.draft_token_count < 0:
|
|
raise ValueError("draft_token_count must be non-negative")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArchitectureAuxStateHook:
|
|
"""Reserved per-shard architecture auxiliary-state hook.
|
|
|
|
Covers V4's CSA/HCA/SWA/indexer/compressor state and any other
|
|
architecture-local state a future adapter needs. RALPH-CONTEXT locks this
|
|
as shard-local, keyed by route session/epoch, and explicitly never carried
|
|
over the WAN seam — so this hook has no wire encoding of its own and must
|
|
never be embedded inside a :class:`BoundaryBundle`.
|
|
"""
|
|
|
|
kind: str = ""
|
|
state: Mapping[str, Any] | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LoadRequest:
|
|
"""One exact artifact/recipe/range identity for a worker to load."""
|
|
|
|
artifact_path: str
|
|
shard_start: int
|
|
shard_end: int
|
|
total_layers: int
|
|
recipe: Mapping[str, Any] = field(default_factory=dict)
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.artifact_path:
|
|
raise ValueError("load request requires an artifact path")
|
|
if self.shard_start < 0 or self.shard_end < self.shard_start:
|
|
raise ValueError("shard_start must be <= shard_end and non-negative")
|
|
if self.total_layers <= self.shard_end:
|
|
raise ValueError("total_layers must exceed shard_end (shard_end is inclusive)")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LoadResult:
|
|
status: StructuredStatus
|
|
effective_start: int = 0
|
|
architecture: str = ""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EngineCapabilities:
|
|
status: StructuredStatus
|
|
shard_start: int = 0
|
|
shard_end: int = 0
|
|
effective_start: int = 0
|
|
total_layers: int = 0
|
|
architecture: str = ""
|
|
max_concurrent_sessions: int = 0
|
|
max_context_tokens: int = 0
|
|
supports_mtp: bool = False
|
|
|
|
@property
|
|
def is_head(self) -> bool:
|
|
return self.shard_start == 0
|
|
|
|
@property
|
|
def is_tail(self) -> bool:
|
|
return self.shard_end >= self.total_layers - 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PrefillRequest:
|
|
"""A prefill step. Exactly one of ``token_ids`` (head) or ``input`` (middle/tail) is set."""
|
|
|
|
session_id: str
|
|
route_epoch: int
|
|
position: int
|
|
idempotency_step: int
|
|
token_ids: tuple[int, ...] | None = None
|
|
input: BoundaryBundle | None = None
|
|
cache_expectation: CacheExpectation = CacheExpectation.NONE
|
|
mtp: MtpHook = field(default_factory=MtpHook)
|
|
architecture_aux_state: ArchitectureAuxStateHook | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
_require_exactly_one_input(self.token_ids, self.input)
|
|
if not self.session_id:
|
|
raise ValueError("prefill request requires a session id")
|
|
if self.route_epoch < 0 or self.position < 0 or self.idempotency_step < 0:
|
|
raise ValueError("route_epoch, position, and idempotency_step must be non-negative")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DecodeRequest:
|
|
"""A decode step. Exactly one of ``token_id`` (head) or ``input`` (middle/tail) is set."""
|
|
|
|
session_id: str
|
|
route_epoch: int
|
|
position: int
|
|
idempotency_step: int
|
|
token_id: int | None = None
|
|
input: BoundaryBundle | None = None
|
|
mtp: MtpHook = field(default_factory=MtpHook)
|
|
architecture_aux_state: ArchitectureAuxStateHook | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
_require_exactly_one_input(
|
|
None if self.token_id is None else (self.token_id,), self.input
|
|
)
|
|
if not self.session_id:
|
|
raise ValueError("decode request requires a session id")
|
|
if self.route_epoch < 0 or self.position < 0 or self.idempotency_step < 0:
|
|
raise ValueError("route_epoch, position, and idempotency_step must be non-negative")
|
|
|
|
|
|
def _require_exactly_one_input(
|
|
token_ids: tuple[int, ...] | None, bundle: BoundaryBundle | None
|
|
) -> None:
|
|
if (token_ids is None) == (bundle is None):
|
|
raise ValueError("exactly one of token ids or a boundary bundle must be set")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StepResult:
|
|
"""The result of a prefill or decode step.
|
|
|
|
``output`` is a :class:`BoundaryBundle` for a head/middle shard handing an
|
|
activation to the next hop, or a :class:`TokenOutput` for a tail shard
|
|
that sampled a token. It is ``None`` only when ``status.code`` is not
|
|
``OK``.
|
|
"""
|
|
|
|
status: StructuredStatus
|
|
cache_result: CacheResult = CacheResult.NOT_REQUESTED
|
|
output: BoundaryBundle | TokenOutput | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.status.code is StatusCode.OK and self.output is None:
|
|
raise ValueError("a successful step result must carry an output")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HealthResult:
|
|
status: StructuredStatus
|
|
serving: bool = False
|
|
state: str = "UNKNOWN"
|
|
active_sessions: int = 0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MetricsResult:
|
|
status: StructuredStatus
|
|
active_sessions: int = 0
|
|
queued_frames: int = 0
|
|
inflight_bytes: int = 0
|
|
kv_entries: int = 0
|
|
generated_tokens: int = 0
|
|
cancelled_sessions: int = 0
|
|
|
|
|
|
class ShardEngine(abc.ABC):
|
|
"""The contract every shard execution engine (fake or real) must implement.
|
|
|
|
Every method returns a project-owned result carrying a
|
|
:class:`~meshnet_node.shard_lifecycle.StructuredStatus` rather than
|
|
raising for expected, protocol-visible outcomes (a cache miss, a stale
|
|
epoch, an unknown session); an :class:`EngineError` is reserved for
|
|
genuine programming errors at the call site (malformed request objects),
|
|
which the request dataclasses' own ``__post_init__`` validation already
|
|
catches before an implementation ever sees them.
|
|
"""
|
|
|
|
@abc.abstractmethod
|
|
def load(self, request: LoadRequest) -> LoadResult:
|
|
"""Load one exact artifact/recipe/range identity. Idempotent per engine instance."""
|
|
|
|
@abc.abstractmethod
|
|
def capabilities(self) -> EngineCapabilities:
|
|
"""Report this engine's authoritative range and limits after ``load``."""
|
|
|
|
@abc.abstractmethod
|
|
def prefill(self, request: PrefillRequest) -> StepResult:
|
|
"""Run one prefill step for a session."""
|
|
|
|
@abc.abstractmethod
|
|
def decode(self, request: DecodeRequest) -> StepResult:
|
|
"""Run one decode step for a session."""
|
|
|
|
@abc.abstractmethod
|
|
def cancel(self, session_id: str, *, work_id: str = "", reason: str = "") -> StructuredStatus:
|
|
"""Cancel a session (or one work item within it) in flight."""
|
|
|
|
@abc.abstractmethod
|
|
def release(self, session_id: str) -> StructuredStatus:
|
|
"""Release a session's held state. Idempotent."""
|
|
|
|
@abc.abstractmethod
|
|
def health(self) -> HealthResult:
|
|
"""Report liveness/serving state. Must never raise."""
|
|
|
|
@abc.abstractmethod
|
|
def metrics(self) -> MetricsResult:
|
|
"""Report point-in-time operational counters. Must never raise."""
|