feat: define shard lifecycle status contract
This commit is contained in:
Dobromir Popov
472
packages/node/meshnet_node/shard_lifecycle.py
Normal file
472
packages/node/meshnet_node/shard_lifecycle.py
Normal file
@@ -0,0 +1,472 @@
|
||||
"""Versioned Shard lifecycle and structured status contract.
|
||||
|
||||
This module defines the semantic contract consumed by a future generated gRPC
|
||||
binding. It deliberately contains no Meshnet routing, authentication policy,
|
||||
billing, or llama.cpp types.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
import json
|
||||
import threading
|
||||
from typing import Any, Callable, ClassVar
|
||||
|
||||
|
||||
SCHEMA_NAME = "meshnet.shard-runtime"
|
||||
SCHEMA_VERSION = 1
|
||||
SUPPORTED_VERSIONS = frozenset({SCHEMA_VERSION})
|
||||
DEFAULT_MAX_INFLIGHT_FRAMES = 32
|
||||
DEFAULT_MAX_INFLIGHT_BYTES = 8 * 1024 * 1024
|
||||
|
||||
|
||||
class RpcName(str, Enum):
|
||||
CAPABILITY = "GetCapability"
|
||||
HEALTH = "CheckHealth"
|
||||
SESSION = "OpenSession"
|
||||
CANCEL = "CancelSession"
|
||||
RELEASE = "ReleaseSession"
|
||||
METRICS = "GetMetrics"
|
||||
|
||||
|
||||
class StatusCode(str, Enum):
|
||||
OK = "OK"
|
||||
INVALID_ARGUMENT = "INVALID_ARGUMENT"
|
||||
UNSUPPORTED_VERSION = "UNSUPPORTED_VERSION"
|
||||
FAILED_PRECONDITION = "FAILED_PRECONDITION"
|
||||
MALFORMED_LIFECYCLE = "MALFORMED_LIFECYCLE"
|
||||
NOT_FOUND = "NOT_FOUND"
|
||||
ALREADY_EXISTS = "ALREADY_EXISTS"
|
||||
CANCELLED = "CANCELLED"
|
||||
DEADLINE_EXCEEDED = "DEADLINE_EXCEEDED"
|
||||
RESOURCE_EXHAUSTED = "RESOURCE_EXHAUSTED"
|
||||
UNAUTHENTICATED = "UNAUTHENTICATED"
|
||||
PERMISSION_DENIED = "PERMISSION_DENIED"
|
||||
DATA_LOSS = "DATA_LOSS"
|
||||
UNAVAILABLE = "UNAVAILABLE"
|
||||
INTERNAL = "INTERNAL"
|
||||
|
||||
|
||||
class LifecycleState(str, Enum):
|
||||
OPEN = "OPEN"
|
||||
PREFILLING = "PREFILLING"
|
||||
DECODING = "DECODING"
|
||||
CANCELLING = "CANCELLING"
|
||||
CANCELLED = "CANCELLED"
|
||||
RELEASING = "RELEASING"
|
||||
RELEASED = "RELEASED"
|
||||
FAILED = "FAILED"
|
||||
|
||||
|
||||
class CacheExpectation(str, Enum):
|
||||
NONE = "NONE"
|
||||
OPTIONAL = "OPTIONAL"
|
||||
REQUIRED = "REQUIRED"
|
||||
|
||||
|
||||
class CacheResult(str, Enum):
|
||||
NOT_REQUESTED = "NOT_REQUESTED"
|
||||
HIT = "HIT"
|
||||
MISS = "MISS"
|
||||
INVALIDATED = "INVALIDATED"
|
||||
STORED = "STORED"
|
||||
|
||||
|
||||
class SessionPhase(str, Enum):
|
||||
PREFILL = "PREFILL"
|
||||
DECODE = "DECODE"
|
||||
|
||||
|
||||
class LifecycleContractError(ValueError):
|
||||
"""A protocol violation represented by a structured status."""
|
||||
|
||||
def __init__(self, status: "StructuredStatus") -> None:
|
||||
self.status = status
|
||||
super().__init__(status.message)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StructuredStatus:
|
||||
code: StatusCode
|
||||
message: str
|
||||
retryable: bool = False
|
||||
details: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"code": self.code.value,
|
||||
"message": self.message,
|
||||
"retryable": self.retryable,
|
||||
"details": dict(sorted(self.details.items())),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, value: dict[str, Any]) -> "StructuredStatus":
|
||||
return cls(
|
||||
code=StatusCode(str(value["code"])),
|
||||
message=str(value["message"]),
|
||||
retryable=bool(value.get("retryable", False)),
|
||||
details={str(k): str(v) for k, v in value.get("details", {}).items()},
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DeadlinePolicy:
|
||||
capability_seconds: float = 5.0
|
||||
health_seconds: float = 2.0
|
||||
session_open_seconds: float = 10.0
|
||||
session_idle_seconds: float = 30.0
|
||||
cancel_seconds: float = 2.0
|
||||
release_seconds: float = 5.0
|
||||
metrics_seconds: float = 5.0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if any(value <= 0 for value in self.__dict__.values()):
|
||||
raise ValueError("all RPC deadlines must be positive")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TlsAuthHooks:
|
||||
"""Transport hooks only; Meshnet remains the identity and billing authority."""
|
||||
|
||||
tls_required: bool = True
|
||||
server_name: str = "shard.meshnet"
|
||||
client_certificate_hook: str | None = None
|
||||
peer_identity_hook: str | None = None
|
||||
auth_metadata_hook: str | None = None
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.tls_required and not self.server_name:
|
||||
raise ValueError("TLS server_name is required when TLS is enabled")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FlowControlLimits:
|
||||
max_inflight_frames: int = DEFAULT_MAX_INFLIGHT_FRAMES
|
||||
max_inflight_bytes: int = DEFAULT_MAX_INFLIGHT_BYTES
|
||||
max_frame_bytes: int = 1024 * 1024
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if min(self.max_inflight_frames, self.max_inflight_bytes, self.max_frame_bytes) <= 0:
|
||||
raise ValueError("flow-control limits must be positive")
|
||||
|
||||
|
||||
class FlowControl:
|
||||
"""Bounded sender window; callers block or fail instead of growing unbounded."""
|
||||
|
||||
def __init__(self, limits: FlowControlLimits = FlowControlLimits()) -> None:
|
||||
self.limits = limits
|
||||
self._condition = threading.Condition()
|
||||
self._frames = 0
|
||||
self._bytes = 0
|
||||
|
||||
@property
|
||||
def outstanding(self) -> tuple[int, int]:
|
||||
with self._condition:
|
||||
return self._frames, self._bytes
|
||||
|
||||
def acquire(self, size: int, *, wait: bool = False, cancelled: Callable[[], bool] | None = None) -> None:
|
||||
if size < 0 or size > self.limits.max_frame_bytes:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.RESOURCE_EXHAUSTED, "frame exceeds bounded flow-control window",
|
||||
details={"max_frame_bytes": str(self.limits.max_frame_bytes)},
|
||||
))
|
||||
with self._condition:
|
||||
while (
|
||||
self._frames >= self.limits.max_inflight_frames
|
||||
or self._bytes + size > self.limits.max_inflight_bytes
|
||||
):
|
||||
if cancelled and cancelled():
|
||||
raise LifecycleContractError(StructuredStatus(StatusCode.CANCELLED, "flow-control wait cancelled"))
|
||||
if not wait:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.RESOURCE_EXHAUSTED, "flow-control window is full", retryable=True,
|
||||
))
|
||||
self._condition.wait(timeout=0.05)
|
||||
self._frames += 1
|
||||
self._bytes += size
|
||||
|
||||
def release(self, size: int) -> None:
|
||||
with self._condition:
|
||||
self._frames = max(0, self._frames - 1)
|
||||
self._bytes = max(0, self._bytes - max(0, size))
|
||||
self._condition.notify_all()
|
||||
|
||||
|
||||
class CancellationToken:
|
||||
def __init__(self) -> None:
|
||||
self._event = threading.Event()
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
return self._event.is_set()
|
||||
|
||||
def cancel(self) -> None:
|
||||
self._event.set()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CapabilityRequest:
|
||||
schema_version: int = SCHEMA_VERSION
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CapabilityResponse:
|
||||
status: StructuredStatus
|
||||
schema_version: int
|
||||
supported_versions: tuple[int, ...]
|
||||
artifact_fingerprint: str = ""
|
||||
runtime_fingerprint: str = ""
|
||||
shard_start: int = 0
|
||||
shard_end: int = 0
|
||||
effective_start: int = 0
|
||||
max_sessions: int = 0
|
||||
max_frame_bytes: int = DEFAULT_MAX_INFLIGHT_BYTES
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.status.code is StatusCode.OK and self.schema_version not in SUPPORTED_VERSIONS:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.UNSUPPORTED_VERSION, "worker selected an unsupported schema version",
|
||||
))
|
||||
if self.shard_end < self.shard_start or self.effective_start < self.shard_start:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.INVALID_ARGUMENT, "invalid authoritative shard range",
|
||||
))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HealthRequest:
|
||||
schema_version: int = SCHEMA_VERSION
|
||||
include_metrics: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HealthResponse:
|
||||
status: StructuredStatus
|
||||
serving: bool
|
||||
state: str
|
||||
active_sessions: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionRequest:
|
||||
schema_version: int
|
||||
request_id: str
|
||||
work_id: str
|
||||
route_session: str
|
||||
route_epoch: int
|
||||
artifact_fingerprint: str
|
||||
runtime_fingerprint: str
|
||||
shard_start: int
|
||||
shard_end: int
|
||||
effective_start: int
|
||||
cache_expectation: CacheExpectation = CacheExpectation.NONE
|
||||
deadline_seconds: float = 30.0
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.schema_version not in SUPPORTED_VERSIONS:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.UNSUPPORTED_VERSION, "unsupported session schema version",
|
||||
details={"requested": str(self.schema_version), "supported": "1"},
|
||||
))
|
||||
if not self.request_id or not self.work_id or not self.route_session:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.INVALID_ARGUMENT, "request, work, and route-session IDs are required",
|
||||
))
|
||||
if self.route_epoch < 0 or self.shard_start < 0 or self.shard_end <= self.shard_start:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.INVALID_ARGUMENT, "invalid route epoch or shard range",
|
||||
))
|
||||
if not self.shard_start <= self.effective_start <= self.shard_end:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.INVALID_ARGUMENT, "effective start must be inside the shard range",
|
||||
))
|
||||
if self.deadline_seconds <= 0:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.INVALID_ARGUMENT, "session deadline must be positive",
|
||||
))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionFrame:
|
||||
phase: SessionPhase
|
||||
position: int
|
||||
idempotency_step: int
|
||||
payload: Any
|
||||
cache_expectation: CacheExpectation = CacheExpectation.NONE
|
||||
size_bytes: int = 0
|
||||
|
||||
def validate(self, limits: FlowControlLimits) -> None:
|
||||
if self.position < 0 or self.idempotency_step < 0:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.INVALID_ARGUMENT, "position and idempotency step must be non-negative",
|
||||
))
|
||||
if self.size_bytes < 0 or self.size_bytes > limits.max_frame_bytes:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.RESOURCE_EXHAUSTED, "session frame exceeds max_frame_bytes",
|
||||
))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionResult:
|
||||
status: StructuredStatus
|
||||
cache_result: CacheResult = CacheResult.NOT_REQUESTED
|
||||
position: int = 0
|
||||
idempotency_step: int = 0
|
||||
payload: Any = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CancelRequest:
|
||||
schema_version: int
|
||||
request_id: str
|
||||
work_id: str
|
||||
route_session: str
|
||||
route_epoch: int
|
||||
reason: str = ""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReleaseRequest:
|
||||
schema_version: int
|
||||
request_id: str
|
||||
work_id: str
|
||||
route_session: str
|
||||
route_epoch: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetricsRequest:
|
||||
schema_version: int = SCHEMA_VERSION
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetricsResponse:
|
||||
status: StructuredStatus
|
||||
active_sessions: int
|
||||
queued_frames: int
|
||||
inflight_bytes: int
|
||||
kv_entries: int
|
||||
generated_tokens: int
|
||||
cancelled_sessions: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionLifecycle:
|
||||
"""Fail-closed state machine for one Route Session Activation Seam."""
|
||||
|
||||
request: SessionRequest
|
||||
state: LifecycleState = LifecycleState.OPEN
|
||||
cancellation: CancellationToken = field(default_factory=CancellationToken)
|
||||
last_idempotency_step: int = -1
|
||||
_seen_steps: set[int] = field(default_factory=set, init=False, repr=False)
|
||||
|
||||
_ALLOWED: ClassVar[dict[LifecycleState, frozenset[LifecycleState]]] = {
|
||||
LifecycleState.OPEN: frozenset({LifecycleState.PREFILLING, LifecycleState.CANCELLING, LifecycleState.RELEASING, LifecycleState.FAILED}),
|
||||
LifecycleState.PREFILLING: frozenset({LifecycleState.PREFILLING, LifecycleState.DECODING, LifecycleState.CANCELLING, LifecycleState.RELEASING, LifecycleState.FAILED}),
|
||||
LifecycleState.DECODING: frozenset({LifecycleState.DECODING, LifecycleState.CANCELLING, LifecycleState.RELEASING, LifecycleState.FAILED}),
|
||||
LifecycleState.CANCELLING: frozenset({LifecycleState.CANCELLED, LifecycleState.RELEASING, LifecycleState.FAILED}),
|
||||
LifecycleState.CANCELLED: frozenset({LifecycleState.RELEASING, LifecycleState.RELEASED}),
|
||||
LifecycleState.RELEASING: frozenset({LifecycleState.RELEASED, LifecycleState.FAILED}),
|
||||
LifecycleState.RELEASED: frozenset(),
|
||||
LifecycleState.FAILED: frozenset({LifecycleState.RELEASING, LifecycleState.RELEASED}),
|
||||
}
|
||||
|
||||
def _transition(self, target: LifecycleState) -> None:
|
||||
if target not in self._ALLOWED[self.state]:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.MALFORMED_LIFECYCLE,
|
||||
f"cannot transition from {self.state.value} to {target.value}",
|
||||
details={"state": self.state.value, "target": target.value},
|
||||
))
|
||||
self.state = target
|
||||
|
||||
def apply(self, frame: SessionFrame) -> None:
|
||||
frame.validate(FlowControlLimits())
|
||||
if self.cancellation.cancelled and frame.phase is not SessionPhase.PREFILL:
|
||||
raise LifecycleContractError(StructuredStatus(StatusCode.CANCELLED, "session cancellation propagated"))
|
||||
if frame.idempotency_step in self._seen_steps:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.ALREADY_EXISTS, "duplicate idempotency step", details={"step": str(frame.idempotency_step)},
|
||||
))
|
||||
if frame.idempotency_step <= self.last_idempotency_step:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.MALFORMED_LIFECYCLE, "idempotency steps must increase monotonically",
|
||||
))
|
||||
target = LifecycleState.PREFILLING if frame.phase is SessionPhase.PREFILL else LifecycleState.DECODING
|
||||
self._transition(target)
|
||||
self._seen_steps.add(frame.idempotency_step)
|
||||
self.last_idempotency_step = frame.idempotency_step
|
||||
|
||||
def cancel(self, reason: str = "") -> StructuredStatus:
|
||||
if self.state in {LifecycleState.RELEASED, LifecycleState.RELEASING}:
|
||||
return StructuredStatus(StatusCode.FAILED_PRECONDITION, "session is already releasing or released")
|
||||
if self.state is LifecycleState.CANCELLED:
|
||||
return StructuredStatus(StatusCode.OK, "session already cancelled")
|
||||
self._transition(LifecycleState.CANCELLING)
|
||||
self.cancellation.cancel()
|
||||
self._transition(LifecycleState.CANCELLED)
|
||||
return StructuredStatus(StatusCode.CANCELLED, reason or "session cancelled")
|
||||
|
||||
def release(self) -> StructuredStatus:
|
||||
if self.state is LifecycleState.RELEASED:
|
||||
return StructuredStatus(StatusCode.OK, "session already released")
|
||||
self._transition(LifecycleState.RELEASING)
|
||||
self.cancellation.cancel()
|
||||
self._transition(LifecycleState.RELEASED)
|
||||
return StructuredStatus(StatusCode.OK, "session released")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShardRpcContract:
|
||||
"""Service/method and operational rules for generated gRPC bindings."""
|
||||
|
||||
schema: str = SCHEMA_NAME
|
||||
version: int = SCHEMA_VERSION
|
||||
methods: tuple[RpcName, ...] = tuple(RpcName)
|
||||
deadlines: DeadlinePolicy = DeadlinePolicy()
|
||||
flow_control: FlowControlLimits = FlowControlLimits()
|
||||
tls_auth: TlsAuthHooks = TlsAuthHooks()
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.schema != SCHEMA_NAME or self.version not in SUPPORTED_VERSIONS:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.UNSUPPORTED_VERSION, "unsupported Shard RPC contract version",
|
||||
))
|
||||
if self.methods != tuple(RpcName):
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.FAILED_PRECONDITION, "contract must expose the complete lifecycle RPC set",
|
||||
))
|
||||
self.tls_auth.validate()
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
self.validate()
|
||||
data = {
|
||||
"schema": self.schema,
|
||||
"version": self.version,
|
||||
"methods": [method.value for method in self.methods],
|
||||
"deadlines": self.deadlines.__dict__,
|
||||
"flow_control": self.flow_control.__dict__,
|
||||
"tls_auth": self.tls_auth.__dict__,
|
||||
}
|
||||
return json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, raw: bytes) -> "ShardRpcContract":
|
||||
data = json.loads(raw)
|
||||
if int(data.get("version", -1)) not in SUPPORTED_VERSIONS:
|
||||
raise LifecycleContractError(StructuredStatus(
|
||||
StatusCode.UNSUPPORTED_VERSION, "unsupported Shard RPC contract version",
|
||||
))
|
||||
methods = tuple(RpcName(item) for item in data.get("methods", ()))
|
||||
contract = cls(
|
||||
schema=str(data.get("schema", "")),
|
||||
version=int(data["version"]),
|
||||
methods=methods,
|
||||
deadlines=DeadlinePolicy(**data.get("deadlines", {})),
|
||||
flow_control=FlowControlLimits(**data.get("flow_control", {})),
|
||||
tls_auth=TlsAuthHooks(**data.get("tls_auth", {})),
|
||||
)
|
||||
contract.validate()
|
||||
return contract
|
||||
Reference in New Issue
Block a user