Compare commits
5 Commits
989b55970b
...
a35d86f343
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a35d86f343 | ||
|
|
9b257d9a1b | ||
|
|
3611b2cf9e | ||
|
|
efd1cf4ef6 | ||
|
|
ab466ce6b6 |
101
.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md
Normal file
101
.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md
Normal file
@@ -0,0 +1,101 @@
|
||||
# DGR-021 evidence — versioned named-tensor activation envelope
|
||||
|
||||
**Completed:** 2026-07-17
|
||||
**Branch:** `distributed-gguf-runtime`
|
||||
**Authority:** `.scratch/distributed-gguf-runtime/prd.json`
|
||||
**Dependency:** DGR-018 (`evidence/DGR-018/README.md`) — canonical backlog schema / issue projection contract
|
||||
|
||||
## Objective
|
||||
|
||||
Establish the backend-neutral activation envelope used by direct and relayed Shard traffic, with stable versioning, named tensors, bounded fragmentation, checksum validation, and reserved extensibility for future state.
|
||||
|
||||
## Changes
|
||||
|
||||
### `packages/node/meshnet_node/protocol.py` (new)
|
||||
|
||||
Added a self-contained activation-envelope module with:
|
||||
|
||||
- `SCHEMA_NAME = "meshnet.activation-stream"` and `SCHEMA_VERSION = 1`
|
||||
- `TensorFragment`
|
||||
- bounded byte fragments with offset, compression tag, checksum, and extension preservation
|
||||
- deterministic `to_dict()` / `from_dict()` round-trip
|
||||
- `NamedTensor`
|
||||
- named tensor metadata: `name`, `shape`, `dtype`, `byte_order`, `compression`, `checksum`, `fragments`
|
||||
- fragmentation via `from_bytes(..., max_fragment_bytes=...)`
|
||||
- checksum validation over reconstructed tensor bytes
|
||||
- unknown-field preservation via `extensions`
|
||||
- `ActivationEnvelope`
|
||||
- top-level fields for `request_id`, `work_id`, `route_session`, `route_epoch`, `shard_start`, `effective_start`, `phase`, `position`, and `idempotency_step`
|
||||
- reserved extension fields for `token_id_sideband`, `architecture_state`, `recurrent_state`, and `mtp`
|
||||
- deterministic canonical serialization (`to_bytes`) and round-trip parsing (`from_bytes`)
|
||||
- size-limit enforcement (`to_bytes(max_bytes=...)`)
|
||||
- conversion from a live `TensorPayload` into the envelope and back again
|
||||
|
||||
### `packages/node/meshnet_node/model_backend.py`
|
||||
|
||||
Extended `TensorPayload` with envelope conversion helpers:
|
||||
|
||||
- `TensorPayload.to_envelope(...)`
|
||||
- `TensorPayload.from_envelope(...)`
|
||||
|
||||
These keep the existing activation payload interface intact while exposing the new versioned envelope as the shared protocol layer.
|
||||
|
||||
### `tests/test_activation_envelope.py` (new)
|
||||
|
||||
Added focused deterministic tests covering:
|
||||
|
||||
- deterministic envelope serialization and round-trip parsing
|
||||
- tensor fragmentation and checksum validation
|
||||
- unknown-field preservation at both envelope and tensor levels
|
||||
- size-limit rejection
|
||||
- `TensorPayload` ↔ envelope round-trip
|
||||
|
||||
### `.scratch/distributed-gguf-runtime/prd.json`
|
||||
|
||||
Marked `DGR-021.passes = true` and added completion notes recording the envelope implementation and verification commands.
|
||||
|
||||
## Commands and results
|
||||
|
||||
```bash
|
||||
pytest -q tests/test_activation_envelope.py
|
||||
```
|
||||
|
||||
```text
|
||||
5 passed in 0.06s
|
||||
```
|
||||
|
||||
```bash
|
||||
pytest -q tests/test_activation_envelope.py tests/test_kv_cache_distributed.py -k 'session_is_stable_and_decode_payloads_are_single_token or large_prefill_activation_survives_zstd_compressed_hop'
|
||||
```
|
||||
|
||||
```text
|
||||
.. [100%]
|
||||
2 passed, 21 deselected in 1.84s
|
||||
```
|
||||
|
||||
```bash
|
||||
python3 -m compileall packages/node/meshnet_node tests/test_activation_envelope.py
|
||||
```
|
||||
|
||||
```text
|
||||
Listing 'packages/node/meshnet_node'...
|
||||
Listing 'packages/node/meshnet_node/native_protocol'...
|
||||
Compiling 'tests/test_activation_envelope.py'...
|
||||
```
|
||||
|
||||
```bash
|
||||
git diff --check
|
||||
```
|
||||
|
||||
```text
|
||||
No whitespace errors
|
||||
```
|
||||
|
||||
## Limitations
|
||||
|
||||
- The envelope is implemented as a canonical deterministic JSON contract with dataclasses and conversion hooks, not generated `.proto` classes. The environment had `protobuf` available but not the `grpc_tools` generation toolchain, so I did not materialize a compiled proto artifact here.
|
||||
- The direct/relayed HTTP/WebSocket transports remain byte-oriented; the envelope is the shared structured contract layered above those transports.
|
||||
|
||||
## Dependency handoff
|
||||
|
||||
DGR-022 and later shard-control stories can reuse the envelope contract and its `TensorPayload` conversion hooks as the stable activation metadata layer. Future work that requires generated protobuf code can replace the JSON serialization with a generated wire codec without changing the top-level field contract defined here.
|
||||
46
.scratch/distributed-gguf-runtime/evidence/DGR-022/README.md
Normal file
46
.scratch/distributed-gguf-runtime/evidence/DGR-022/README.md
Normal file
@@ -0,0 +1,46 @@
|
||||
# DGR-022 evidence — Shard lifecycle and structured status RPC contract
|
||||
|
||||
**Completed:** 2026-07-17
|
||||
**Branch:** `ralph/distributed-gguf-runtime`
|
||||
**Authority:** `.scratch/distributed-gguf-runtime/prd.json`
|
||||
|
||||
## Outcome
|
||||
|
||||
Implemented the versioned, backend-neutral lifecycle/status contract consumed by a future generated gRPC binding. The contract keeps Meshnet routing, identity, authentication policy, billing, and llama.cpp ownership outside the worker contract.
|
||||
|
||||
## Implemented
|
||||
|
||||
- `packages/node/meshnet_node/shard_lifecycle.py`
|
||||
- capability, health, session, cancellation, release, and metrics RPC names
|
||||
- schema version negotiation and fail-closed unsupported-version handling
|
||||
- structured status/error taxonomy with retryability and details
|
||||
- lifecycle state machine for prefill/decode/cancel/release transitions
|
||||
- monotonic idempotency-step enforcement and duplicate rejection
|
||||
- bounded frame/byte flow control with cancellation-aware waits
|
||||
- explicit cache expectation/result types
|
||||
- deadline policy and TLS/auth transport hooks
|
||||
- deterministic contract serialization round-trip
|
||||
- `tests/test_shard_lifecycle.py`
|
||||
- contract round-trip and RPC coverage
|
||||
- unsupported-version rejection
|
||||
- malformed transition and idempotency rejection
|
||||
- cancellation/release behavior
|
||||
- bounded flow-control behavior
|
||||
- TLS hook and incomplete-contract fail-closed behavior
|
||||
|
||||
## Verification
|
||||
|
||||
```text
|
||||
$ PYTHONPATH=packages/node pytest -q tests/test_shard_lifecycle.py tests/test_activation_envelope.py
|
||||
17 passed in 0.10s
|
||||
```
|
||||
|
||||
The existing DGR-021 activation-envelope tests remain green alongside DGR-022.
|
||||
|
||||
## Scope limitation
|
||||
|
||||
This story defines the lifecycle/status contract only. Generated Python/C++ protobuf bindings and the concrete `shard_runtime.proto` generation pipeline are DGR-023 and remain separate.
|
||||
|
||||
## Dependency handoff
|
||||
|
||||
DGR-023 may consume the RPC names, status taxonomy, version identity, deadlines, flow-control limits, and TLS/auth hooks when the canonical `.proto` schema and toolchain are provisioned.
|
||||
@@ -28,20 +28,20 @@
|
||||
"DGR-021": {
|
||||
"number": 5,
|
||||
"url": "https://git.d-popov.com/popov/neuron-tai/issues/5",
|
||||
"state": "open",
|
||||
"status": "ready"
|
||||
"state": "closed",
|
||||
"status": "completed"
|
||||
},
|
||||
"DGR-022": {
|
||||
"number": 6,
|
||||
"url": "https://git.d-popov.com/popov/neuron-tai/issues/6",
|
||||
"state": "open",
|
||||
"status": "blocked"
|
||||
"status": "ready"
|
||||
},
|
||||
"DGR-023": {
|
||||
"number": 7,
|
||||
"url": "https://git.d-popov.com/popov/neuron-tai/issues/7",
|
||||
"state": "open",
|
||||
"status": "blocked"
|
||||
"status": "ready"
|
||||
},
|
||||
"DGR-024": {
|
||||
"number": 8,
|
||||
@@ -53,7 +53,7 @@
|
||||
"number": 9,
|
||||
"url": "https://git.d-popov.com/popov/neuron-tai/issues/9",
|
||||
"state": "open",
|
||||
"status": "blocked"
|
||||
"status": "ready"
|
||||
},
|
||||
"DGR-026": {
|
||||
"number": 10,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
<!-- GENERATED FROM prd.json — DO NOT EDIT AS AN INDEPENDENT SOURCE. prd.json IS AUTHORITATIVE. -->
|
||||
# DGR-021: Define the versioned named-tensor stream envelope
|
||||
|
||||
- **Status / triage:** specification only; `ready-for-agent`; `passes: false`
|
||||
- **Status / triage:** completed; `passes: true`
|
||||
- **Execution mode:** `AFK`
|
||||
- **Milestone:** `M1`
|
||||
- **Dependencies:** `DGR-018`
|
||||
@@ -18,11 +18,11 @@ Fresh Ralph session: read `.scratch/distributed-gguf-runtime/RALPH-CONTEXT.md`,
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- [ ] Define schema version, request/work ID, route session/epoch, shard range/effective start, phase, position, and idempotency step.
|
||||
- [ ] Define named tensors with shape, dtype, byte order, bounded fragments, compression identity, and checksum.
|
||||
- [ ] Reserve extensible fields for token-ID sidebands, architecture state, recurrent state, and MTP without claiming implementations.
|
||||
- [ ] Add deterministic serialization, fragmentation, checksum, unknown-field, and size-limit tests.
|
||||
- [ ] Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff.
|
||||
- [x] Define schema version, request/work ID, route session/epoch, shard range/effective start, phase, position, and idempotency step.
|
||||
- [x] Define named tensors with shape, dtype, byte order, bounded fragments, compression identity, and checksum.
|
||||
- [x] Reserve extensible fields for token-ID sidebands, architecture state, recurrent state, and MTP without claiming implementations.
|
||||
- [x] Add deterministic serialization, fragmentation, checksum, unknown-field, and size-limit tests.
|
||||
- [x] Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff.
|
||||
|
||||
## Shared quality gates
|
||||
|
||||
@@ -36,4 +36,4 @@ Fresh Ralph session: read `.scratch/distributed-gguf-runtime/RALPH-CONTEXT.md`,
|
||||
|
||||
## Evidence handoff
|
||||
|
||||
Write and verify `.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md`. Until every criterion and applicable gate has real evidence, this story remains `passes: false`. Legacy evidence is provenance only, not completion credit.
|
||||
Verified evidence: `.scratch/distributed-gguf-runtime/evidence/DGR-021/README.md`. Legacy evidence remains provenance only and grants no implementation completion credit.
|
||||
|
||||
@@ -441,7 +441,8 @@
|
||||
"Add deterministic serialization, fragmentation, checksum, unknown-field, and size-limit tests.",
|
||||
"Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff."
|
||||
],
|
||||
"passes": false,
|
||||
"passes": true,
|
||||
"completionNotes": "Added a versioned activation envelope with deterministic JSON serialization, bounded tensor fragmentation, checksum validation, unknown-field preservation, and TensorPayload conversion hooks; verified by targeted pytest and compileall runs.",
|
||||
"notes": "Generated source issue: .scratch/distributed-gguf-runtime/issues/021-define-the-versioned-named-tensor-stream-envelope.md; prd.json is authoritative.",
|
||||
"blocks": [
|
||||
"DGR-022",
|
||||
@@ -482,8 +483,8 @@
|
||||
"Add compatibility tests for supported versions and fail-closed tests for unsupported versions and malformed lifecycle transitions.",
|
||||
"Applicable shared quality gates in `prd.json` pass, and the evidence handoff records exact commands/results, changed files, limitations, and dependency handoff."
|
||||
],
|
||||
"passes": false,
|
||||
"notes": "Generated source issue: .scratch/distributed-gguf-runtime/issues/022-define-shard-lifecycle-and-structured-status-rpcs.md; prd.json is authoritative.",
|
||||
"passes": true,
|
||||
"notes": "Completed from isolated DGR-022 Ralph worktree; verified with 17 focused pytest cases and retained DGR-021 envelope compatibility.",
|
||||
"blocks": [
|
||||
"DGR-024",
|
||||
"DGR-033",
|
||||
|
||||
15
docs/archived/README.md
Normal file
15
docs/archived/README.md
Normal file
@@ -0,0 +1,15 @@
|
||||
# Archived task programs
|
||||
|
||||
These task programs are historical, completed, superseded, or explicitly not part of the active development queue. They are preserved for provenance and must not be treated as runnable work.
|
||||
|
||||
## Archived programs
|
||||
|
||||
- `alpha-hardening/` — historical alpha hardening and settlement/authentication task set.
|
||||
- `dashboard-test-runner/` — historical dashboard test-runner work.
|
||||
- `distributed-inference-performance/` — superseded distributed-inference performance backlog.
|
||||
- `node-capability-admission/` — historical capability-admission task set.
|
||||
- `proxy-stream-cancellation/` — historical proxy cancellation task.
|
||||
- `qwen3.6-27b-demand-placement/` — superseded Qwen demand-placement planning.
|
||||
- `routing-compatibility-regression/` — historical routing compatibility regression work.
|
||||
|
||||
The active distributed GGUF backlog remains under `.scratch/distributed-gguf-runtime/`. Architectural decisions under `docs/adr/` remain active and were not moved.
|
||||
@@ -86,6 +86,59 @@ class TensorPayload:
|
||||
# Number of tokens already cached before this payload's tokens (decode steps).
|
||||
past_len: int | None = None
|
||||
|
||||
def to_envelope(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
request_id: str,
|
||||
work_id: str,
|
||||
route_session: str,
|
||||
route_epoch: int,
|
||||
shard_start: int,
|
||||
effective_start: int,
|
||||
phase: str,
|
||||
position: int,
|
||||
idempotency_step: int,
|
||||
byte_order: str = "little",
|
||||
compression: str = "identity",
|
||||
max_fragment_bytes: int | None = None,
|
||||
token_id_sideband: list[int] | None = None,
|
||||
architecture_state: dict[str, Any] | None = None,
|
||||
recurrent_state: dict[str, Any] | None = None,
|
||||
mtp: dict[str, Any] | None = None,
|
||||
extensions: dict[str, Any] | None = None,
|
||||
):
|
||||
from .protocol import ActivationEnvelope, DEFAULT_FRAGMENT_BYTES
|
||||
|
||||
return ActivationEnvelope.from_tensor_payload(
|
||||
payload=self,
|
||||
name=name,
|
||||
request_id=request_id,
|
||||
work_id=work_id,
|
||||
route_session=route_session,
|
||||
route_epoch=route_epoch,
|
||||
shard_start=shard_start,
|
||||
effective_start=effective_start,
|
||||
phase=phase,
|
||||
position=position,
|
||||
idempotency_step=idempotency_step,
|
||||
byte_order=byte_order,
|
||||
compression=compression,
|
||||
max_fragment_bytes=max_fragment_bytes or DEFAULT_FRAGMENT_BYTES,
|
||||
token_id_sideband=token_id_sideband,
|
||||
architecture_state=architecture_state,
|
||||
recurrent_state=recurrent_state,
|
||||
mtp=mtp,
|
||||
extensions=extensions,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_envelope(cls, envelope, *, tensor_name: str = "activations"):
|
||||
tensor_payload = envelope.to_tensor_payload(tensor_name=tensor_name)
|
||||
if not isinstance(tensor_payload, cls):
|
||||
raise TypeError("envelope did not produce a TensorPayload")
|
||||
return tensor_payload
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TailTokenResult:
|
||||
|
||||
376
packages/node/meshnet_node/protocol.py
Normal file
376
packages/node/meshnet_node/protocol.py
Normal file
@@ -0,0 +1,376 @@
|
||||
"""Versioned activation-stream envelope for shard hops.
|
||||
|
||||
The transport still moves raw bytes over HTTP/WebSocket, but the payload now has
|
||||
a stable, extensible envelope that names each tensor, preserves unknown fields,
|
||||
and can round-trip deterministically across direct and relayed hops.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
SCHEMA_NAME = "meshnet.activation-stream"
|
||||
SCHEMA_VERSION = 1
|
||||
DEFAULT_FRAGMENT_BYTES = 64 * 1024
|
||||
|
||||
|
||||
def _canonical_json(data: Any) -> bytes:
|
||||
return json.dumps(data, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
|
||||
|
||||
def _sha256_hex(data: bytes) -> str:
|
||||
return hashlib.sha256(data).hexdigest()
|
||||
|
||||
|
||||
def _normalize_shape(shape: list[int] | tuple[int, ...]) -> list[int]:
|
||||
normalized = [int(dim) for dim in shape]
|
||||
if not normalized or any(dim <= 0 for dim in normalized):
|
||||
raise ValueError("shape must be a non-empty list of positive integers")
|
||||
return normalized
|
||||
|
||||
|
||||
def _fragment_bytes(body: bytes, max_fragment_bytes: int) -> tuple[bytes, ...]:
|
||||
if max_fragment_bytes <= 0:
|
||||
raise ValueError("max_fragment_bytes must be positive")
|
||||
if not body:
|
||||
return (b"",)
|
||||
return tuple(body[offset : offset + max_fragment_bytes] for offset in range(0, len(body), max_fragment_bytes))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorFragment:
|
||||
"""One bounded chunk of a named tensor."""
|
||||
|
||||
offset: int
|
||||
body: bytes
|
||||
checksum: str
|
||||
compression: str = "identity"
|
||||
extensions: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_bytes(
|
||||
cls,
|
||||
body: bytes,
|
||||
*,
|
||||
offset: int,
|
||||
compression: str = "identity",
|
||||
extensions: dict[str, Any] | None = None,
|
||||
) -> "TensorFragment":
|
||||
return cls(
|
||||
offset=int(offset),
|
||||
body=bytes(body),
|
||||
checksum=_sha256_hex(body),
|
||||
compression=compression,
|
||||
extensions=dict(extensions or {}),
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
data = {
|
||||
"offset": self.offset,
|
||||
"compression": self.compression,
|
||||
"checksum": self.checksum,
|
||||
"body_base64": base64.b64encode(self.body).decode("ascii"),
|
||||
}
|
||||
data.update(self.extensions)
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "TensorFragment":
|
||||
known = {"offset", "compression", "checksum", "body_base64"}
|
||||
body = base64.b64decode(data.get("body_base64", ""))
|
||||
fragment = cls(
|
||||
offset=int(data["offset"]),
|
||||
body=body,
|
||||
checksum=str(data.get("checksum") or _sha256_hex(body)),
|
||||
compression=str(data.get("compression") or "identity"),
|
||||
extensions={k: v for k, v in data.items() if k not in known},
|
||||
)
|
||||
fragment.validate()
|
||||
return fragment
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.checksum != _sha256_hex(self.body):
|
||||
raise ValueError("fragment checksum mismatch")
|
||||
if self.offset < 0:
|
||||
raise ValueError("fragment offset must be non-negative")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NamedTensor:
|
||||
"""A tensor named within a versioned activation envelope."""
|
||||
|
||||
name: str
|
||||
shape: list[int]
|
||||
dtype: str
|
||||
byte_order: str
|
||||
checksum: str
|
||||
fragments: tuple[TensorFragment, ...]
|
||||
compression: str = "identity"
|
||||
extensions: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_bytes(
|
||||
cls,
|
||||
*,
|
||||
name: str,
|
||||
body: bytes,
|
||||
shape: list[int] | tuple[int, ...],
|
||||
dtype: str,
|
||||
byte_order: str = "little",
|
||||
compression: str = "identity",
|
||||
max_fragment_bytes: int = DEFAULT_FRAGMENT_BYTES,
|
||||
extensions: dict[str, Any] | None = None,
|
||||
) -> "NamedTensor":
|
||||
normalized_shape = _normalize_shape(shape)
|
||||
fragments = tuple(
|
||||
TensorFragment.from_bytes(fragment, offset=offset, compression=compression)
|
||||
for offset, fragment in enumerate(_fragment_bytes(body, max_fragment_bytes))
|
||||
for offset in (offset * max_fragment_bytes,)
|
||||
)
|
||||
return cls(
|
||||
name=str(name),
|
||||
shape=normalized_shape,
|
||||
dtype=str(dtype),
|
||||
byte_order=str(byte_order),
|
||||
checksum=_sha256_hex(body),
|
||||
fragments=fragments,
|
||||
compression=compression,
|
||||
extensions=dict(extensions or {}),
|
||||
)
|
||||
|
||||
def body(self) -> bytes:
|
||||
ordered = sorted(self.fragments, key=lambda frag: frag.offset)
|
||||
body = b"".join(fragment.body for fragment in ordered)
|
||||
if _sha256_hex(body) != self.checksum:
|
||||
raise ValueError(f"tensor {self.name!r} checksum mismatch")
|
||||
return body
|
||||
|
||||
def validate(self) -> None:
|
||||
for fragment in self.fragments:
|
||||
fragment.validate()
|
||||
self.body()
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
data = {
|
||||
"name": self.name,
|
||||
"shape": list(self.shape),
|
||||
"dtype": self.dtype,
|
||||
"byte_order": self.byte_order,
|
||||
"compression": self.compression,
|
||||
"checksum": self.checksum,
|
||||
"fragments": [fragment.to_dict() for fragment in self.fragments],
|
||||
}
|
||||
data.update(self.extensions)
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "NamedTensor":
|
||||
known = {
|
||||
"name",
|
||||
"shape",
|
||||
"dtype",
|
||||
"byte_order",
|
||||
"compression",
|
||||
"checksum",
|
||||
"fragments",
|
||||
}
|
||||
tensor = cls(
|
||||
name=str(data["name"]),
|
||||
shape=_normalize_shape(list(data["shape"])),
|
||||
dtype=str(data["dtype"]),
|
||||
byte_order=str(data.get("byte_order", "little")),
|
||||
compression=str(data.get("compression", "identity")),
|
||||
checksum=str(data["checksum"]),
|
||||
fragments=tuple(TensorFragment.from_dict(fragment) for fragment in data.get("fragments", [])),
|
||||
extensions={k: v for k, v in data.items() if k not in known},
|
||||
)
|
||||
tensor.validate()
|
||||
return tensor
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ActivationEnvelope:
|
||||
"""Versioned envelope for shard activation traffic."""
|
||||
|
||||
request_id: str
|
||||
work_id: str
|
||||
route_session: str
|
||||
route_epoch: int
|
||||
shard_start: int
|
||||
effective_start: int
|
||||
phase: str
|
||||
position: int
|
||||
idempotency_step: int
|
||||
tensors: tuple[NamedTensor, ...]
|
||||
version: int = SCHEMA_VERSION
|
||||
schema: str = SCHEMA_NAME
|
||||
token_id_sideband: list[int] | None = None
|
||||
architecture_state: dict[str, Any] | None = None
|
||||
recurrent_state: dict[str, Any] | None = None
|
||||
mtp: dict[str, Any] | None = None
|
||||
extensions: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_tensor_payload(
|
||||
cls,
|
||||
*,
|
||||
payload: Any,
|
||||
name: str,
|
||||
request_id: str,
|
||||
work_id: str,
|
||||
route_session: str,
|
||||
route_epoch: int,
|
||||
shard_start: int,
|
||||
effective_start: int,
|
||||
phase: str,
|
||||
position: int,
|
||||
idempotency_step: int,
|
||||
byte_order: str = "little",
|
||||
compression: str = "identity",
|
||||
max_fragment_bytes: int = DEFAULT_FRAGMENT_BYTES,
|
||||
token_id_sideband: list[int] | None = None,
|
||||
architecture_state: dict[str, Any] | None = None,
|
||||
recurrent_state: dict[str, Any] | None = None,
|
||||
mtp: dict[str, Any] | None = None,
|
||||
extensions: dict[str, Any] | None = None,
|
||||
) -> "ActivationEnvelope":
|
||||
tensor = NamedTensor.from_bytes(
|
||||
name=name,
|
||||
body=payload.body,
|
||||
shape=payload.shape,
|
||||
dtype="bfloat16",
|
||||
byte_order=byte_order,
|
||||
compression=compression,
|
||||
max_fragment_bytes=max_fragment_bytes,
|
||||
extensions={
|
||||
"attention_mask_header": payload.attention_mask_header,
|
||||
"position_ids_header": payload.position_ids_header,
|
||||
**({"past_len": payload.past_len} if payload.past_len is not None else {}),
|
||||
},
|
||||
)
|
||||
return cls(
|
||||
request_id=request_id,
|
||||
work_id=work_id,
|
||||
route_session=route_session,
|
||||
route_epoch=int(route_epoch),
|
||||
shard_start=int(shard_start),
|
||||
effective_start=int(effective_start),
|
||||
phase=str(phase),
|
||||
position=int(position),
|
||||
idempotency_step=int(idempotency_step),
|
||||
tensors=(tensor,),
|
||||
token_id_sideband=list(token_id_sideband) if token_id_sideband is not None else None,
|
||||
architecture_state=architecture_state,
|
||||
recurrent_state=recurrent_state,
|
||||
mtp=mtp,
|
||||
extensions=dict(extensions or {}),
|
||||
)
|
||||
|
||||
def to_tensor_payload(self, *, tensor_name: str = "activations") -> Any:
|
||||
from .model_backend import TensorPayload
|
||||
|
||||
tensor = self.tensor(tensor_name)
|
||||
return TensorPayload(
|
||||
body=tensor.body(),
|
||||
shape=list(tensor.shape),
|
||||
attention_mask_header=tensor.extensions.get("attention_mask_header"),
|
||||
position_ids_header=tensor.extensions.get("position_ids_header"),
|
||||
past_len=tensor.extensions.get("past_len"),
|
||||
)
|
||||
|
||||
def tensor(self, name: str = "activations") -> NamedTensor:
|
||||
for tensor in self.tensors:
|
||||
if tensor.name == name:
|
||||
return tensor
|
||||
raise KeyError(name)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
data = {
|
||||
"schema": self.schema,
|
||||
"version": self.version,
|
||||
"request_id": self.request_id,
|
||||
"work_id": self.work_id,
|
||||
"route_session": self.route_session,
|
||||
"route_epoch": self.route_epoch,
|
||||
"shard_start": self.shard_start,
|
||||
"effective_start": self.effective_start,
|
||||
"phase": self.phase,
|
||||
"position": self.position,
|
||||
"idempotency_step": self.idempotency_step,
|
||||
"tensors": [tensor.to_dict() for tensor in self.tensors],
|
||||
}
|
||||
if self.token_id_sideband is not None:
|
||||
data["token_id_sideband"] = list(self.token_id_sideband)
|
||||
if self.architecture_state is not None:
|
||||
data["architecture_state"] = self.architecture_state
|
||||
if self.recurrent_state is not None:
|
||||
data["recurrent_state"] = self.recurrent_state
|
||||
if self.mtp is not None:
|
||||
data["mtp"] = self.mtp
|
||||
data.update(self.extensions)
|
||||
return data
|
||||
|
||||
def to_bytes(self, *, max_bytes: int | None = None) -> bytes:
|
||||
raw = _canonical_json(self.to_dict())
|
||||
if max_bytes is not None and len(raw) > max_bytes:
|
||||
raise ValueError("activation envelope exceeds the size limit")
|
||||
return raw
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, data: bytes) -> "ActivationEnvelope":
|
||||
payload = json.loads(data)
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("activation envelope must be a JSON object")
|
||||
known = {
|
||||
"schema",
|
||||
"version",
|
||||
"request_id",
|
||||
"work_id",
|
||||
"route_session",
|
||||
"route_epoch",
|
||||
"shard_start",
|
||||
"effective_start",
|
||||
"phase",
|
||||
"position",
|
||||
"idempotency_step",
|
||||
"tensors",
|
||||
"token_id_sideband",
|
||||
"architecture_state",
|
||||
"recurrent_state",
|
||||
"mtp",
|
||||
}
|
||||
envelope = cls(
|
||||
schema=str(payload.get("schema", SCHEMA_NAME)),
|
||||
version=int(payload.get("version", SCHEMA_VERSION)),
|
||||
request_id=str(payload["request_id"]),
|
||||
work_id=str(payload["work_id"]),
|
||||
route_session=str(payload["route_session"]),
|
||||
route_epoch=int(payload["route_epoch"]),
|
||||
shard_start=int(payload["shard_start"]),
|
||||
effective_start=int(payload["effective_start"]),
|
||||
phase=str(payload["phase"]),
|
||||
position=int(payload["position"]),
|
||||
idempotency_step=int(payload["idempotency_step"]),
|
||||
tensors=tuple(NamedTensor.from_dict(item) for item in payload.get("tensors", [])),
|
||||
token_id_sideband=payload.get("token_id_sideband"),
|
||||
architecture_state=payload.get("architecture_state"),
|
||||
recurrent_state=payload.get("recurrent_state"),
|
||||
mtp=payload.get("mtp"),
|
||||
extensions={k: v for k, v in payload.items() if k not in known},
|
||||
)
|
||||
envelope.validate()
|
||||
return envelope
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.version != SCHEMA_VERSION:
|
||||
raise ValueError("unsupported activation envelope version")
|
||||
if self.schema != SCHEMA_NAME:
|
||||
raise ValueError("unsupported activation envelope schema")
|
||||
if self.phase not in {"prefill", "decode"}:
|
||||
raise ValueError("phase must be prefill or decode")
|
||||
for tensor in self.tensors:
|
||||
tensor.validate()
|
||||
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
|
||||
130
tests/test_activation_envelope.py
Normal file
130
tests/test_activation_envelope.py
Normal file
@@ -0,0 +1,130 @@
|
||||
"""DGR-021: versioned activation envelope for shard traffic."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_node.model_backend import TensorPayload
|
||||
from meshnet_node.protocol import ActivationEnvelope, NamedTensor, TensorFragment
|
||||
|
||||
|
||||
def _payload(body: bytes = b"\x01\x02\x03\x04", *, shape=(1, 2, 1)) -> TensorPayload:
|
||||
return TensorPayload(
|
||||
body=body,
|
||||
shape=list(shape),
|
||||
attention_mask_header="1,2:AAAA",
|
||||
position_ids_header="1,2:BBBB",
|
||||
past_len=7,
|
||||
)
|
||||
|
||||
|
||||
def test_envelope_roundtrips_deterministically():
|
||||
payload = _payload()
|
||||
env_a = payload.to_envelope(
|
||||
name="activations",
|
||||
request_id="req-1",
|
||||
work_id="work-1",
|
||||
route_session="route-1",
|
||||
route_epoch=9,
|
||||
shard_start=12,
|
||||
effective_start=18,
|
||||
phase="prefill",
|
||||
position=0,
|
||||
idempotency_step=0,
|
||||
extensions={"zeta": 3, "alpha": "x"},
|
||||
)
|
||||
env_b = ActivationEnvelope.from_bytes(env_a.to_bytes())
|
||||
|
||||
assert env_a.to_bytes() == env_b.to_bytes()
|
||||
assert env_b.request_id == "req-1"
|
||||
assert env_b.work_id == "work-1"
|
||||
assert env_b.route_session == "route-1"
|
||||
assert env_b.shard_start == 12
|
||||
assert env_b.effective_start == 18
|
||||
assert env_b.phase == "prefill"
|
||||
assert env_b.tensors[0].body() == payload.body
|
||||
assert env_b.tensors[0].extensions["past_len"] == 7
|
||||
assert env_b.extensions == {"alpha": "x", "zeta": 3}
|
||||
|
||||
|
||||
def test_fragmentation_and_checksums_are_bounded():
|
||||
body = b"0123456789abcdef"
|
||||
tensor = NamedTensor.from_bytes(
|
||||
name="activations",
|
||||
body=body,
|
||||
shape=[1, 8, 1],
|
||||
dtype="bfloat16",
|
||||
max_fragment_bytes=5,
|
||||
)
|
||||
|
||||
assert len(tensor.fragments) == 4
|
||||
assert all(isinstance(fragment, TensorFragment) for fragment in tensor.fragments)
|
||||
assert all(len(fragment.body) <= 5 for fragment in tensor.fragments)
|
||||
assert tensor.body() == body
|
||||
assert tensor.checksum == hashlib.sha256(body).hexdigest()
|
||||
|
||||
|
||||
def test_unknown_fields_are_preserved_on_roundtrip():
|
||||
env = _payload().to_envelope(
|
||||
name="activations",
|
||||
request_id="req-2",
|
||||
work_id="work-2",
|
||||
route_session="route-2",
|
||||
route_epoch=1,
|
||||
shard_start=0,
|
||||
effective_start=0,
|
||||
phase="decode",
|
||||
position=3,
|
||||
idempotency_step=2,
|
||||
)
|
||||
raw = json.loads(env.to_bytes())
|
||||
raw["future_top_level"] = {"x": 1}
|
||||
raw["tensors"][0]["future_tensor_field"] = "ok"
|
||||
roundtrip = ActivationEnvelope.from_bytes(json.dumps(raw).encode())
|
||||
|
||||
assert roundtrip.extensions["future_top_level"] == {"x": 1}
|
||||
assert roundtrip.tensors[0].extensions["future_tensor_field"] == "ok"
|
||||
|
||||
|
||||
def test_size_limit_rejects_large_envelopes():
|
||||
env = _payload(body=b"x" * 64).to_envelope(
|
||||
name="activations",
|
||||
request_id="req-3",
|
||||
work_id="work-3",
|
||||
route_session="route-3",
|
||||
route_epoch=1,
|
||||
shard_start=0,
|
||||
effective_start=0,
|
||||
phase="decode",
|
||||
position=5,
|
||||
idempotency_step=4,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="size limit"):
|
||||
env.to_bytes(max_bytes=32)
|
||||
|
||||
|
||||
def test_tensor_payload_roundtrips_through_envelope():
|
||||
payload = _payload(body=b"abcde", shape=(1, 5, 1))
|
||||
env = payload.to_envelope(
|
||||
name="activations",
|
||||
request_id="req-4",
|
||||
work_id="work-4",
|
||||
route_session="route-4",
|
||||
route_epoch=2,
|
||||
shard_start=6,
|
||||
effective_start=11,
|
||||
phase="decode",
|
||||
position=8,
|
||||
idempotency_step=9,
|
||||
)
|
||||
|
||||
restored = TensorPayload.from_envelope(env)
|
||||
assert restored.body == payload.body
|
||||
assert restored.shape == payload.shape
|
||||
assert restored.attention_mask_header == payload.attention_mask_header
|
||||
assert restored.position_ids_header == payload.position_ids_header
|
||||
assert restored.past_len == payload.past_len
|
||||
182
tests/test_shard_lifecycle.py
Normal file
182
tests/test_shard_lifecycle.py
Normal file
@@ -0,0 +1,182 @@
|
||||
"""DGR-022 lifecycle and structured-status contract tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_node.shard_lifecycle import (
|
||||
CacheExpectation,
|
||||
CacheResult,
|
||||
CapabilityResponse,
|
||||
CancellationToken,
|
||||
FlowControl,
|
||||
FlowControlLimits,
|
||||
LifecycleContractError,
|
||||
LifecycleState,
|
||||
RpcName,
|
||||
SessionFrame,
|
||||
SessionLifecycle,
|
||||
SessionPhase,
|
||||
SessionRequest,
|
||||
ShardRpcContract,
|
||||
StatusCode,
|
||||
StructuredStatus,
|
||||
)
|
||||
|
||||
|
||||
def _request(**overrides) -> SessionRequest:
|
||||
values = dict(
|
||||
schema_version=1,
|
||||
request_id="request-1",
|
||||
work_id="work-1",
|
||||
route_session="route-1",
|
||||
route_epoch=3,
|
||||
artifact_fingerprint="artifact-sha",
|
||||
runtime_fingerprint="runtime-sha",
|
||||
shard_start=0,
|
||||
shard_end=8,
|
||||
effective_start=0,
|
||||
cache_expectation=CacheExpectation.OPTIONAL,
|
||||
)
|
||||
values.update(overrides)
|
||||
return SessionRequest(**values)
|
||||
|
||||
|
||||
def _frame(phase: SessionPhase, step: int, position: int = 0) -> SessionFrame:
|
||||
return SessionFrame(
|
||||
phase=phase,
|
||||
position=position,
|
||||
idempotency_step=step,
|
||||
payload=b"frame",
|
||||
size_bytes=5,
|
||||
)
|
||||
|
||||
|
||||
def test_contract_roundtrip_exposes_all_lifecycle_rpcs_and_operational_hooks():
|
||||
contract = ShardRpcContract()
|
||||
restored = ShardRpcContract.from_bytes(contract.to_bytes())
|
||||
|
||||
assert restored == contract
|
||||
assert restored.methods == tuple(RpcName)
|
||||
assert restored.deadlines.session_idle_seconds > restored.deadlines.health_seconds
|
||||
assert restored.flow_control.max_inflight_frames == 32
|
||||
assert restored.tls_auth.tls_required is True
|
||||
|
||||
|
||||
def test_supported_version_roundtrip_and_status_are_stable():
|
||||
status = StructuredStatus(StatusCode.RESOURCE_EXHAUSTED, "window full", retryable=True, details={"limit": "32"})
|
||||
restored = StructuredStatus.from_dict(status.to_dict())
|
||||
|
||||
assert restored == status
|
||||
assert CapabilityResponse(
|
||||
status=StructuredStatus(StatusCode.OK, "ready"),
|
||||
schema_version=1,
|
||||
supported_versions=(1,),
|
||||
).validate() is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", [0, 2, 99])
|
||||
def test_unsupported_versions_fail_closed(version):
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
ShardRpcContract.from_bytes(json.dumps({
|
||||
"schema": "meshnet.shard-runtime",
|
||||
"version": version,
|
||||
"methods": [method.value for method in RpcName],
|
||||
}).encode())
|
||||
assert exc.value.status.code is StatusCode.UNSUPPORTED_VERSION
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
_request(schema_version=version).validate()
|
||||
assert exc.value.status.code is StatusCode.UNSUPPORTED_VERSION
|
||||
|
||||
|
||||
def test_malformed_lifecycle_transitions_fail_closed():
|
||||
session = SessionLifecycle(_request())
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
session.apply(_frame(SessionPhase.DECODE, 0))
|
||||
assert exc.value.status.code is StatusCode.MALFORMED_LIFECYCLE
|
||||
assert session.state is LifecycleState.OPEN
|
||||
|
||||
session.apply(_frame(SessionPhase.PREFILL, 0))
|
||||
session.apply(_frame(SessionPhase.DECODE, 1, position=1))
|
||||
assert session.state is LifecycleState.DECODING
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
session.apply(_frame(SessionPhase.PREFILL, 2))
|
||||
assert exc.value.status.code is StatusCode.MALFORMED_LIFECYCLE
|
||||
|
||||
|
||||
def test_duplicate_and_non_monotonic_idempotency_steps_are_rejected():
|
||||
session = SessionLifecycle(_request())
|
||||
session.apply(_frame(SessionPhase.PREFILL, 4))
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
session.apply(_frame(SessionPhase.PREFILL, 4))
|
||||
assert exc.value.status.code is StatusCode.ALREADY_EXISTS
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
session.apply(_frame(SessionPhase.PREFILL, 3))
|
||||
assert exc.value.status.code is StatusCode.MALFORMED_LIFECYCLE
|
||||
|
||||
|
||||
def test_cancel_propagates_to_waiters_and_release_is_idempotent():
|
||||
session = SessionLifecycle(_request())
|
||||
assert session.cancel("client disconnected").code is StatusCode.CANCELLED
|
||||
assert session.cancellation.cancelled is True
|
||||
assert session.state is LifecycleState.CANCELLED
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
session.apply(_frame(SessionPhase.DECODE, 0))
|
||||
assert exc.value.status.code is StatusCode.CANCELLED
|
||||
|
||||
assert session.release().code is StatusCode.OK
|
||||
assert session.state is LifecycleState.RELEASED
|
||||
assert session.release().code is StatusCode.OK
|
||||
|
||||
|
||||
def test_bounded_flow_control_rejects_oversize_and_full_windows():
|
||||
flow = FlowControl(FlowControlLimits(max_inflight_frames=1, max_inflight_bytes=10, max_frame_bytes=8))
|
||||
flow.acquire(8)
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
flow.acquire(1)
|
||||
assert exc.value.status.code is StatusCode.RESOURCE_EXHAUSTED
|
||||
assert exc.value.status.retryable is True
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
flow.acquire(9)
|
||||
assert exc.value.status.code is StatusCode.RESOURCE_EXHAUSTED
|
||||
|
||||
flow.release(8)
|
||||
assert flow.outstanding == (0, 0)
|
||||
|
||||
|
||||
def test_flow_control_wait_honours_cancellation():
|
||||
token = CancellationToken()
|
||||
flow = FlowControl(FlowControlLimits(max_inflight_frames=1, max_inflight_bytes=10, max_frame_bytes=8))
|
||||
flow.acquire(4)
|
||||
token.cancel()
|
||||
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
flow.acquire(1, wait=True, cancelled=lambda: token.cancelled)
|
||||
assert exc.value.status.code is StatusCode.CANCELLED
|
||||
|
||||
|
||||
def test_cache_expectation_and_result_are_explicit():
|
||||
request = _request(cache_expectation=CacheExpectation.REQUIRED)
|
||||
assert request.cache_expectation is CacheExpectation.REQUIRED
|
||||
assert CacheResult.MISS.value == "MISS"
|
||||
|
||||
|
||||
def test_invalid_contract_methods_and_tls_hooks_fail_closed():
|
||||
with pytest.raises(LifecycleContractError) as exc:
|
||||
ShardRpcContract(methods=(RpcName.HEALTH,)).validate()
|
||||
assert exc.value.status.code is StatusCode.FAILED_PRECONDITION
|
||||
|
||||
with pytest.raises(ValueError, match="server_name"):
|
||||
ShardRpcContract(
|
||||
tls_auth=type(ShardRpcContract().tls_auth)(tls_required=True, server_name="")
|
||||
).validate()
|
||||
Reference in New Issue
Block a user