feat: checkpoint batching and release-gate stories
This commit is contained in:
Dobromir Popov
1024
packages/node/meshnet_node/batch_scheduler.py
Normal file
1024
packages/node/meshnet_node/batch_scheduler.py
Normal file
File diff suppressed because it is too large
Load Diff
893
packages/node/meshnet_node/failure_semantics.py
Normal file
893
packages/node/meshnet_node/failure_semantics.py
Normal file
@@ -0,0 +1,893 @@
|
||||
"""Bounded failure, cancellation, and restart semantics for Shard streams (DGR-013).
|
||||
|
||||
Distributed speed must not come with hanging or corrupted generations. This module
|
||||
hardens the per-Route-Session decode stream that runs over the DGR-007 Hot KV State
|
||||
manager (isolated ``(session, epoch)`` KV) and the DGR-012 continuous-batch
|
||||
scheduler. It is deliberately backend-agnostic and pure-Python: it drives the same
|
||||
``KvBoundaryAdapter`` the default deterministic gate uses, so the whole matrix stays
|
||||
download-free, GPU-free, and API-credit-free while exercising the *real* KV
|
||||
isolation path (the pinned llama.cpp worker, DGR-008, implements the identical
|
||||
adapter contract).
|
||||
|
||||
The guarantees, mapped to the story's acceptance criteria:
|
||||
|
||||
* **Deadlines and heartbeat/health loss terminate blocked stream operations.**
|
||||
:class:`DeadlineGuard` bounds every step against an absolute deadline and a
|
||||
heartbeat-timeout; when either is breached it raises :class:`StreamTerminated`
|
||||
so a blocked stream never hangs.
|
||||
* **Cancellation propagates across every Shard and releases local KV and queued
|
||||
buffers.** :class:`ShardCancellationGroup` fans a single cancel across every
|
||||
node-local KV manager serving a Route Session and releases queued activation
|
||||
buffers; the DGR-012 scheduler's :meth:`~meshnet_node.batch_scheduler.
|
||||
ContinuousBatchScheduler.cancel` drops queued/active work on this node.
|
||||
* **Duplicate steps are idempotent; uncertain mutations are never replayed
|
||||
silently.** :class:`IdempotencyLedger` records each committed
|
||||
``(session, epoch, step)`` and returns the recorded token for a duplicate
|
||||
delivery instead of re-running it. A step whose outcome is *uncertain* (the
|
||||
worker died mid-mutation) is marked uncertain and can never be silently
|
||||
replayed — a replay attempt raises :class:`UncertainMutationError`, forcing an
|
||||
explicit verify-or-restart.
|
||||
* **Alpha failover restarts from token zero on a newly compatible route rather
|
||||
than importing unverified KV.** :class:`RestartController` opens a *new* route
|
||||
epoch, releases every shard's prior-epoch KV, and the restart re-prefills the
|
||||
whole prompt from token zero. The old epoch becomes stale (rejected by the KV
|
||||
manager); unverified KV is never migrated (RALPH runtime decision #14).
|
||||
* **Billing/work records distinguish completed, cancelled, failed, and unverified
|
||||
work.** :class:`WorkLedger` records a typed :class:`WorkRecord` per attempt;
|
||||
only :attr:`WorkStatus.COMPLETED` records are billable, so cancelled, failed,
|
||||
and uncertain (unverified) work is accounted but never charged.
|
||||
|
||||
:class:`HardenedSessionRunner` composes these into one drivable stream: it runs a
|
||||
single session's prefill+decode through the adapter under a deadline/heartbeat
|
||||
guard and a cancellation token, records the typed work outcome, and — via
|
||||
:meth:`HardenedSessionRunner.run_with_failover` — restarts a transient failure
|
||||
from token zero on a fresh epoch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field, replace
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Mapping, Sequence
|
||||
|
||||
from meshnet_node.batch_scheduler import DoneReason, GenerationRequest
|
||||
from meshnet_node.boundary_adapter import BoundaryContractError, TailOutput
|
||||
from meshnet_node.hot_kv_state import (
|
||||
CacheMiss,
|
||||
HotKvStateManager,
|
||||
IncompatibleCacheRecipeError,
|
||||
KvBoundaryAdapter,
|
||||
KvCacheMissError,
|
||||
StaleRouteEpochError,
|
||||
)
|
||||
|
||||
|
||||
class FailureSemanticsError(RuntimeError):
|
||||
"""Base class for failure/cancellation/restart errors."""
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Typed outcomes: failure kinds and billing/work statuses.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class FailureKind(str, Enum):
|
||||
"""Why a stream step failed. Stable strings for the protocol's structured status."""
|
||||
|
||||
# Bounded termination of a blocked op.
|
||||
DEADLINE_EXCEEDED = "deadline-exceeded"
|
||||
HEARTBEAT_LOST = "heartbeat-lost"
|
||||
# Transport / worker loss (transient — a restart from token zero may succeed).
|
||||
WORKER_DEATH = "worker-death"
|
||||
STREAM_RESET = "stream-reset"
|
||||
# Protocol violations (deterministic — a restart would fail identically).
|
||||
MALFORMED_BUNDLE = "malformed-bundle"
|
||||
STALE_EPOCH = "stale-epoch"
|
||||
INCOMPATIBLE_RECIPE = "incompatible-recipe"
|
||||
# KV state expected by the caller is gone; re-prefill from token zero.
|
||||
CACHE_MISS = "cache-miss"
|
||||
# Explicit client cancellation.
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
# Failure kinds that a from-token-zero restart on a fresh route may recover from.
|
||||
# A protocol violation or an explicit bound (deadline/cancel) is NOT restartable —
|
||||
# retrying it would hang or fail identically, so we surface it instead.
|
||||
_RESTARTABLE = frozenset(
|
||||
{
|
||||
FailureKind.WORKER_DEATH,
|
||||
FailureKind.STREAM_RESET,
|
||||
FailureKind.CACHE_MISS,
|
||||
}
|
||||
)
|
||||
|
||||
# Failure kinds whose mutation outcome is *uncertain* — the KV may or may not have
|
||||
# advanced, so the confirmed work is billed as UNVERIFIED and never replayed
|
||||
# silently. Only an *unexpected* error raised while a step was executing is
|
||||
# uncertain (mapped to WORKER_DEATH). A stream reset, deadline, or cache miss
|
||||
# detected at a step boundary is certain: nothing committed for that step.
|
||||
_UNCERTAIN = frozenset({FailureKind.WORKER_DEATH})
|
||||
|
||||
|
||||
class WorkStatus(str, Enum):
|
||||
"""The billing-relevant outcome class of a unit of work (AC: billing records).
|
||||
|
||||
Only :attr:`COMPLETED` work is billable. Cancelled, failed, and unverified
|
||||
work is recorded distinctly so a client is never charged for a generation that
|
||||
hung, was cancelled, or whose mutations could not be verified.
|
||||
"""
|
||||
|
||||
COMPLETED = "completed"
|
||||
CANCELLED = "cancelled"
|
||||
FAILED = "failed"
|
||||
UNVERIFIED = "unverified"
|
||||
|
||||
|
||||
def work_status_for(kind: FailureKind) -> WorkStatus:
|
||||
"""Map a terminal failure kind to its billing/work status."""
|
||||
if kind is FailureKind.CANCELLED:
|
||||
return WorkStatus.CANCELLED
|
||||
if kind in _UNCERTAIN:
|
||||
return WorkStatus.UNVERIFIED
|
||||
return WorkStatus.FAILED
|
||||
|
||||
|
||||
def classify_exception(exc: BaseException) -> FailureKind:
|
||||
"""Classify a raised error into a :class:`FailureKind`.
|
||||
|
||||
Protocol violations map to their specific kind; a :class:`StreamTerminated`
|
||||
carries its own kind; any *unexpected* error is treated as worker death
|
||||
(an uncertain, transient loss), never silently ignored.
|
||||
"""
|
||||
if isinstance(exc, StreamTerminated):
|
||||
return exc.kind
|
||||
if isinstance(exc, OperationCancelled):
|
||||
return FailureKind.CANCELLED
|
||||
if isinstance(exc, StaleRouteEpochError):
|
||||
return FailureKind.STALE_EPOCH
|
||||
if isinstance(exc, IncompatibleCacheRecipeError):
|
||||
return FailureKind.INCOMPATIBLE_RECIPE
|
||||
if isinstance(exc, BoundaryContractError):
|
||||
return FailureKind.MALFORMED_BUNDLE
|
||||
if isinstance(exc, KvCacheMissError):
|
||||
return FailureKind.CACHE_MISS
|
||||
return FailureKind.WORKER_DEATH
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Deadlines and heartbeat/health loss.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class StreamTerminated(FailureSemanticsError):
|
||||
"""A blocked stream op was terminated by a deadline or heartbeat/health loss."""
|
||||
|
||||
def __init__(self, kind: FailureKind, detail: str = "") -> None:
|
||||
self.kind = kind
|
||||
self.detail = detail
|
||||
suffix = f": {detail}" if detail else ""
|
||||
super().__init__(f"stream terminated ({kind.value}){suffix}")
|
||||
|
||||
|
||||
class OperationCancelled(FailureSemanticsError):
|
||||
"""Raised when a step observes its :class:`CancellationToken` is cancelled."""
|
||||
|
||||
def __init__(self, reason: str = "client-cancel") -> None:
|
||||
self.reason = reason
|
||||
super().__init__(f"operation cancelled: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeadlineGuard:
|
||||
"""Bounds a blocked stream op against an absolute deadline and heartbeat loss.
|
||||
|
||||
``deadline`` is an absolute time on ``clock``'s scale (``None`` disables it).
|
||||
``heartbeat_timeout`` is the maximum tolerated gap since the last observed
|
||||
heartbeat; when the peer stops sending heartbeats (its health is lost) the gap
|
||||
grows past the timeout and :meth:`check` raises rather than blocking forever.
|
||||
Both bounds are checked with an injected ``clock`` so the matrix is
|
||||
deterministic.
|
||||
"""
|
||||
|
||||
deadline: float | None = None
|
||||
heartbeat_timeout: float | None = None
|
||||
clock: Callable[[], float] = time.monotonic
|
||||
_last_heartbeat: float = field(default=0.0, init=False)
|
||||
_started: bool = field(default=False, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.heartbeat_timeout is not None and self.heartbeat_timeout <= 0:
|
||||
raise FailureSemanticsError("heartbeat_timeout must be positive")
|
||||
|
||||
def start(self) -> None:
|
||||
self._last_heartbeat = self.clock()
|
||||
self._started = True
|
||||
|
||||
def heartbeat(self) -> None:
|
||||
"""Record that the peer is alive (resets the heartbeat gap)."""
|
||||
self._last_heartbeat = self.clock()
|
||||
|
||||
def check(self) -> None:
|
||||
"""Raise :class:`StreamTerminated` if the deadline or heartbeat lapsed."""
|
||||
if not self._started:
|
||||
self.start()
|
||||
now = self.clock()
|
||||
if self.deadline is not None and now >= self.deadline:
|
||||
raise StreamTerminated(
|
||||
FailureKind.DEADLINE_EXCEEDED,
|
||||
f"deadline {self.deadline} reached at {now}",
|
||||
)
|
||||
if self.heartbeat_timeout is not None:
|
||||
gap = now - self._last_heartbeat
|
||||
if gap > self.heartbeat_timeout:
|
||||
raise StreamTerminated(
|
||||
FailureKind.HEARTBEAT_LOST,
|
||||
f"no heartbeat for {gap} > {self.heartbeat_timeout}",
|
||||
)
|
||||
|
||||
def remaining(self) -> float | None:
|
||||
if self.deadline is None:
|
||||
return None
|
||||
return self.deadline - self.clock()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Cancellation that propagates across shards and releases KV + queued buffers.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class CancellationToken:
|
||||
"""A thread-safe one-shot cancellation flag shared by a Route Session's steps."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._cancelled = False
|
||||
self._reason = ""
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def cancel(self, reason: str = "client-cancel") -> None:
|
||||
with self._lock:
|
||||
if not self._cancelled:
|
||||
self._cancelled = True
|
||||
self._reason = reason
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
with self._lock:
|
||||
return self._cancelled
|
||||
|
||||
@property
|
||||
def reason(self) -> str:
|
||||
with self._lock:
|
||||
return self._reason
|
||||
|
||||
def raise_if_cancelled(self) -> None:
|
||||
with self._lock:
|
||||
if self._cancelled:
|
||||
raise OperationCancelled(self._reason)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CancellationOutcome:
|
||||
"""What a :meth:`ShardCancellationGroup.cancel` released (for observability)."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
shards_released: int
|
||||
buffers_released: int
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"route_epoch": self.route_epoch,
|
||||
"shards_released": self.shards_released,
|
||||
"buffers_released": self.buffers_released,
|
||||
}
|
||||
|
||||
|
||||
class ShardCancellationGroup:
|
||||
"""Fan one cancellation across every node-local Shard of a Route Session.
|
||||
|
||||
A Route Session spans a chain of Shards, each with its own local Hot KV State
|
||||
manager (KV is never migrated between nodes). Cancelling the session must free
|
||||
*all* of that state: this group releases the ``(session, epoch)`` KV on every
|
||||
registered manager and invokes every registered queued-buffer release callback
|
||||
(the pending activation bundles a node holds for the session). Release is
|
||||
idempotent, so cancelling twice is safe.
|
||||
"""
|
||||
|
||||
def __init__(self, session_id: str, route_epoch: int) -> None:
|
||||
if not isinstance(session_id, str) or not session_id.strip():
|
||||
raise FailureSemanticsError("session_id must be a non-empty string")
|
||||
self.session_id = session_id
|
||||
self.route_epoch = int(route_epoch)
|
||||
self._managers: list[HotKvStateManager] = []
|
||||
self._buffers: list[Callable[[], None]] = []
|
||||
self._lock = threading.Lock()
|
||||
self._cancelled = False
|
||||
|
||||
def add_shard(self, manager: HotKvStateManager) -> "ShardCancellationGroup":
|
||||
with self._lock:
|
||||
self._managers.append(manager)
|
||||
return self
|
||||
|
||||
def add_queued_buffer(
|
||||
self, release: Callable[[], None]
|
||||
) -> "ShardCancellationGroup":
|
||||
"""Register a queued activation buffer's release callback."""
|
||||
with self._lock:
|
||||
self._buffers.append(release)
|
||||
return self
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
with self._lock:
|
||||
return self._cancelled
|
||||
|
||||
def cancel(self) -> CancellationOutcome:
|
||||
"""Release every shard's KV and every queued buffer for this session."""
|
||||
with self._lock:
|
||||
managers = list(self._managers)
|
||||
buffers = list(self._buffers)
|
||||
self._buffers.clear()
|
||||
self._cancelled = True
|
||||
shards_released = 0
|
||||
for manager in managers:
|
||||
if manager.release(self.session_id, self.route_epoch):
|
||||
shards_released += 1
|
||||
buffers_released = 0
|
||||
for release in buffers:
|
||||
release()
|
||||
buffers_released += 1
|
||||
return CancellationOutcome(
|
||||
session_id=self.session_id,
|
||||
route_epoch=self.route_epoch,
|
||||
shards_released=shards_released,
|
||||
buffers_released=buffers_released,
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Idempotency: duplicate steps are no-ops; uncertain mutations never replay.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class StepPhase(str, Enum):
|
||||
IN_FLIGHT = "in-flight"
|
||||
COMMITTED = "committed"
|
||||
UNCERTAIN = "uncertain"
|
||||
|
||||
|
||||
class UncertainMutationError(FailureSemanticsError):
|
||||
"""Raised when a caller tries to replay a step whose outcome is uncertain.
|
||||
|
||||
A step is uncertain when its mutation may or may not have been applied (worker
|
||||
death / stream reset mid-append). Replaying it silently could double-apply KV
|
||||
or bill unverified work, so the ledger refuses: the caller must verify against
|
||||
the actual KV length or restart from token zero on a fresh epoch instead.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StepKey:
|
||||
"""Identity of one idempotent stream step within a route epoch."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
step_index: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StepDisposition:
|
||||
"""What :meth:`IdempotencyLedger.begin` decided for a step."""
|
||||
|
||||
fresh: bool
|
||||
token: int | None = None
|
||||
|
||||
@property
|
||||
def duplicate(self) -> bool:
|
||||
return not self.fresh
|
||||
|
||||
|
||||
class IdempotencyLedger:
|
||||
"""Records committed/uncertain stream steps so duplicates never re-mutate.
|
||||
|
||||
Keyed by ``(session, epoch, step_index)`` — the protocol's idempotency step.
|
||||
|
||||
* :meth:`begin` on a *fresh* key marks it in-flight and returns "execute".
|
||||
* :meth:`begin` on a *committed* key returns the recorded token so a duplicate
|
||||
delivery is a no-op (idempotent replay).
|
||||
* :meth:`begin` on an *in-flight* or *uncertain* key raises
|
||||
:class:`UncertainMutationError` — a concurrent duplicate or a replay of an
|
||||
unverified mutation is never silently applied.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._phase: dict[StepKey, StepPhase] = {}
|
||||
self._token: dict[StepKey, int] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def begin(self, key: StepKey) -> StepDisposition:
|
||||
with self._lock:
|
||||
phase = self._phase.get(key)
|
||||
if phase is None:
|
||||
self._phase[key] = StepPhase.IN_FLIGHT
|
||||
return StepDisposition(fresh=True)
|
||||
if phase is StepPhase.COMMITTED:
|
||||
return StepDisposition(fresh=False, token=self._token[key])
|
||||
# IN_FLIGHT (concurrent duplicate) or UNCERTAIN (post-crash replay):
|
||||
# both are unverified and must not be silently re-applied.
|
||||
raise UncertainMutationError(
|
||||
f"step {key.step_index} for session {key.session_id[:8]} epoch "
|
||||
f"{key.route_epoch} is {phase.value}; refusing silent replay"
|
||||
)
|
||||
|
||||
def commit(self, key: StepKey, token: int) -> None:
|
||||
with self._lock:
|
||||
self._phase[key] = StepPhase.COMMITTED
|
||||
self._token[key] = int(token)
|
||||
|
||||
def mark_uncertain(self, key: StepKey, detail: str = "") -> None:
|
||||
with self._lock:
|
||||
# A committed step is verified; never downgrade it.
|
||||
if self._phase.get(key) is StepPhase.COMMITTED:
|
||||
return
|
||||
self._phase[key] = StepPhase.UNCERTAIN
|
||||
|
||||
def phase_of(self, key: StepKey) -> StepPhase | None:
|
||||
with self._lock:
|
||||
return self._phase.get(key)
|
||||
|
||||
def committed_token(self, key: StepKey) -> int | None:
|
||||
with self._lock:
|
||||
return self._token.get(key)
|
||||
|
||||
def has_uncertain(self) -> bool:
|
||||
with self._lock:
|
||||
return any(p is StepPhase.UNCERTAIN for p in self._phase.values())
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Restart / alpha failover: from token zero on a fresh compatible route.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
class RestartController:
|
||||
"""Alpha failover that restarts from token zero, never importing prior KV.
|
||||
|
||||
RALPH runtime decision #14: when the alpha (the head owning embedding + final
|
||||
head) fails, the route retries from token zero; unverified KV is never
|
||||
migrated. :meth:`failover` opens the *next* route epoch and releases every
|
||||
node-local shard's prior-epoch KV, so the restart begins with empty caches. The
|
||||
KV manager then treats the failed epoch as stale (a later reference to it is
|
||||
rejected), which is what keeps a half-computed cache from being reused.
|
||||
"""
|
||||
|
||||
def __init__(self, managers: Sequence[HotKvStateManager]) -> None:
|
||||
self._managers = list(managers)
|
||||
|
||||
def failover(self, session_id: str, failed_epoch: int) -> int:
|
||||
"""Advance to a fresh epoch and drop the failed epoch's KV on every shard."""
|
||||
new_epoch = int(failed_epoch) + 1
|
||||
for manager in self._managers:
|
||||
manager.release(session_id, failed_epoch)
|
||||
return new_epoch
|
||||
|
||||
def assert_fresh_start(self, session_id: str, new_epoch: int) -> None:
|
||||
"""Verify no shard carries KV for the new epoch (a true token-zero restart).
|
||||
|
||||
Any residual KV under the new epoch would be unverified imported state;
|
||||
this fails closed so a restart can never silently attend over it.
|
||||
"""
|
||||
for manager in self._managers:
|
||||
result = manager.resolve(session_id, new_epoch)
|
||||
if not isinstance(result, CacheMiss):
|
||||
raise FailureSemanticsError(
|
||||
f"restart epoch {new_epoch} for session {session_id[:8]} is not "
|
||||
"empty; refusing to import unverified KV"
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Billing / work records.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkRecord:
|
||||
"""A typed unit of served work, distinguishing what may be billed.
|
||||
|
||||
``tokens`` counts only *committed* generated tokens. Only a
|
||||
:attr:`WorkStatus.COMPLETED` record is billable; cancelled/failed/unverified
|
||||
records carry their confirmed token count for observability but are excluded
|
||||
from billing so uncompleted or unverified work is never charged.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
status: WorkStatus
|
||||
tokens: int
|
||||
failure_kind: FailureKind | None = None
|
||||
detail: str = ""
|
||||
|
||||
@property
|
||||
def billable(self) -> bool:
|
||||
return self.status is WorkStatus.COMPLETED
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"route_epoch": self.route_epoch,
|
||||
"status": self.status.value,
|
||||
"tokens": self.tokens,
|
||||
"failure_kind": self.failure_kind.value if self.failure_kind else None,
|
||||
"detail": self.detail,
|
||||
"billable": self.billable,
|
||||
}
|
||||
|
||||
|
||||
class WorkLedger:
|
||||
"""Append-only ledger of :class:`WorkRecord`, split by billing status."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._records: list[WorkRecord] = []
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def record(self, record: WorkRecord) -> WorkRecord:
|
||||
with self._lock:
|
||||
self._records.append(record)
|
||||
return record
|
||||
|
||||
def records(self) -> list[WorkRecord]:
|
||||
with self._lock:
|
||||
return list(self._records)
|
||||
|
||||
def records_for(self, session_id: str) -> list[WorkRecord]:
|
||||
with self._lock:
|
||||
return [r for r in self._records if r.session_id == session_id]
|
||||
|
||||
def billable_records(self) -> list[WorkRecord]:
|
||||
with self._lock:
|
||||
return [r for r in self._records if r.billable]
|
||||
|
||||
def billable_tokens(self) -> int:
|
||||
"""Total tokens that may be charged (completed work only)."""
|
||||
with self._lock:
|
||||
return sum(r.tokens for r in self._records if r.billable)
|
||||
|
||||
def counts_by_status(self) -> dict[str, int]:
|
||||
counts: dict[str, int] = {s.value: 0 for s in WorkStatus}
|
||||
with self._lock:
|
||||
for record in self._records:
|
||||
counts[record.status.value] += 1
|
||||
return counts
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
with self._lock:
|
||||
records = [r.to_dict() for r in self._records]
|
||||
counts: dict[str, int] = {s.value: 0 for s in WorkStatus}
|
||||
for record in records:
|
||||
counts[record["status"]] += 1
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"records": records,
|
||||
"counts_by_status": counts,
|
||||
"billable_tokens": sum(r["tokens"] for r in records if r["billable"]),
|
||||
}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# The hardened single-session stream runner.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunOutcome:
|
||||
"""The typed result of one hardened generation attempt."""
|
||||
|
||||
session_id: str
|
||||
route_epoch: int
|
||||
status: WorkStatus
|
||||
tokens: tuple[int, ...]
|
||||
failure_kind: FailureKind | None
|
||||
detail: str
|
||||
|
||||
@property
|
||||
def completed(self) -> bool:
|
||||
return self.status is WorkStatus.COMPLETED
|
||||
|
||||
@property
|
||||
def token_count(self) -> int:
|
||||
return len(self.tokens)
|
||||
|
||||
@property
|
||||
def restartable(self) -> bool:
|
||||
return self.failure_kind in _RESTARTABLE
|
||||
|
||||
def work_record(self) -> WorkRecord:
|
||||
return WorkRecord(
|
||||
session_id=self.session_id,
|
||||
route_epoch=self.route_epoch,
|
||||
status=self.status,
|
||||
tokens=len(self.tokens),
|
||||
failure_kind=self.failure_kind,
|
||||
detail=self.detail,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FailoverResult:
|
||||
"""The result of a run that may have restarted from token zero after a failure."""
|
||||
|
||||
outcome: RunOutcome
|
||||
attempts: tuple[RunOutcome, ...]
|
||||
restarts: int
|
||||
|
||||
@property
|
||||
def completed(self) -> bool:
|
||||
return self.outcome.completed
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"final_status": self.outcome.status.value,
|
||||
"final_epoch": self.outcome.route_epoch,
|
||||
"restarts": self.restarts,
|
||||
"attempts": [
|
||||
{
|
||||
"route_epoch": a.route_epoch,
|
||||
"status": a.status.value,
|
||||
"failure_kind": a.failure_kind.value if a.failure_kind else None,
|
||||
"tokens": a.token_count,
|
||||
}
|
||||
for a in self.attempts
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class HardenedSessionRunner:
|
||||
"""Drive one Route Session's decode stream with bounded failure semantics.
|
||||
|
||||
The runner owns a single full-shard :class:`KvBoundaryAdapter` (head **and**
|
||||
tail, so a step samples a token) and threads every DGR-013 guarantee through a
|
||||
step loop:
|
||||
|
||||
* every step is bounded by a :class:`DeadlineGuard` and can observe a
|
||||
:class:`CancellationToken`;
|
||||
* every step is idempotent through an :class:`IdempotencyLedger` (a duplicate
|
||||
returns the recorded token; an uncertain mutation is never replayed);
|
||||
* any failure releases this session's KV (cancellation) and is recorded as a
|
||||
typed :class:`WorkRecord` in the :class:`WorkLedger`;
|
||||
* :meth:`run_with_failover` restarts a transient failure from token zero on a
|
||||
fresh epoch via a :class:`RestartController`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
adapter: KvBoundaryAdapter,
|
||||
*,
|
||||
clock: Callable[[], float] | None = None,
|
||||
work_ledger: WorkLedger | None = None,
|
||||
idempotency: IdempotencyLedger | None = None,
|
||||
) -> None:
|
||||
if not (adapter.is_head and adapter.is_tail):
|
||||
raise FailureSemanticsError(
|
||||
"HardenedSessionRunner needs a full (head+tail) shard so decode "
|
||||
"steps sample tokens; got a partial range "
|
||||
f"(head={adapter.is_head} tail={adapter.is_tail})"
|
||||
)
|
||||
self._adapter = adapter
|
||||
self._manager: HotKvStateManager = adapter.manager
|
||||
self._clock = clock or time.monotonic
|
||||
self.work_ledger = work_ledger or WorkLedger()
|
||||
self.idempotency = idempotency or IdempotencyLedger()
|
||||
|
||||
# -- single attempt -------------------------------------------------------
|
||||
|
||||
def run(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
*,
|
||||
deadline: float | None = None,
|
||||
heartbeat_timeout: float | None = None,
|
||||
cancel_token: CancellationToken | None = None,
|
||||
heartbeat: Callable[[int], bool] | None = None,
|
||||
before_step: Callable[[int], None] | None = None,
|
||||
) -> RunOutcome:
|
||||
"""Run one attempt of ``request``; record and return a typed outcome.
|
||||
|
||||
``deadline`` (absolute, on the injected clock) and ``heartbeat_timeout``
|
||||
bound blocked steps. ``cancel_token`` lets a client cancel mid-stream.
|
||||
``heartbeat(step)`` returns ``True`` when a heartbeat was heard before that
|
||||
step (resetting the health timer); ``before_step(step)`` is a fault-
|
||||
injection / clock-advance hook run before each step and may raise
|
||||
:class:`StreamTerminated` (e.g. a stream reset) or
|
||||
:class:`OperationCancelled`.
|
||||
"""
|
||||
sid = request.session_id
|
||||
epoch = request.route_epoch
|
||||
guard = DeadlineGuard(
|
||||
deadline=deadline,
|
||||
heartbeat_timeout=heartbeat_timeout,
|
||||
clock=self._clock,
|
||||
)
|
||||
guard.start()
|
||||
tokens: list[int] = []
|
||||
current_key: StepKey | None = None
|
||||
try:
|
||||
# step 0 is the prefill (emits the first token); steps 1..N are decodes.
|
||||
for step_index in range(request.max_new_tokens):
|
||||
# before_step is the fault-injection / clock-advance hook and may
|
||||
# itself terminate the step (stream reset, cancel); run it first so
|
||||
# a fault it raises takes effect on this step, then re-check the
|
||||
# bounds it may have advanced (deadline / heartbeat / cancel).
|
||||
if before_step is not None:
|
||||
before_step(step_index)
|
||||
if cancel_token is not None:
|
||||
cancel_token.raise_if_cancelled()
|
||||
if heartbeat is not None and heartbeat(step_index):
|
||||
guard.heartbeat()
|
||||
guard.check()
|
||||
|
||||
current_key = StepKey(sid, epoch, step_index)
|
||||
disposition = self.idempotency.begin(current_key)
|
||||
if disposition.duplicate:
|
||||
# Idempotent replay: reuse the recorded token, do not re-mutate.
|
||||
assert disposition.token is not None
|
||||
tokens.append(disposition.token)
|
||||
continue
|
||||
|
||||
token = self._execute_step(request, step_index, tokens)
|
||||
if isinstance(token, CacheMiss):
|
||||
# The expected KV was gone; the append never started, so this is
|
||||
# a certain (not uncertain) miss — restartable from token zero.
|
||||
return self._finish_failure(
|
||||
request,
|
||||
tokens,
|
||||
FailureKind.CACHE_MISS,
|
||||
str(token),
|
||||
cancel_token,
|
||||
)
|
||||
self.idempotency.commit(current_key, token)
|
||||
tokens.append(token)
|
||||
except (StreamTerminated, OperationCancelled) as exc:
|
||||
return self._finish_failure(
|
||||
request, tokens, classify_exception(exc), str(exc), cancel_token
|
||||
)
|
||||
except (
|
||||
BoundaryContractError,
|
||||
StaleRouteEpochError,
|
||||
IncompatibleCacheRecipeError,
|
||||
KvCacheMissError,
|
||||
) as exc:
|
||||
# Deterministic protocol/state errors, all validated before any KV
|
||||
# append committed — certain, not uncertain.
|
||||
return self._finish_failure(
|
||||
request, tokens, classify_exception(exc), str(exc), cancel_token
|
||||
)
|
||||
except UncertainMutationError as exc:
|
||||
# A replay of an unverified step reached the ledger — never silent.
|
||||
return self._finish_failure(
|
||||
request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - unexpected == worker death
|
||||
# An unexpected error mid-step may have left the KV half-mutated; mark
|
||||
# the step uncertain so it can never be silently replayed, then fail
|
||||
# closed as unverified work.
|
||||
if current_key is not None:
|
||||
self.idempotency.mark_uncertain(current_key, str(exc))
|
||||
return self._finish_failure(
|
||||
request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token
|
||||
)
|
||||
|
||||
return self._finish_completed(request, tokens)
|
||||
|
||||
def _execute_step(
|
||||
self, request: GenerationRequest, step_index: int, tokens: list[int]
|
||||
) -> int | CacheMiss:
|
||||
sid = request.session_id
|
||||
epoch = request.route_epoch
|
||||
if step_index == 0:
|
||||
out = self._adapter.prefill(
|
||||
sid, epoch, token_ids=list(request.prompt_token_ids)
|
||||
)
|
||||
else:
|
||||
# expected_seq_len defends the KV layer against a desynchronised decode:
|
||||
# prompt positions plus the tokens already committed this run.
|
||||
expected = request.prompt_len + (step_index - 1)
|
||||
out = self._adapter.decode(
|
||||
sid,
|
||||
epoch,
|
||||
token_ids=[tokens[-1]],
|
||||
expected_seq_len=expected,
|
||||
)
|
||||
if isinstance(out, CacheMiss):
|
||||
return out
|
||||
if not isinstance(out, TailOutput):
|
||||
raise FailureSemanticsError(
|
||||
"full-shard step did not yield a sampled token; got "
|
||||
f"{type(out).__name__}"
|
||||
)
|
||||
return int(out.token_id)
|
||||
|
||||
# -- failover across restarts --------------------------------------------
|
||||
|
||||
def run_with_failover(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
controller: RestartController,
|
||||
*,
|
||||
max_restarts: int = 3,
|
||||
**run_kwargs: Any,
|
||||
) -> FailoverResult:
|
||||
"""Run ``request``, restarting a transient failure from token zero.
|
||||
|
||||
On a restartable failure (worker death, stream reset, cache miss) the
|
||||
controller advances to a fresh epoch and drops the failed epoch's KV; the
|
||||
next attempt re-prefills the whole prompt from token zero. A deterministic
|
||||
failure (deadline, cancel, malformed bundle, stale epoch) is returned as-is
|
||||
— retrying it would hang or fail identically. Per-attempt fault-injection
|
||||
hooks (``before_step`` / ``heartbeat``) are only applied to the *first*
|
||||
attempt so a restart runs clean.
|
||||
"""
|
||||
if max_restarts < 0:
|
||||
raise FailureSemanticsError("max_restarts must be >= 0")
|
||||
epoch = request.route_epoch
|
||||
attempts: list[RunOutcome] = []
|
||||
first_kwargs = run_kwargs
|
||||
for attempt in range(max_restarts + 1):
|
||||
attempt_request = replace(request, route_epoch=epoch)
|
||||
kwargs = first_kwargs if attempt == 0 else {}
|
||||
outcome = self.run(attempt_request, **kwargs)
|
||||
attempts.append(outcome)
|
||||
if outcome.completed or not outcome.restartable or attempt == max_restarts:
|
||||
return FailoverResult(
|
||||
outcome=outcome, attempts=tuple(attempts), restarts=attempt
|
||||
)
|
||||
# Alpha failover: fresh epoch, drop prior-epoch KV on every shard, and
|
||||
# verify the new epoch starts empty (no unverified KV import).
|
||||
epoch = controller.failover(request.session_id, epoch)
|
||||
controller.assert_fresh_start(request.session_id, epoch)
|
||||
# Unreachable: the loop always returns, but keep the type-checker happy.
|
||||
raise FailureSemanticsError("run_with_failover exhausted without returning")
|
||||
|
||||
# -- outcome bookkeeping --------------------------------------------------
|
||||
|
||||
def _finish_completed(
|
||||
self, request: GenerationRequest, tokens: list[int]
|
||||
) -> RunOutcome:
|
||||
outcome = RunOutcome(
|
||||
session_id=request.session_id,
|
||||
route_epoch=request.route_epoch,
|
||||
status=WorkStatus.COMPLETED,
|
||||
tokens=tuple(tokens),
|
||||
failure_kind=None,
|
||||
detail="",
|
||||
)
|
||||
self.work_ledger.record(outcome.work_record())
|
||||
return outcome
|
||||
|
||||
def _finish_failure(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
tokens: list[int],
|
||||
kind: FailureKind,
|
||||
detail: str,
|
||||
cancel_token: CancellationToken | None,
|
||||
) -> RunOutcome:
|
||||
# Cancellation semantics: release this session's local KV so a failed or
|
||||
# cancelled stream never leaks its cache. release() is idempotent.
|
||||
self._manager.release(request.session_id, request.route_epoch)
|
||||
if cancel_token is not None and kind is not FailureKind.CANCELLED:
|
||||
# Ensure downstream shards sharing the token also stop.
|
||||
cancel_token.cancel(kind.value)
|
||||
outcome = RunOutcome(
|
||||
session_id=request.session_id,
|
||||
route_epoch=request.route_epoch,
|
||||
status=work_status_for(kind),
|
||||
tokens=tuple(tokens),
|
||||
failure_kind=kind,
|
||||
detail=detail,
|
||||
)
|
||||
self.work_ledger.record(outcome.work_record())
|
||||
return outcome
|
||||
Reference in New Issue
Block a user