894 lines
34 KiB
Python
894 lines
34 KiB
Python
"""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
|