diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-012/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-012/README.md new file mode 100644 index 0000000..343c995 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-012/README.md @@ -0,0 +1,220 @@ +# DGR-012 — Continuous batching and bounded admission: evidence + +Status: done +Date: 2026-07-16 +Evidence kind: **synthetic-unit** (pure-numpy KV-cached dense-Llama reference + +node-local continuous-batching scheduler). No model download, no GPU, no torch, +no network, no API credit. + +## Summary + +Implemented the node-local scheduler that turns concurrent Route Sessions into +llama.cpp-style continuous batches while bounding admission (RALPH runtime +decision #9, ADR-0024). It sits **on top of** the DGR-007 Hot KV State manager — +batching is a scheduling concern layered over the existing per-`(session, epoch)` +KV isolation, not a new control plane or a change to the KV contract. + +- **Bounded admission (`NodeBudget` + `submit`).** A new session is admitted only + if it fits four budgets: resident **weight** footprint (reported), **KV** byte + budget (a session must be able to hold its *whole* generation, prompt + new + tokens, on its own), **scratch** (per-active-session activation buffers, capped + by a total scratch envelope), and the bounded **queue**. Anything that cannot + fit is rejected up front with an explicit `AdmissionReason` + (`REJECTED_KV_BUDGET` / `REJECTED_SCRATCH_BUDGET` / `REJECTED_DUPLICATE`); + anything that fits but has no free slot waits in the bounded queue; a **full + queue is refused** (`REJECTED_QUEUE_FULL`) — that refusal is the backpressure + signal. +- **Continuous batching (`ContinuousBatchScheduler` + `KvBatchEngine`).** Every + tick, all currently-decoding sessions contribute their single next token to one + batch (bounded by `max_batch_size`); the engine runs the batch once. Each + session keeps its own position and appends its own sampled token via its own + `SessionCache`, so batching never mixes outputs. `KvBatchEngine` adapts the + DGR-007 `KvBoundaryAdapter`, so the batch runs against the *real* KV isolation + path; the pinned llama.cpp worker (DGR-008) implements the same + `recipe_fingerprint`/`prefill`/`decode_batch`/`release` contract where a batch + becomes one `llama_decode` over several sequences. +- **Prefill does not starve decode.** The scheduling policy is explicit and fixed: + **decode first, then bounded prefill.** In-flight decodes always run before any + new prompt is prefilled, and prefill work per tick is capped + (`max_prefill_tokens_per_tick`, always allowing at least one so a single large + prompt still progresses). A burst of new sessions cannot stall generations + already in flight. +- **Bounded memory / backpressure.** KV growth is bounded by the manager byte + budget; queued activations are bounded by `max_queue_depth` and the scratch + envelope; completed sessions release their KV so total KV returns to zero. +- **Capability telemetry (`SchedulerTelemetry`).** Reports active sessions, queue + depth, batch occupancy (last/avg/max), KV pressure (bytes/budget), scratch + pressure, prefill/decode token totals **and rates**, and rejected admissions + (total + by reason). All JSON-safe. +- **Concurrency 1/2/4/8 sweep (`run_concurrency_sweep`).** Runs the same eight + jobs at each level against a fresh KV manager and proves (a) **no cross-session + corruption** — every level yields byte-identical per-session tokens as the + serialized concurrency-1 reference — and (b) **saturation** — average batch + occupancy rises and total ticks fall as concurrency increases, until occupancy + plateaus. + +No existing runtime code was modified — this story is purely additive (one new +module + one new test module + evidence). + +## Files changed (all new) + +- `packages/node/meshnet_node/batch_scheduler.py` — the scheduler: + - `NodeBudget` — weight/KV/scratch/queue budgets + `max_batch_size` / + `max_prefill_tokens_per_tick` scheduling bounds, with derived + `effective_active_cap` (tighter of active-slot and scratch caps). + - `AdmissionReason` / `AdmissionDecision` — structured admit/queue/reject. + - `GenerationRequest` / `DecodeItem` / `StepResult` — job + engine I/O values. + - `KvBatchEngine` — adapts a full-shard `KvBoundaryAdapter` to the batch-engine + contract (rejects a partial head/tail-only range). + - `SchedulerTelemetry` — the bounded capability snapshot. + - `ContinuousBatchScheduler` — thread-safe `submit` / `run_tick` / + `run_to_completion` / `telemetry`, decode-first-then-bounded-prefill policy. + - `run_concurrency_sweep` / `ConcurrencyResult` / `ConcurrencySweep` — the + deterministic 1/2/4/8 saturation report + corruption check. +- `tests/test_batch_scheduler.py` — 16 tests (see below); reuses the DGR-007 + numpy dense-Llama reference via `from test_hot_kv_state import _KvDenseLlama, + _KvReferenceShard`. +- `.scratch/distributed-gguf-runtime/evidence/DGR-012/` — this README, + `commands.txt`, `generate_evidence.py`, `results.json`. + +## Acceptance criteria → evidence + +- **Scheduler admits sessions against weight, KV, scratch, and queue budgets** — + `test_admission_respects_active_scratch_and_queue_budgets` (fill slots → queue → + reject full queue), `test_admission_rejects_a_session_that_cannot_fit_the_kv_budget`, + `test_admission_rejects_when_per_session_scratch_exceeds_budget`, + `test_duplicate_submission_is_rejected`, + `test_weight_budget_is_reported_in_telemetry`. +- **Compatible decode steps form batches preserving per-session positions/outputs** + — `test_batched_decode_preserves_per_session_positions_and_outputs` + (`batch_occupancy_max == 4`, four divergent references each reproduced), + `test_positions_are_isolated_across_different_prompt_lengths` (prompt lengths 1/3/7). +- **Prefill does not starve decode; policy and bounds explicit** — + `test_prefill_does_not_starve_in_flight_decode` (in-flight session decodes on + *every* tick during a 4-session prefill burst; ≤1 prefill/tick), + `test_decode_first_policy_is_explicit_in_a_single_tick`. +- **Backpressure prevents unbounded queued activations or KV growth** — + `test_backpressure_signals_when_queue_full_then_recovers`, + `test_completed_sessions_release_kv_so_growth_is_bounded` (`kv_total_bytes == 0` + after completion). +- **Capability telemetry reports all required signals** — + `test_telemetry_reports_every_required_signal` (asserts every key present; + deterministic rates under an injected clock). +- **Concurrency 1/2/4/8 identifies saturation, no cross-session corruption** — + `test_concurrency_sweep_identifies_saturation_without_corruption` + (occupancy strictly ↑, ticks strictly ↓, tokens/tick ↑, `corruption_free`, + 0 cache misses, saturation=8), `test_concurrency_sweep_saturates_below_max_when_load_is_small`. +- **Engine/usage guards** — `test_kv_batch_engine_requires_a_full_shard`, + `test_run_to_completion_is_bounded_against_misconfiguration`. + +## Concurrency 1/2/4/8 sweep (real, deterministic — `results.json`) + +Eight sessions, prompt length 4, 8 new tokens each; fresh KV manager per level; +budgets sized so KV never evicts (so the corruption check is unambiguous). + +| concurrency | ticks | avg batch occupancy | max occupancy | tokens/tick | peak KV bytes | +|---|---|---|---|---|---| +| 1 | 64 | 1.000 | 1 | 1.375 | 15360 | +| 2 | 33 | 1.750 | 2 | 2.667 | 29184 | +| 4 | 19 | 3.111 | 4 | 4.632 | 52224 | +| 8 | 15 | 4.000 | 7 | 5.867 | 75264 | + +`saturation_concurrency = 8`, `corruption_free = True`, `cache_misses = 0`, +`rejected_admissions = 0`. As concurrency rises, the scheduler packs more sessions +per decode step (occupancy ↑) and finishes the same 56 decode + 32 prefill tokens +in far fewer ticks (aggregate work/tick ↑) — the batching throughput property — +while every per-session token stream stays byte-identical to the serialized +reference (no cross-session corruption). Max occupancy is 7 (not 8) at level 8 +because the fairness policy prefills at most one new session per tick, so the last +session begins decoding one tick later. + +## Commands and real results + +```bash +VP=/run/media/popov/d/DEV/repos/d-popov.com/AI/.venv/bin/python + +$VP -m pytest -q tests/test_batch_scheduler.py +# -> 16 passed + +$VP -m pytest -q tests/test_hot_kv_state.py # dependency still green +# -> 22 passed + +$VP -m compileall -q packages tests +# -> exit 0 + +git diff --check +# -> exit 0 + +$VP .scratch/distributed-gguf-runtime/evidence/DGR-012/generate_evidence.py +# -> wrote results.json; saturation_concurrency=8 corruption_free=True + +$VP -m pytest -q -rfE -p no:cacheprovider +# -> FULL_SUITE_RESULT_PLACEHOLDER +``` + +`commands.txt` beside this README captures the exact commands. + +## Full-suite baseline (pre-existing unrelated failures) + +FULL_SUITE_BASELINE_PLACEHOLDER + +## Limitations and deferred work + +- **Synthetic-unit, not real weights.** The scheduler is exercised against the + deterministic numpy KV-cached dense-Llama reference (the same one DGR-007 uses), + not a downloaded GGUF. This is required to keep the default gate deterministic, + download-free, and GPU-free. Real concurrent throughput on a downloaded + dense-Llama (CPU/ROCm) belongs to DGR-010 (blocked — no certified dense-Llama + artifact on this machine; see `evidence/DGR-010/BLOCKED.md`) and the final + comparison in DGR-014. +- **Batching is a scheduling grouping in this reference.** `KvBatchEngine.decode_batch` + runs each batch member sequentially through the cached decode (each attends only + its own KV, exactly like an independent llama.cpp sequence). The pinned llama.cpp + worker (DGR-008) fuses the batch into one `llama_decode` graph; the scheduling + semantics — one batch per tick, isolated positions/outputs — are identical. The + numbers here are *scheduler* quantities (ticks, batch occupancy, tokens/tick) + that are real and deterministic; **actual kernel-level batching speedup is a + native-worker property and is NOT claimed here** (RALPH performance discipline: + no unmeasured speed claims). It is measured in DGR-008/DGR-010/DGR-014. +- **Greedy sampling only.** Reuses the DGR-006 greedy `SamplingContract`. Greedy + over isolated per-session KV is order-independent, which is exactly why the + corruption check can assert byte-identical outputs across concurrency levels. + Stochastic sampling is out of scope for the deterministic gate. +- **Single loaded shard / single recipe per scheduler.** The scheduler batches + compatible sessions of one loaded shard (one `recipe_fingerprint`), which is the + node-local case. Multi-range routes batch at the head node whose adapter owns the + final head; cross-node coordination stays in the Meshnet control plane. +- **Native / llama.cpp gates N/A.** No native code, CMake, or llama.cpp patch was + touched (same as DGR-005/006/007), so those gates do not apply to this story. + +## Compatibility / migration notes + +- Purely additive: no existing module changed, so no behavior of the Torch/GGUF + backends, tracker, or KV manager is altered. The scheduler is opt-in — a server + constructs it around a `KvBatchEngine` when it wants continuous batching. +- `SchedulerTelemetry.to_dict()` is JSON-safe and aligns with the capability-signal + vocabulary (active sessions, queue depth, batch occupancy, KV pressure, + prefill/decode rates, rejected admissions) that a node advertises upward; it can + be folded into the DGR-009 capability report / heartbeat without schema changes + here. +- `AdmissionReason` values are stable strings suitable for the native protocol's + structured status / backpressure signalling. + +## Handoff for dependent stories + +- **DGR-008 (C++ gRPC worker):** implement the `BatchEngine` contract natively — + `decode_batch` becomes one `llama_decode` over the sessions' filtered sequences; + `prefill`/`release` map to the same KV manager operations. The scheduler, + admission budgets, fairness policy, and telemetry are unchanged; only the engine + swaps from numpy to llama.cpp. +- **DGR-010 (local real two-process acceptance, blocked):** once a certified + dense-Llama artifact is mounted, drive `run_concurrency_sweep` (or the scheduler + directly) with a real `KvBatchEngine` over the GGUF backend to produce + real-hardware occupancy/throughput/KV-pressure numbers under + `MESHNET_ENABLE_REAL_INFERENCE_TESTS=1` / `.venv-rocm`. +- **DGR-013 (failure/cancel/restart):** the `DoneReason.CACHE_MISS` path (a decode + whose KV was evicted marks the session done and re-prefillable) and the KV-release + on completion are the unit basis for the cancellation/cleanup matrix. +- **DGR-014 (release gate):** feed the real-hardware sweep’s aggregate throughput + and saturation point into the immutable DGR-001 comparison; do not reuse these + synthetic numbers as a performance claim. diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-012/commands.txt b/.scratch/distributed-gguf-runtime/evidence/DGR-012/commands.txt new file mode 100644 index 0000000..05e113e --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-012/commands.txt @@ -0,0 +1,24 @@ +# DGR-012 — exact commands (run from the worktree root) +# Default venv (Python 3.14); deterministic, download-free, GPU-free, API-credit-free. +VP=/run/media/popov/d/DEV/repos/d-popov.com/AI/.venv/bin/python + +# Targeted story tests +$VP -m pytest -q tests/test_batch_scheduler.py +# -> 16 passed + +# Dependency (DGR-007) still green — scheduler builds on this KV manager +$VP -m pytest -q tests/test_hot_kv_state.py +# -> 22 passed + +# Python quality gates +$VP -m compileall -q packages tests +# -> exit 0 +git diff --check +# -> exit 0 + +# Regenerate the machine-readable concurrency-sweep evidence +$VP .scratch/distributed-gguf-runtime/evidence/DGR-012/generate_evidence.py +# -> writes results.json; saturation_concurrency=8 corruption_free=True + +# Full deterministic suite (records the pre-existing unrelated failure baseline) +$VP -m pytest -q -rfE -p no:cacheprovider diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-012/generate_evidence.py b/.scratch/distributed-gguf-runtime/evidence/DGR-012/generate_evidence.py new file mode 100644 index 0000000..e9da177 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-012/generate_evidence.py @@ -0,0 +1,117 @@ +"""Regenerate the DGR-012 concurrency-sweep evidence artifact. + +Deterministic, download-free, GPU-free. Run from the repo root with the default +venv so the worktree ``meshnet_node`` package and the DGR-007 numpy reference +(``tests/test_hot_kv_state``) are importable: + + python .scratch/distributed-gguf-runtime/evidence/DGR-012/generate_evidence.py + +Writes ``results.json`` beside this script. +""" + +from __future__ import annotations + +import json +import pathlib +import sys + +_ROOT = pathlib.Path(__file__).resolve().parents[4] +sys.path.insert(0, str(_ROOT / "packages" / "node")) +sys.path.insert(0, str(_ROOT / "tests")) + +from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard # noqa: E402 + +from meshnet_node.batch_scheduler import ( # noqa: E402 + ContinuousBatchScheduler, + GenerationRequest, + KvBatchEngine, + NodeBudget, + run_concurrency_sweep, +) +from meshnet_node.hot_kv_state import ( # noqa: E402 + HotKvStateManager, + KvBoundaryAdapter, + kv_recipe_for, +) + +MODEL = _KvDenseLlama() + + +def make_engine() -> KvBatchEngine: + shard = _KvReferenceShard(MODEL, 0, MODEL.n_layers - 1) + manager = HotKvStateManager(kv_recipe_for(shard)) + return KvBatchEngine(KvBoundaryAdapter(shard, manager)) + + +def main() -> int: + prompts = { + "s0": [1, 2, 3, 4], "s1": [5, 6, 7, 8], "s2": [9, 10, 11, 12], + "s3": [13, 14, 15, 16], "s4": [17, 18, 19, 20], "s5": [21, 22, 23, 24], + "s6": [25, 26, 27, 28], "s7": [29, 30, 31, 32], + } + n_new = 8 + requests = [ + GenerationRequest(sid, 0, tuple(p), n_new) for sid, p in prompts.items() + ] + sweep = run_concurrency_sweep( + make_engine, requests, concurrency_levels=(1, 2, 4, 8) + ) + + # A representative telemetry snapshot mid-run at concurrency 4 (shows the live + # capability signals a node advertises upward). + engine = make_engine() + scheduler = ContinuousBatchScheduler( + engine, + NodeBudget( + max_active_sessions=4, max_batch_size=4, max_queue_depth=8, + scratch_bytes_per_session=1, scratch_budget_bytes=4, + ), + ) + for request in requests: + scheduler.submit(request) + for _ in range(6): + scheduler.run_tick() + mid_run_telemetry = scheduler.telemetry().to_dict() + + artifact = { + "schema_version": 1, + "evidence_kind": "synthetic-unit", + "model": { + "reference": "pure-numpy KV-cached dense-Llama (tests/test_hot_kv_state)", + "n_layers": MODEL.n_layers, + "hidden": MODEL.hidden, + "n_heads": MODEL.n_heads, + "vocab": MODEL.vocab, + }, + "workload": { + "sessions": len(prompts), + "prompt_len": 4, + "max_new_tokens": n_new, + }, + "concurrency_sweep": sweep.to_dict(), + "mid_run_telemetry_concurrency_4": mid_run_telemetry, + } + + out = pathlib.Path(__file__).with_name("results.json") + out.write_text(json.dumps(artifact, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(f"wrote {out}") + print( + "saturation_concurrency=%d corruption_free=%s" + % (sweep.saturation_concurrency, sweep.corruption_free) + ) + for result in sweep.results: + print( + " c=%d ticks=%d avg_occ=%.3f tokens/tick=%.3f peak_kv=%dB" + % ( + result.concurrency, + result.ticks, + result.avg_batch_occupancy, + result.tokens_per_tick, + result.peak_kv_bytes, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-012/results.json b/.scratch/distributed-gguf-runtime/evidence/DGR-012/results.json new file mode 100644 index 0000000..f1a64d1 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-012/results.json @@ -0,0 +1,179 @@ +{ + "concurrency_sweep": { + "corruption_free": true, + "reference_outputs": { + "s0": [ + 27, + 8, + 27, + 8, + 27, + 8, + 1, + 1 + ], + "s1": [ + 26, + 39, + 39, + 39, + 39, + 3, + 39, + 39 + ], + "s2": [ + 12, + 12, + 12, + 12, + 12, + 12, + 30, + 12 + ], + "s3": [ + 29, + 41, + 42, + 47, + 47, + 42, + 47, + 42 + ], + "s4": [ + 23, + 11, + 44, + 29, + 29, + 29, + 41, + 29 + ], + "s5": [ + 35, + 11, + 0, + 1, + 11, + 0, + 11, + 15 + ], + "s6": [ + 39, + 39, + 28, + 39, + 39, + 39, + 28, + 28 + ], + "s7": [ + 39, + 39, + 39, + 39, + 39, + 39, + 8, + 47 + ] + }, + "results": [ + { + "avg_batch_occupancy": 1.0, + "cache_misses": 0, + "concurrency": 1, + "decode_batches": 56, + "decode_tokens": 56, + "max_batch_occupancy": 1, + "peak_kv_bytes": 15360, + "prefill_tokens": 32, + "rejected_admissions": 0, + "ticks": 64, + "tokens_per_tick": 1.375 + }, + { + "avg_batch_occupancy": 1.75, + "cache_misses": 0, + "concurrency": 2, + "decode_batches": 32, + "decode_tokens": 56, + "max_batch_occupancy": 2, + "peak_kv_bytes": 29184, + "prefill_tokens": 32, + "rejected_admissions": 0, + "ticks": 33, + "tokens_per_tick": 2.6667 + }, + { + "avg_batch_occupancy": 3.1111, + "cache_misses": 0, + "concurrency": 4, + "decode_batches": 18, + "decode_tokens": 56, + "max_batch_occupancy": 4, + "peak_kv_bytes": 52224, + "prefill_tokens": 32, + "rejected_admissions": 0, + "ticks": 19, + "tokens_per_tick": 4.6316 + }, + { + "avg_batch_occupancy": 4.0, + "cache_misses": 0, + "concurrency": 8, + "decode_batches": 14, + "decode_tokens": 56, + "max_batch_occupancy": 7, + "peak_kv_bytes": 75264, + "prefill_tokens": 32, + "rejected_admissions": 0, + "ticks": 15, + "tokens_per_tick": 5.8667 + } + ], + "saturation_concurrency": 8, + "schema_version": 1 + }, + "evidence_kind": "synthetic-unit", + "mid_run_telemetry_concurrency_4": { + "active_sessions": 4, + "batch_occupancy_avg": 4.0, + "batch_occupancy_last": 4, + "batch_occupancy_max": 4, + "completed_sessions": 0, + "decode_tokens_per_sec": 1637.355, + "decode_tokens_total": 20, + "kv_budget_bytes": 67108864, + "kv_pressure": 0.0008, + "kv_total_bytes": 55296, + "prefill_tokens_per_sec": 1309.884, + "prefill_tokens_total": 16, + "queue_depth": 4, + "rejected_admissions_total": 0, + "rejected_by_reason": {}, + "scratch_budget_bytes": 4, + "scratch_pressure": 1.0, + "scratch_used_bytes": 4, + "ticks": 6, + "weight_bytes": 0 + }, + "model": { + "hidden": 32, + "n_heads": 4, + "n_layers": 6, + "reference": "pure-numpy KV-cached dense-Llama (tests/test_hot_kv_state)", + "vocab": 48 + }, + "schema_version": 1, + "workload": { + "max_new_tokens": 8, + "prompt_len": 4, + "sessions": 8 + } +} diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-013/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-013/README.md new file mode 100644 index 0000000..4258c21 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-013/README.md @@ -0,0 +1,223 @@ +# DGR-013 — Harden failure, cancellation, and restart semantics: evidence + +Status: done +Date: 2026-07-16 +Evidence kind: **synthetic-unit** (pure-numpy KV-cached dense-Llama reference + +node-local hardened stream). No model download, no GPU, no torch, no network, no +API credit. + +## Summary + +Implemented bounded, explicit failure/cancellation/restart semantics for the +per-Route-Session decode stream, layered on the DGR-007 Hot KV State manager +(isolated `(session, epoch)` KV) and the DGR-012 continuous-batch scheduler. The +goal (RALPH product objective) is that distributed speed never comes with hanging +or corrupted generations: every blocked op is bounded, every cancel frees state, +duplicate steps are idempotent, uncertain mutations are never silently replayed, +alpha failover restarts from token zero, and billing distinguishes what actually +completed. + +Everything runs against the same deterministic numpy dense-Llama reference the +default gate uses (`tests/test_hot_kv_state.py::_KvDenseLlama` / `_KvReferenceShard`), +so the whole failure matrix is deterministic, download-free, GPU-free, and +API-credit-free while exercising the **real** KV isolation path +(`KvBoundaryAdapter` + `HotKvStateManager`). The pinned llama.cpp worker (DGR-008) +implements the identical adapter contract, so the semantics carry over to native +execution unchanged. + +### What was built (`packages/node/meshnet_node/failure_semantics.py`, new) + +- **`DeadlineGuard` + `StreamTerminated`** — bounds every step against an absolute + deadline and a heartbeat-timeout on an injected clock. A reached deadline or a + lost heartbeat (peer health loss) raises `StreamTerminated(kind)` so a blocked + stream terminates instead of hanging. (**AC: deadlines/heartbeat terminate + blocked ops.**) +- **`CancellationToken`, `ShardCancellationGroup`, `CancellationOutcome`** — one + cancel fans across **every** node-local Shard of a Route Session, releasing the + `(session, epoch)` KV on each shard's manager and invoking every queued-buffer + release callback (the pending activation bundles). Idempotent. The DGR-012 + scheduler also gains a `cancel()` that drops queued/active work on this node and + frees its KV. (**AC: cancellation propagates across every Shard, releases KV + + queued buffers.**) +- **`IdempotencyLedger`, `StepKey`, `StepDisposition`, `UncertainMutationError`** — + records each committed `(session, epoch, step)`; a duplicate delivery returns the + recorded token with no re-mutation. A step whose mutation outcome is *uncertain* + (worker died mid-step) is marked uncertain and can **never** be replayed + silently — `begin()` on an uncertain (or still in-flight) step raises + `UncertainMutationError`, forcing verify-or-restart. (**AC: duplicate steps + idempotent; uncertain mutations never replayed silently.**) +- **`RestartController`** — alpha failover: opens the *next* route epoch, releases + every shard's prior-epoch KV, and `assert_fresh_start` fails closed if any shard + still holds new-epoch KV. The restart re-prefills the whole prompt from token + zero; the failed epoch becomes stale (KV manager rejects it). Unverified KV is + never migrated (RALPH runtime decision #14). (**AC: alpha failover restarts from + token zero rather than importing unverified KV.**) +- **`WorkStatus`, `WorkRecord`, `WorkLedger`** — a typed per-attempt work record + with four distinct statuses: `completed`, `cancelled`, `failed`, `unverified`. + Only `completed` records are billable; cancelled/failed/unverified tokens are + recorded for observability but never charged. JSON-safe for the tracker billing + handoff (`packages/tracker/meshnet_tracker/billing.py` charges only completed, + verified work). (**AC: billing/work records distinguish completed/cancelled/ + failed/unverified.**) +- **`HardenedSessionRunner`** — composes all of the above to drive one session's + prefill+decode through the adapter under a deadline/heartbeat guard + cancel + token, records the typed outcome, and `run_with_failover` restarts a transient + failure from token zero on a fresh epoch. +- **`FailureKind` + `classify_exception` + `work_status_for`** — stable-string + classification of worker death, stream reset, malformed bundle, stale epoch, + cache miss, deadline, heartbeat loss, and cancel, plus the failure→billing-status + mapping. Suitable for the native protocol's structured status. + +### Scheduler extension (`packages/node/meshnet_node/batch_scheduler.py`, DGR-012 file, additive) + +Purely additive so the DGR-012 gate stays green (16/16): +- `DoneReason.CANCELLED` / `DoneReason.FAILED` terminal reasons. +- `ContinuousBatchScheduler.cancel(session_id, *, reason)` — drops a queued + session from the bounded queue or releases an active session's KV, moving it to + the done set with a non-completed reason (never counted as completed work). +- `SchedulerTelemetry.cancelled_sessions` / `failed_sessions` counters. + +## Files changed + +- `packages/node/meshnet_node/failure_semantics.py` — new module (the whole + failure/cancel/restart layer above). +- `packages/node/meshnet_node/batch_scheduler.py` — additive `cancel()` + two + `DoneReason` members + two telemetry counters (DGR-012 file; its 16 tests still + pass unchanged). +- `tests/test_failure_semantics.py` — new, 22 tests (matrix below); reuses the + DGR-007 numpy reference via `from test_hot_kv_state import _KvDenseLlama, + _KvReferenceShard`. +- `.scratch/distributed-gguf-runtime/evidence/DGR-013/` — this README, + `commands.txt`, `generate_evidence.py`, `results.json`. +- `.ralph-tui/progress.md` — appended the DGR-013 note. +- `.scratch/distributed-gguf-runtime/issues/13-...md` — set `Status: done`. + +## Acceptance criteria → evidence + +| Criterion | Tests (`tests/test_failure_semantics.py`) | +|---|---| +| Deadlines/heartbeat loss terminate blocked stream ops | `test_deadline_terminates_a_blocked_stream_and_releases_kv`, `test_heartbeat_loss_terminates_a_blocked_stream`, `test_deadline_guard_reports_remaining_and_resets_on_heartbeat` | +| Cancellation propagates across every Shard, releases KV + queued buffers | `test_cancellation_token_terminates_stream_and_releases_kv`, `test_shard_cancellation_group_releases_every_shard_and_queued_buffers`, `test_scheduler_cancel_drains_queue_and_releases_active_kv`, `test_scheduler_cancel_rejects_a_completed_reason` | +| Duplicate steps idempotent; uncertain mutations never replayed silently | `test_duplicate_step_delivery_is_idempotent_no_remutation`, `test_idempotent_run_replays_tokens_without_advancing_kv`, `test_uncertain_mutation_is_never_replayed_silently`, `test_in_flight_duplicate_is_treated_as_uncertain` | +| Alpha failover restarts from token zero, no unverified KV import | `test_alpha_failover_restarts_from_token_zero_and_completes`, `test_failover_refuses_to_import_unverified_kv`, `test_non_restartable_failure_is_not_retried` | +| Worker death, stream reset, malformed bundle, stale epoch, cache miss | `test_worker_death_midstream_is_unverified_and_marks_step_uncertain`, `test_stream_reset_is_restartable_failure`, `test_malformed_bundle_is_classified_and_does_not_corrupt_kv`, `test_stale_epoch_reference_is_rejected_and_classified`, `test_cache_miss_midstream_is_restartable` | +| Billing/work records distinguish completed/cancelled/failed/unverified | `test_work_ledger_distinguishes_all_four_statuses`, `test_work_status_and_classification_mapping`, plus the clean-run billability check `test_clean_run_matches_stateless_reference_and_is_billable` | + +## Failure matrix (real, deterministic — `results.json`) + +Generated by `generate_evidence.py` against the numpy dense-Llama (prompt `[7,3,9,1]`, +8 new tokens): + +| scenario | status | failure_kind | tokens | restartable | KV released | +|---|---|---|---|---|---| +| clean | completed | — | 8 | — | (held, then reaped) | +| deadline | failed | deadline-exceeded | 2 | no | yes | +| heartbeat_loss | failed | heartbeat-lost | 3 | no | yes | +| cancel | cancelled | cancelled | 3 | no | yes | +| worker_death | unverified | worker-death | 3 | yes | yes | +| stream_reset | failed | stream-reset | — | yes | yes | +| stale_epoch | failed | stale-epoch | — | no | (never opened) | +| cache_miss | failed | cache-miss | 4 | yes | (already evicted) | +| alpha_failover | **completed** (epoch 1) | — | 8 | — | old epoch stale | + +Alpha failover: attempt 0 (epoch 0) dies mid-step → `unverified`; the controller +advances to epoch 1, drops epoch-0 KV, and the restart re-prefills from token zero +→ `completed`, reproducing the byte-identical stateless reference. The old epoch is +now stale (a reference to it raises `StaleRouteEpochError`). Work ledger: +`{completed: 2, cancelled: 1, failed: 0, unverified: 2}`, `billable_tokens = 16` +(only the two completed streams — the failover restart and the clean run — are +billed; the cancelled and the two unverified attempts are not). + +## Commands and real results + +See `commands.txt`. Key results: + +``` +tests/test_failure_semantics.py -> 22 passed +tests/test_batch_scheduler.py -> 16 passed (DGR-012 unchanged) +tests/test_hot_kv_state.py -> 22 passed (DGR-007) +tests/test_gguf_backend.py -> 2 passed (DGR-009) +python -m compileall -q packages tests -> exit 0 +git diff --check -> exit 0 +python -m pytest -q -> 16 failed, 792 passed, 14 skipped in 253.93s +``` + +## Full-suite baseline (pre-existing, unrelated failures) + +The 16 failures are **pre-existing and unrelated to DGR-013**. None import +`failure_semantics` or `batch_scheduler`; they live in the tracker/control-plane, +node-startup, doctor, calibration, and route-benchmark suites and fail on the +model-download / control-plane / recipe-admission paths (e.g. +`UnsupportedRecipeParam: worker_transport` from the DGR-009 native recipe against +the Torch backend, and Torch/HF-model startup that this deterministic sandbox does +not provide). Removing the two DGR-013 files and re-running the failing tests +reproduces the identical failures (see `commands.txt`, 4-test spot check → same +4 failures), so DGR-013 introduces no new failure. + +Exact failing set (16): + +``` +tests/test_dynamic_routing.py::test_admin_can_replace_a_served_model_and_release_it +tests/test_manual_route_benchmark.py::test_pinned_route_uses_named_node +tests/test_manual_route_benchmark.py::test_unknown_route_node_is_400 +tests/test_manual_route_benchmark.py::test_invalid_route_shape_is_400 +tests/test_manual_route_benchmark.py::test_clients_without_route_are_unaffected +tests/test_manual_route_benchmark.py::test_benchmark_records_one_and_two_node_routes +tests/test_node_doctor.py::test_the_shipped_recipes_are_all_applicable_by_the_backend +tests/test_node_doctor.py::test_cli_doctor_flags_select_what_is_validated +tests/test_node_startup.py::test_preset_model_with_hf_repo_loads_torch_backend +tests/test_node_startup.py::test_real_model_startup_registers_downloaded_inventory_without_checksum +tests/test_toploc_calibration_dispatch.py::test_calibration_run_dispatches_only_solo_capable_nodes +tests/test_toploc_calibration_dispatch.py::test_calibration_run_persists_corpus_and_results_endpoint_reports_it +tests/test_toploc_calibration_dispatch.py::test_calibration_run_node_without_commitment_endpoint_is_skipped_not_failed +tests/test_tracker_capability_admission.py::test_an_enforcing_tracker_never_routes_a_node_whose_proof_does_not_cover_it[invalid] +tests/test_tracker_routing.py::test_torch_node_applies_tracker_load_shard_directive +tests/test_tracker_routing.py::test_shard_heal_cycle_surviving_node_covers_dead_peers_gap +``` + +## Limitations and deferred work + +- **Synthetic-unit, not real weights.** Semantics are exercised against the + deterministic numpy dense-Llama, not a downloaded GGUF, to keep the default gate + deterministic/download-free/GPU-free. Real worker-death/stream-reset behavior on + a live llama.cpp worker over gRPC belongs to DGR-008/DGR-010 (DGR-010 is blocked + — no certified dense-Llama artifact on this machine; see + `evidence/DGR-010/BLOCKED.md`). +- **Single-node per-session stream.** `HardenedSessionRunner` drives one full-shard + session (the node-local case); multi-node cancellation is modelled by + `ShardCancellationGroup` fanning across each node's KV manager. The cross-node + propagation *transport* (cancel frames over gRPC/relay) is the native protocol's + job (DGR-002/008); this story owns the local release + record semantics the + transport triggers. +- **Fault injection is deterministic.** Worker death is a shard that raises on the + Nth step; stream reset / deadline / heartbeat are injected via an explicit clock + and hook. This is what makes the matrix reproducible; live fault behavior is a + native/real-hardware property. +- **Greedy sampling only.** Reuses the DGR-006 greedy `SamplingContract`; the + idempotent-replay equality check depends on order-independent greedy decode. +- **Native / llama.cpp gates N/A.** No native code, CMake, or llama.cpp patch was + touched (same as DGR-005/006/007/012), so those gates do not apply. + +## Compatibility / migration notes + +- `failure_semantics.py` is a new, additive module — no existing behavior changes. +- `batch_scheduler.py` changes are additive (new enum members, one method, two + telemetry fields); the DGR-012 contract and its 16 tests are unchanged. +- `WorkRecord.to_dict()` / `WorkLedger.to_dict()` are JSON-safe and map cleanly to + the tracker `BillingLedger.charge_request` inputs: report `node_work` only for + `billable` (completed) records so cancelled/failed/unverified work is never + charged. `FailureKind` / `WorkStatus` are stable strings suitable for the native + protocol's structured status and the capability/heartbeat report. + +## Handoff for dependent stories + +- **DGR-008 (C++ gRPC worker):** implement the same contract natively — the worker + maps a transport deadline/heartbeat to `StreamTerminated`, a dropped stream to a + restartable failure, and a mid-`llama_decode` crash to an *uncertain* step + (mark-uncertain, never silent replay). `RestartController.failover` maps to + opening a fresh llama sequence under the new `(session, epoch)`; the failed + sequence's KV is dropped, never migrated. +- **DGR-010/DGR-014 (real acceptance / release gate):** drive the same failure + scenarios against the live worker to produce real cleanup/latency numbers, and + feed the `WorkLedger` status split into the billing/attribution comparison — + only `completed` work is charged. diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-013/commands.txt b/.scratch/distributed-gguf-runtime/evidence/DGR-013/commands.txt new file mode 100644 index 0000000..b1547d9 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-013/commands.txt @@ -0,0 +1,36 @@ +# DGR-013 — exact commands and real results (worktree venv) +VP=/run/media/popov/d/DEV/repos/d-popov.com/AI/.venv/bin/python + +# Targeted story tests (this story) +$VP -m pytest -q tests/test_failure_semantics.py +# -> 22 passed + +# Dependency gates stay green +$VP -m pytest -q tests/test_batch_scheduler.py # DGR-012 +# -> 16 passed +$VP -m pytest -q tests/test_hot_kv_state.py # DGR-007 +# -> 22 passed +$VP -m pytest -q tests/test_gguf_backend.py # DGR-009 +# -> 2 passed + +# Quality gates +$VP -m compileall -q packages tests +# -> exit 0 +git diff --check +# -> exit 0 + +# Machine-readable evidence +$VP .scratch/distributed-gguf-runtime/evidence/DGR-013/generate_evidence.py +# -> wrote results.json; work statuses {'completed':2,'cancelled':1,'failed':0,'unverified':2} billable_tokens=16 + +# Full deterministic suite +$VP -m pytest -q -p no:cacheprovider +# -> 16 failed, 792 passed, 14 skipped in 253.93s + +# Clean-tree reproduction of the 16 pre-existing failures (DGR-013 files removed) +# rm packages/node/meshnet_node/failure_semantics.py tests/test_failure_semantics.py +$VP -m pytest -q tests/test_dynamic_routing.py::test_admin_can_replace_a_served_model_and_release_it \ + tests/test_node_doctor.py::test_the_shipped_recipes_are_all_applicable_by_the_backend \ + tests/test_tracker_routing.py::test_torch_node_applies_tracker_load_shard_directive \ + tests/test_node_startup.py::test_preset_model_with_hf_repo_loads_torch_backend +# -> 4 failed (same failures reproduce without any DGR-013 change) diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-013/generate_evidence.py b/.scratch/distributed-gguf-runtime/evidence/DGR-013/generate_evidence.py new file mode 100644 index 0000000..ea4496f --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-013/generate_evidence.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python +"""Generate deterministic DGR-013 failure/cancel/restart evidence (results.json). + +Runs the real hardened per-session stream (``HardenedSessionRunner`` over the +DGR-007 ``KvBoundaryAdapter`` + ``HotKvStateManager``) through each failure mode +with the same pure-numpy dense-Llama reference the default gate uses. No model +download, no GPU, no torch, no network, no API credit. + +Run from the repo root with the worktree venv: + + .venv/bin/python .scratch/distributed-gguf-runtime/evidence/DGR-013/generate_evidence.py +""" + +from __future__ import annotations + +import json +import os +import sys + +import numpy as np + +# Make the worktree packages and the DGR-007 numpy reference importable, exactly +# as pytest's prepend-import + conftest do. +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) +sys.path.insert(0, os.path.join(ROOT, "packages", "node")) +sys.path.insert(0, os.path.join(ROOT, "tests")) + +from meshnet_node.hot_kv_state import ( # noqa: E402 + HotKvStateConfig, + HotKvStateManager, + KvBoundaryAdapter, + StaleRouteEpochError, + kv_recipe_for, +) +from meshnet_node.batch_scheduler import GenerationRequest # noqa: E402 +from meshnet_node.failure_semantics import ( # noqa: E402 + CancellationToken, + FailureKind, + HardenedSessionRunner, + RestartController, + StreamTerminated, + WorkLedger, + WorkStatus, +) + +from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard # noqa: E402 + + +class _FaultyShard(_KvReferenceShard): + def __init__(self, model, start, end, *, fail_at_call=None): + super().__init__(model, start, end) + self._fail_at_call = fail_at_call + self.calls = 0 + + def run_layers_cached(self, hidden, *, positions, past_kv): + self.calls += 1 + if self._fail_at_call is not None and self.calls == self._fail_at_call: + raise RuntimeError("worker died mid-step") + return super().run_layers_cached(hidden, positions=positions, past_kv=past_kv) + + +class _Clock: + def __init__(self): + self.now = 0.0 + + def __call__(self): + return self.now + + def advance(self, d): + self.now += d + + +def _adapter(model, *, config=None, shard=None): + shard = shard or _KvReferenceShard(model, 0, model.n_layers - 1) + manager = HotKvStateManager(kv_recipe_for(shard), config=config) + return KvBoundaryAdapter(shard, manager) + + +def _gen(sid, prompt, n, epoch=0): + return GenerationRequest( + session_id=sid, route_epoch=epoch, + prompt_token_ids=tuple(prompt), max_new_tokens=n, + ) + + +def _kv_released(manager, sid, epoch): + from meshnet_node.hot_kv_state import CacheMiss + return isinstance(manager.resolve(sid, epoch), CacheMiss) + + +def main() -> None: + model = _KvDenseLlama() + prompt = [7, 3, 9, 1] + n_new = 8 + ledger = WorkLedger() + scenarios = [] + + # 1. Clean baseline. + ad = _adapter(model) + r = HardenedSessionRunner(ad, work_ledger=ledger).run(_gen("clean", prompt, n_new)) + scenarios.append({ + "scenario": "clean", + "status": r.status.value, + "tokens": r.token_count, + "matches_reference": list(r.tokens) == model.stateless_greedy(prompt, n_new), + "kv_released": _kv_released(ad.manager, "clean", 0), + }) + + # 2. Deadline terminates a blocked stream. + clk = _Clock() + ad = _adapter(model) + r = HardenedSessionRunner(ad, clock=clk).run( + _gen("deadline", prompt, 50), deadline=3.0, + before_step=lambda _s: clk.advance(1.0), + ) + scenarios.append({ + "scenario": "deadline", "status": r.status.value, + "failure_kind": r.failure_kind.value, "tokens": r.token_count, + "kv_released": _kv_released(ad.manager, "deadline", 0), + }) + + # 3. Heartbeat/health loss terminates a blocked stream. + clk = _Clock() + ad = _adapter(model) + r = HardenedSessionRunner(ad, clock=clk).run( + _gen("heartbeat", prompt, 50), heartbeat_timeout=1.5, + heartbeat=lambda step: step < 2, + before_step=lambda _s: clk.advance(1.0), + ) + scenarios.append({ + "scenario": "heartbeat_loss", "status": r.status.value, + "failure_kind": r.failure_kind.value, "tokens": r.token_count, + "kv_released": _kv_released(ad.manager, "heartbeat", 0), + }) + + # 4. Explicit client cancellation releases KV. + ad = _adapter(model) + tok = CancellationToken() + r = HardenedSessionRunner(ad, work_ledger=ledger).run( + _gen("cancel", prompt, 50), cancel_token=tok, + before_step=lambda step: tok.cancel("client-hangup") if step == 3 else None, + ) + scenarios.append({ + "scenario": "cancel", "status": r.status.value, + "failure_kind": r.failure_kind.value, "tokens": r.token_count, + "kv_released": _kv_released(ad.manager, "cancel", 0), + }) + + # 5. Worker death mid-step -> unverified. + ad = _adapter(model, shard=_FaultyShard(model, 0, model.n_layers - 1, fail_at_call=4)) + r = HardenedSessionRunner(ad, work_ledger=ledger).run(_gen("worker", prompt, n_new)) + scenarios.append({ + "scenario": "worker_death", "status": r.status.value, + "failure_kind": r.failure_kind.value, "tokens": r.token_count, + "restartable": r.restartable, "kv_released": _kv_released(ad.manager, "worker", 0), + }) + + # 6. Stream reset -> failed, restartable. + ad = _adapter(model) + def reset(step): + if step == 2: + raise StreamTerminated(FailureKind.STREAM_RESET, "peer reset") + r = HardenedSessionRunner(ad).run(_gen("reset", prompt, n_new), before_step=reset) + scenarios.append({ + "scenario": "stream_reset", "status": r.status.value, + "failure_kind": r.failure_kind.value, "restartable": r.restartable, + }) + + # 7. Stale epoch -> failed. + ad = _adapter(model) + ad.manager.open("stale", 5) + r = HardenedSessionRunner(ad).run(_gen("stale", prompt, n_new, epoch=3)) + scenarios.append({ + "scenario": "stale_epoch", "status": r.status.value, + "failure_kind": r.failure_kind.value, + }) + + # 8. Cache miss mid-stream -> restartable. + ad = _adapter(model) + mgr = ad.manager + r = HardenedSessionRunner(ad).run( + _gen("miss", prompt, 12), + before_step=lambda step: mgr.release("miss", 0) if step == 4 else None, + ) + scenarios.append({ + "scenario": "cache_miss", "status": r.status.value, + "failure_kind": r.failure_kind.value, "tokens": r.token_count, + "restartable": r.restartable, + }) + + # 9. Alpha failover: restart from token zero, no unverified KV import. + faulty = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3) + ad = _adapter(model, shard=faulty) + runner = HardenedSessionRunner(ad, work_ledger=ledger) + controller = RestartController([ad.manager]) + fo = runner.run_with_failover(_gen("failover", prompt, n_new, epoch=0), controller, + max_restarts=2) + old_epoch_stale = False + try: + ad.manager.resolve("failover", 0) + except StaleRouteEpochError: + old_epoch_stale = True + scenarios.append({ + "scenario": "alpha_failover", + "final_status": fo.outcome.status.value, + "final_epoch": fo.outcome.route_epoch, + "restarts": fo.restarts, + "restarted_from_token_zero": list(fo.outcome.tokens) == model.stateless_greedy(prompt, n_new), + "old_epoch_stale": old_epoch_stale, + "attempt_statuses": [a.status.value for a in fo.attempts], + }) + + result = { + "schema_version": 1, + "evidence_kind": "synthetic-unit", + "model": { + "architecture": model.architecture_adapter, + "n_layers": model.n_layers, "vocab": model.vocab, "hidden": model.hidden, + }, + "scenarios": scenarios, + "work_ledger": ledger.to_dict(), + } + + out_path = os.path.join(os.path.dirname(__file__), "results.json") + with open(out_path, "w") as fh: + json.dump(result, fh, indent=2) + fh.write("\n") + counts = ledger.counts_by_status() + print(f"wrote {out_path}") + print(f"work statuses: {counts} billable_tokens={ledger.billable_tokens()}") + + +if __name__ == "__main__": + main() diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-013/results.json b/.scratch/distributed-gguf-runtime/evidence/DGR-013/results.json new file mode 100644 index 0000000..2070d3e --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-013/results.json @@ -0,0 +1,135 @@ +{ + "schema_version": 1, + "evidence_kind": "synthetic-unit", + "model": { + "architecture": "dense-llama", + "n_layers": 6, + "vocab": 48, + "hidden": 32 + }, + "scenarios": [ + { + "scenario": "clean", + "status": "completed", + "tokens": 8, + "matches_reference": true, + "kv_released": false + }, + { + "scenario": "deadline", + "status": "failed", + "failure_kind": "deadline-exceeded", + "tokens": 2, + "kv_released": true + }, + { + "scenario": "heartbeat_loss", + "status": "failed", + "failure_kind": "heartbeat-lost", + "tokens": 3, + "kv_released": true + }, + { + "scenario": "cancel", + "status": "cancelled", + "failure_kind": "cancelled", + "tokens": 3, + "kv_released": true + }, + { + "scenario": "worker_death", + "status": "unverified", + "failure_kind": "worker-death", + "tokens": 3, + "restartable": true, + "kv_released": true + }, + { + "scenario": "stream_reset", + "status": "failed", + "failure_kind": "stream-reset", + "restartable": true + }, + { + "scenario": "stale_epoch", + "status": "failed", + "failure_kind": "stale-epoch" + }, + { + "scenario": "cache_miss", + "status": "failed", + "failure_kind": "cache-miss", + "tokens": 4, + "restartable": true + }, + { + "scenario": "alpha_failover", + "final_status": "completed", + "final_epoch": 1, + "restarts": 1, + "restarted_from_token_zero": true, + "old_epoch_stale": true, + "attempt_statuses": [ + "unverified", + "completed" + ] + } + ], + "work_ledger": { + "schema_version": 1, + "records": [ + { + "session_id": "clean", + "route_epoch": 0, + "status": "completed", + "tokens": 8, + "failure_kind": null, + "detail": "", + "billable": true + }, + { + "session_id": "cancel", + "route_epoch": 0, + "status": "cancelled", + "tokens": 3, + "failure_kind": "cancelled", + "detail": "operation cancelled: client-hangup", + "billable": false + }, + { + "session_id": "worker", + "route_epoch": 0, + "status": "unverified", + "tokens": 3, + "failure_kind": "worker-death", + "detail": "worker died mid-step", + "billable": false + }, + { + "session_id": "failover", + "route_epoch": 0, + "status": "unverified", + "tokens": 2, + "failure_kind": "worker-death", + "detail": "worker died mid-step", + "billable": false + }, + { + "session_id": "failover", + "route_epoch": 1, + "status": "completed", + "tokens": 8, + "failure_kind": null, + "detail": "", + "billable": true + } + ], + "counts_by_status": { + "completed": 2, + "cancelled": 1, + "failed": 0, + "unverified": 2 + }, + "billable_tokens": 16 + } +} diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-014/BLOCKED.md b/.scratch/distributed-gguf-runtime/evidence/DGR-014/BLOCKED.md new file mode 100644 index 0000000..f855943 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-014/BLOCKED.md @@ -0,0 +1,55 @@ +# DGR-014 — Blocked handoff + +Status: blocked +Date: 2026-07-16 + +## Blocker + +This release-gate story cannot be completed in the current workspace state because the prerequisite real-model comparison chain is still missing its certified dense-Llama artifact on mounted storage. + +Verified blockers: + +- `DGR-011` is still not passed in `.scratch/distributed-gguf-runtime/prd.json`. +- `DGR-011` is explicitly blocked in `.scratch/distributed-gguf-runtime/evidence/DGR-011/BLOCKED.md`. +- `DGR-011` depends on `DGR-010`, and `DGR-010` is blocked because there is no certified dense-Llama artifact available on the mounted drive. +- Current mounted-model storage still only shows Qwen artifacts and llama.cpp vocab GGUFs, not the certified dense-Llama GGUF/safetensors pair needed for a comparable real run. + +## Verified current state + +- The DGR-001 performance contract exists and defines the benchmark lanes, metrics, and stop condition that later release gates must keep unchanged. +- The DGR-012 scheduler and DGR-013 failure semantics evidence are present and usable as supporting context, but they do not satisfy the real final comparison required here. +- `packages/node/meshnet_node/performance_contract.py` already contains the contract metadata and a live endpoint benchmark shim, but there is no recorded DGR-014 release-gate run and no final immutable comparison artifact. +- `evidence/DGR-014/README.md` does not exist yet because the acceptance criteria could not be completed. + +## Commands run + +```bash +sed -n '1,260p' .claude/memory/MEMORY.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/RALPH-CONTEXT.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/issues/14-enforce-the-gguf-versus-safetensors-release-gate.md +sed -n '1,260p' .ralph-tui/progress.md +git status --short +sed -n '1,260p' .scratch/distributed-gguf-runtime/prd.json +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-001/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-012/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-013/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-011/BLOCKED.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-010/BLOCKED.md +find /run/media/popov/d/DEV/models /run/media/popov/d/DEV/llamacpp/llama.cpp/models -maxdepth 4 \( -iname '*llama*' -o -iname '*deepseek*' -o -iname '*dense*' -o -name '*.gguf' -o -name '*.safetensors' -o -name 'config.json' \) +``` + +## Known limitations + +- No certified dense-Llama artifact is mounted, so the real distributed safetensors-versus-GGUF comparison cannot be executed. +- No immutable release-gate evidence can be produced without that artifact and the completed DGR-011 route comparison. +- No code was changed in this iteration. + +## Compatibility notes + +- The DGR-001 contract remains the source of truth for thresholds and metric names. +- Any future DGR-014 run must keep those thresholds unchanged and compare the same certified model/hardware/network scenario for both routes. + +## Dependent-story handoff + +- Finish `DGR-010` and `DGR-011` first with a certified dense-Llama artifact on mounted storage. +- Then run the current distributed safetensors and distributed GGUF routes on the same comparable scenario, record the final numbers in `evidence/DGR-014/README.md`, and update the issue status only after the gate passes. diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-015/BLOCKED.md b/.scratch/distributed-gguf-runtime/evidence/DGR-015/BLOCKED.md new file mode 100644 index 0000000..8be09f6 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-015/BLOCKED.md @@ -0,0 +1,78 @@ +# DGR-015 — Blocked handoff + +Status: blocked +Date: 2026-07-16 + +## Blocker + +This story cannot be completed in the current workspace state because its +mandatory prerequisite, DGR-014, is still not passed. + +Verified blocker chain: + +- `.scratch/distributed-gguf-runtime/prd.json` still marks `DGR-014` as + `"passes": false`, so DGR-015 is not released for completion. +- `.scratch/distributed-gguf-runtime/evidence/DGR-014/BLOCKED.md` records the + release-gate blocker: the certified dense-Llama artifact required for the + comparable real-model comparison is not mounted on this machine. +- `DGR-014` depends on `DGR-011`, which is also blocked because `DGR-010` + cannot run without that same certified dense-Llama artifact. +- The current codebase still fails closed for `qwen3` / `qwen3-moe` in + `packages/node/meshnet_node/boundary_adapter.py`, which is correct for the + current state but means no Qwen3 family recipe is certified yet. + +## Verified current state + +- Dense-Llama boundary semantics, Hot KV isolation, batching, and failure + semantics are already implemented and covered by prior stories. +- Qwen3 strings are present in tracker/model metadata, but they are not yet + backed by a certified architecture adapter or real-model acceptance evidence. +- No `evidence/DGR-015/README.md` exists yet because the acceptance criteria + could not be completed. + +## Commands run + +```bash +sed -n '1,260p' .claude/memory/MEMORY.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/RALPH-CONTEXT.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/issues/15-add-and-certify-a-qwen3-qwen3-moe-adapter.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/architecture.md +sed -n '1,260p' CONTEXT.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/prd.json +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-014/BLOCKED.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-013/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-012/README.md +sed -n '1,260p' packages/node/meshnet_node/boundary_adapter.py +sed -n '1,260p' packages/node/meshnet_node/model_catalog.py +sed -n '1,220p' packages/node/meshnet_node/model_metadata.json +sed -n '1,260p' packages/tracker/meshnet_tracker/capability.py +sed -n '1,260p' packages/tracker/meshnet_tracker/server.py +rg -n "qwen3|qwen3-moe|Qwen3|MoE|router|top-k|shared expert|shared_expert|expert" packages/node/meshnet_node packages/tracker/meshnet_tracker tests -g '!**/__pycache__/**' +git status --short +``` + +## Known limitations + +- No certified dense-Llama artifact is mounted, so DGR-014 cannot complete and + DGR-015 remains blocked behind it. +- No real consumer-hardware Qwen3 acceptance run was possible in this workspace. +- No code was changed in this iteration. + +## Compatibility notes + +- The current boundary adapter intentionally fails closed for uncertified + architectures. That is the correct behavior until a dedicated Qwen3 adapter is + implemented and certified. +- Existing dense-Llama coverage and Hot KV semantics remain the source of truth + for the shared protocol and cache behavior. + +## Dependent-story handoff + +- Finish `DGR-010`, `DGR-011`, and `DGR-014` first with a certified dense-Llama + artifact on mounted storage. +- Once the release gate passes, implement the Qwen3 family adapter as a separate + certified architecture rather than by extending dense-Llama with unchecked name + substitutions. +- Record the real-model Qwen3 parity, admission, memory, and communication + evidence in `evidence/DGR-015/README.md`, then update the issue status only + after the gate passes. diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-016/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-016/README.md new file mode 100644 index 0000000..49ada7f --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-016/README.md @@ -0,0 +1,145 @@ +# DGR-016 — Upstream llama.cpp collaboration package + +Status: partial, blocked by DGR-010 +Date: 2026-07-16 + +## Summary + +Assembled the upstream-facing collaboration package for llama.cpp without +pulling Meshnet routing or control-plane logic into the upstream ask. + +Durable outputs created for this story: + +- `api-note.md` with the generic hook split and patch-per-concern proposal +- `outreach.md` with a maintainer-facing draft for Georgi/llama.cpp + +The package is grounded in the existing research artifacts and the already +implemented deterministic tests for: + +- range-aware GGUF ownership and introspection +- architecture boundary input/output +- layer-filtered KV/session ownership +- reproducible pinned worker build wiring + +The story itself remains blocked because DGR-010 is still marked `passes: false` +and only has a blocked handoff, not a completed real-model acceptance README. + +## Files changed + +- `.scratch/distributed-gguf-runtime/evidence/DGR-016/README.md` +- `.scratch/distributed-gguf-runtime/evidence/DGR-016/api-note.md` +- `.scratch/distributed-gguf-runtime/evidence/DGR-016/outreach.md` + +## Commands run and real results + +### Dependency and context review + +```bash +sed -n '1,260p' .scratch/distributed-gguf-runtime/RALPH-CONTEXT.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/issues/16-produce-the-upstream-llama-cpp-collaboration-package.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-010/BLOCKED.md +sed -n '1,260p' docs/adr/0024-distributed-gguf-runtime.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/architecture.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/decision-framework.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/implementation-strategy.md +sed -n '1,260p' CONTEXT.md +``` + +Result: + +- confirmed the runtime target is a small pinned llama.cpp worker with Meshnet + kept outside upstream +- confirmed DGR-010 is still blocked because there is no certified dense-Llama + artifact on mounted storage + +### Package-relevant targeted pytest + +```bash +python -m pytest -q tests/test_llama_worker_build.py tests/test_gguf_backend.py tests/test_gguf_ownership.py tests/test_boundary_adapter.py tests/test_hot_kv_state.py +``` + +Result: + +- `50 passed in 0.90s` + +### Broader focused pytest slice + +```bash +python -m pytest -q tests/test_llama_worker_build.py tests/test_native_shard_protocol.py tests/test_gguf_backend.py tests/test_boundary_adapter.py tests/test_gguf_ownership.py tests/test_hot_kv_state.py tests/test_kv_cache_distributed.py +``` + +Result: + +- `58 passed, 1 skipped, 9 failed, 12 errors in 1.27s` +- failures were pre-existing environment issues, not this documentation-only + package: + - `tests/test_native_shard_protocol.py` imported generated protobuf code built + against gencode 7.35.0 while the active runtime is 6.33.6 + - `tests/test_kv_cache_distributed.py` hit sandbox socket `PermissionError` + when trying to bind localhost servers + +### Research evidence review + +```bash +sed -n '1,260p' docs/research/distributed-gguf-landscape.md +sed -n '1,260p' docs/research/distributed-gguf-github-followup.md +sed -n '1,220p' .scratch/distributed-gguf-runtime/evidence/DGR-004/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-006/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-007/README.md +sed -n '1,260p' .scratch/distributed-gguf-runtime/evidence/DGR-009/README.md +``` + +Result: + +- confirmed Nakshatra and prima.cpp are the right source/test donors for the + upstream ask +- confirmed the generic API surface is range loading, boundary I/O, and KV + ownership, not Meshnet policy + +### Package assembly + +No code generation, downloads, or model execution were required for this story. +The package is documentation-only and deterministic. + +```bash +python -m compileall -q packages tests +git diff --check +``` + +Result: + +- both commands exited 0 + +## Correctness / performance / hardware classification + +- Correctness evidence: research-only, no live model execution +- Performance evidence: none in this story +- Hardware evidence: none in this story + +## Known limitations and deferred work + +- DGR-010 remains blocked, so this package cannot be treated as the final + release-ready upstream handoff. +- The outreach draft is human-ready but not sent. +- The doc package does not change llama.cpp source code; it only prepares the + upstream ask and test mapping. + +## Compatibility / migration notes + +- Exact upstream pin for the eventual patch series: `b3c9d1b846cc80a6360adb6aeaa4fcd8c4c8dcac` +- The proposed patch split is: + 1. range-aware loading and ownership introspection + 2. boundary input/output and named tensor bundles + 3. layer-filtered KV and local sequence ownership +- Meshnet routing, billing, relay transport, and volunteer-network policy stay + outside llama.cpp. +- The deterministic examples already exist in the tree and can be trimmed into + upstream-facing MREs when the human maintainer sends the package. + +## Dependent-story handoff + +- DGR-010 must clear before any real-model validation can be cited as the final + end-to-end proof for this upstream package. +- Once DGR-010 has a completed evidence README, the package can be refreshed + with the real-model context and sent to the llama.cpp maintainers as a + smaller review bundle. diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-016/api-note.md b/.scratch/distributed-gguf-runtime/evidence/DGR-016/api-note.md new file mode 100644 index 0000000..ab73302 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-016/api-note.md @@ -0,0 +1,90 @@ +# DGR-016 API note: narrow llama.cpp hooks, no Meshnet policy + +This note is the upstream-facing shape for the collaboration package. + +## Goal + +Keep the llama.cpp ask small: + +- expose generic model-layer hooks that are useful to any local or remote + layer-worker setup; +- keep Meshnet routing, session ownership, billing, and relay transport out of + llama.cpp; +- preserve one patch per concern so the series rebases cleanly on the pinned + upstream commit. + +## Concern 1: range-aware loading and authoritative tensor ownership + +Requested surface: + +- accept a contiguous `[start_layer, end_layer)` range; +- expose whether the worker owns embeddings, final norm, and final head; +- make the loaded range authoritative from the model state, not from CLI + claims; +- allow unowned tensors to be absent rather than fabricated. + +Why this is upstreamable: + +- it is generic loader and introspection plumbing; +- it helps any local partitioned inference mode; +- it does not require any Meshnet identity, route, or transport type. + +Minimal examples/tests: + +- `tests/test_gguf_ownership.py` +- `tests/test_llama_worker_build.py` + +## Concern 2: architecture boundary input/output + +Requested surface: + +- accept a versioned boundary bundle carrying one or more named tensors; +- support an unnormalized residual stream as the intermediate handoff; +- keep final norm, LM head, and sampling on the tail shard only; +- keep the bundle format explicit about name, shape, dtype, byte order, and + fragments. + +Why this is upstreamable: + +- it matches both dense Llama and other certified adapter families; +- it does not assume Meshnet or any specific wire protocol; +- it gives a stable ABI for a layer-worker boundary. + +Minimal examples/tests: + +- `tests/test_boundary_adapter.py` +- `tests/test_native_shard_protocol.py` + +## Concern 3: layer-filtered KV and session mapping + +Requested surface: + +- let the worker own KV only for its layer range; +- map a stable session/context identifier to the local sequence; +- allow cache miss, stale epoch, truncate, release, and eviction semantics; +- reject incompatible cache recipes rather than trying to heal them silently. + +Why this is upstreamable: + +- it is a local sequence/KV API, not a network scheduler; +- it is useful to any supervisor that needs one process per layer range; +- it keeps session semantics outside llama.cpp while still making the worker + stateful in a controlled way. + +Minimal examples/tests: + +- `tests/test_hot_kv_state.py` +- `tests/test_kv_cache_distributed.py` + +## Suggested patch split + +Keep the series narrow and independently reviewable against the exact pinned +commit `b3c9d1b846cc80a6360adb6aeaa4fcd8c4c8dcac`: + +1. `range-aware-loading` and ownership introspection. +2. `boundary-input-output` and named tensor bundle handoff. +3. `layer-filtered-kv` and sequence ownership. + +The current Meshnet worker scaffold remains a project-owned wrapper and is not +part of the upstream ask. + diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-016/outreach.md b/.scratch/distributed-gguf-runtime/evidence/DGR-016/outreach.md new file mode 100644 index 0000000..95ed9d7 --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-016/outreach.md @@ -0,0 +1,43 @@ +# DGR-016 outreach draft + +Subject: Narrow llama.cpp hooks for range loading, boundary I/O, and local KV ownership + +Hi Georgi and llama.cpp maintainers, + +We have been building a distributed GGUF route on top of a Meshnet control +plane, and the narrow upstreamable seam is now clear enough to summarize. + +We are not asking llama.cpp to own Meshnet routing, billing, relay transport, +or any volunteer-network policy. The upstream ask is limited to generic local +hooks that make partitioned inference easier to implement and easier to review: + +1. Range-aware loading and ownership introspection for contiguous layer ranges. +2. Architecture-defined boundary input/output using an explicit named-tensor + bundle. +3. Layer-filtered KV ownership and stable local sequence mapping. + +Why we think this is generally useful: + +- Nakshatra already demonstrates the value of a narrow layer-worker seam and + partial GGUF loading. +- prima.cpp shows the same idea from a different angle with selective loading, + local KV, and boundary residual transport. +- Both projects suggest the same conclusion: the missing API is not Meshnet + specific, it is a local runtime seam that any layer-partitioned supervisor can + use. + +The package we would upstream is intentionally split into one concern per patch +so review stays small: + +- range-aware loading and tensor ownership; +- boundary I/O for intermediate residual state; +- layer-filtered KV and sequence ownership. + +If useful, we can send the concrete MRE/test mapping next. We already have +deterministic examples covering the loader, boundary contract, and KV/session +semantics in the Meshnet tree, and we can trim them into upstream-focused test +cases. + +Thanks, +Meshnet maintainers + diff --git a/.scratch/distributed-gguf-runtime/issues/13-harden-failure-cancellation-and-restart-semantics.md b/.scratch/distributed-gguf-runtime/issues/13-harden-failure-cancellation-and-restart-semantics.md index 790838c..0aca94f 100644 --- a/.scratch/distributed-gguf-runtime/issues/13-harden-failure-cancellation-and-restart-semantics.md +++ b/.scratch/distributed-gguf-runtime/issues/13-harden-failure-cancellation-and-restart-semantics.md @@ -1,6 +1,6 @@ # 13 — Harden failure, cancellation, and restart semantics -Status: ready-for-agent +Status: done ## Mandatory fresh-session context diff --git a/packages/node/meshnet_node/batch_scheduler.py b/packages/node/meshnet_node/batch_scheduler.py new file mode 100644 index 0000000..1ce7ea5 --- /dev/null +++ b/packages/node/meshnet_node/batch_scheduler.py @@ -0,0 +1,1024 @@ +"""Continuous batching and bounded admission for concurrent Route Sessions (DGR-012). + +RALPH runtime decision #9: concurrency on a node uses *continuous batching of +compatible active sessions* — not a separate scheduler or control plane. This +module is the node-local scheduler that sits on top of the isolated Hot KV State +manager (DGR-007) and turns many concurrent single-token decode steps into one +batch per tick, while keeping every session's positions, KV, and sampled output +isolated (decisions #7/#8, ADR-0022/0024). + +The design is deliberately backend-agnostic. The scheduler talks to a +:class:`BatchEngine` duck type (``recipe_fingerprint`` / ``prefill`` / +``decode_batch`` / ``release``); the default deterministic test suite drives it +with a pure-numpy dense-Llama engine, and the pinned llama.cpp worker (DGR-008) +implements the same contract where a batch becomes one ``llama_decode`` over +several sequences. :class:`KvBatchEngine` adapts the DGR-007 +:class:`~meshnet_node.hot_kv_state.KvBoundaryAdapter` to this contract so the +scheduler runs against real KV isolation with no new cache code. + +What the scheduler guarantees (the acceptance contract): + +* **Bounded admission.** A new session is admitted only if it fits the node's + weight, KV, scratch, and queue budgets (:class:`NodeBudget`). Anything that + cannot fit is rejected with an explicit :class:`AdmissionReason`; anything that + fits but has no free active slot waits in a bounded queue. When the queue is + full, admission is refused — that refusal *is* the backpressure signal. +* **Continuous batching.** Every tick, all sessions currently decoding contribute + their single next token to one batch (bounded by ``max_batch_size``). The engine + runs the batch once; each session keeps its own position and appends its own + sampled token, so batching never mixes outputs. +* **Prefill does not starve decode.** The scheduling policy is explicit and fixed: + *decode first, then bounded prefill*. Ongoing decodes always run before any new + prompt is prefilled, and prefill work per tick is capped + (``max_prefill_tokens_per_tick``) so a burst of new sessions cannot monopolise + the node and stall in-flight generations. +* **Bounded memory.** KV growth is bounded by the manager's byte budget; queued + activations are bounded by ``max_queue_depth`` and the scratch budget. Neither + the queue nor the KV store grows without limit. +* **Telemetry.** :meth:`ContinuousBatchScheduler.telemetry` reports active + sessions, queue depth, batch occupancy, KV pressure, prefill/decode token rates, + and rejected admissions — the capability signals a node advertises upward. + +Everything here is pure Python + the numpy-backed manager, so the default gate +stays deterministic, download-free, GPU-free, and API-credit-free. Real +kernel-level batching speedup is a native-worker property measured in +DGR-008/DGR-010/DGR-014; this module owns the *scheduling* behaviour and proves, +via the 1/2/4/8 concurrency sweep, that batching raises aggregate work-per-tick +without cross-session corruption. +""" + +from __future__ import annotations + +import threading +import time +from collections import deque +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Callable, Iterable, Mapping, Sequence + +from meshnet_node.hot_kv_state import ( + CacheMiss, + HotKvStateManager, + KvBoundaryAdapter, +) + + +class SchedulerError(RuntimeError): + """Base class for scheduler configuration/usage errors.""" + + +# --------------------------------------------------------------------------- # +# Node budget and admission. +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class NodeBudget: + """Explicit bounds the node admits and schedules against. + + Four budget dimensions gate admission (the story's "weight, KV, scratch, and + queue budgets") plus the scheduling bounds that keep batching fair: + + * ``weight_bytes`` — resident weight footprint of the loaded shard. This is a + fixed, one-time cost; the scheduler treats it as already resident and simply + reports it (a node that cannot hold its shard weight never starts). It is + validated non-negative and surfaced in telemetry. + * ``kv_budget_bytes`` — the Hot KV State byte budget. A session is admissible + only if its *whole* generation (prompt + all new tokens) could fit this + budget on its own; cross-session pressure is then handled by the manager's + LRU/byte eviction. This mirrors ``HotKvStateConfig.budget_bytes`` and should + match the manager the scheduler was given. + * ``scratch_bytes_per_session`` / ``scratch_budget_bytes`` — per-active-session + activation scratch (the transient residual/attention buffers a decode needs) + and the total scratch envelope. Admission keeps + ``active * scratch_per_session <= scratch_budget`` so concurrent activations + are bounded, not just KV. + * ``max_active_sessions`` — hard cap on sessions occupying an execution slot. + * ``max_queue_depth`` — bounded waiting room for admitted-but-not-yet-running + requests. A full queue is the backpressure boundary. + * ``max_batch_size`` — largest decode batch formed per tick. + * ``max_prefill_tokens_per_tick`` — prefill token budget per tick, so prefill + cannot starve decode. + """ + + weight_bytes: int = 0 + kv_budget_bytes: int = 64 * 1024 * 1024 + scratch_bytes_per_session: int = 1 * 1024 * 1024 + scratch_budget_bytes: int = 16 * 1024 * 1024 + max_active_sessions: int = 8 + max_queue_depth: int = 64 + max_batch_size: int = 8 + max_prefill_tokens_per_tick: int = 512 + + def __post_init__(self) -> None: + if self.weight_bytes < 0: + raise SchedulerError("weight_bytes must be >= 0") + if self.kv_budget_bytes <= 0: + raise SchedulerError("kv_budget_bytes must be positive") + if self.scratch_bytes_per_session <= 0: + raise SchedulerError("scratch_bytes_per_session must be positive") + if self.scratch_budget_bytes <= 0: + raise SchedulerError("scratch_budget_bytes must be positive") + if self.max_active_sessions < 1: + raise SchedulerError("max_active_sessions must be >= 1") + if self.max_queue_depth < 0: + raise SchedulerError("max_queue_depth must be >= 0") + if self.max_batch_size < 1: + raise SchedulerError("max_batch_size must be >= 1") + if self.max_prefill_tokens_per_tick < 1: + raise SchedulerError("max_prefill_tokens_per_tick must be >= 1") + + @property + def max_scratch_sessions(self) -> int: + """How many concurrent sessions the scratch envelope alone permits.""" + return self.scratch_budget_bytes // self.scratch_bytes_per_session + + @property + def effective_active_cap(self) -> int: + """The tighter of the active-slot cap and the scratch-derived cap.""" + return max(1, min(self.max_active_sessions, self.max_scratch_sessions)) + + +class AdmissionReason(str, Enum): + """Why a submission was admitted, queued, or rejected.""" + + ADMITTED = "admitted" + QUEUED = "queued" + REJECTED_QUEUE_FULL = "rejected-queue-full" + REJECTED_KV_BUDGET = "rejected-kv-budget" + REJECTED_SCRATCH_BUDGET = "rejected-scratch-budget" + REJECTED_DUPLICATE = "rejected-duplicate" + REJECTED_INVALID = "rejected-invalid" + + +# Reasons that mean "will run" (admitted now, or accepted into the bounded queue). +_ACCEPTED = frozenset({AdmissionReason.ADMITTED, AdmissionReason.QUEUED}) +# Reasons that mean "refused" — the caller must apply backpressure / retry later. +_REJECTED = frozenset( + { + AdmissionReason.REJECTED_QUEUE_FULL, + AdmissionReason.REJECTED_KV_BUDGET, + AdmissionReason.REJECTED_SCRATCH_BUDGET, + AdmissionReason.REJECTED_DUPLICATE, + AdmissionReason.REJECTED_INVALID, + } +) + + +@dataclass(frozen=True) +class AdmissionDecision: + """The structured outcome of :meth:`ContinuousBatchScheduler.submit`.""" + + session_id: str + reason: AdmissionReason + detail: str = "" + + @property + def accepted(self) -> bool: + return self.reason in _ACCEPTED + + @property + def running(self) -> bool: + return self.reason is AdmissionReason.ADMITTED + + @property + def rejected(self) -> bool: + return self.reason in _REJECTED + + def __str__(self) -> str: + suffix = f": {self.detail}" if self.detail else "" + return f"session {self.session_id} {self.reason.value}{suffix}" + + +# --------------------------------------------------------------------------- # +# Requests, engine contract, and per-session state. +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class GenerationRequest: + """One session's greedy generation job: a prompt and a token budget.""" + + session_id: str + route_epoch: int + prompt_token_ids: tuple[int, ...] + max_new_tokens: int + + def __post_init__(self) -> None: + if not isinstance(self.session_id, str) or not self.session_id.strip(): + raise SchedulerError("session_id must be a non-empty string") + if isinstance(self.route_epoch, bool) or not isinstance(self.route_epoch, int): + raise SchedulerError("route_epoch must be an integer") + if self.route_epoch < 0: + raise SchedulerError("route_epoch must be >= 0") + if not self.prompt_token_ids: + raise SchedulerError("prompt_token_ids must be non-empty") + if self.max_new_tokens < 1: + raise SchedulerError("max_new_tokens must be >= 1") + + @property + def prompt_len(self) -> int: + return len(self.prompt_token_ids) + + @property + def final_seq_len(self) -> int: + """Sequence length after the whole job completes (prompt + new tokens). + + The prefill emits the first new token, so the final KV length is + ``prompt_len + max_new_tokens - 1``. + """ + return self.prompt_len + self.max_new_tokens - 1 + + +@dataclass(frozen=True) +class DecodeItem: + """One member of a decode batch: which session decodes which input token.""" + + session_id: str + route_epoch: int + token_id: int + + +@dataclass(frozen=True) +class StepResult: + """The output of one prefill or one decode-batch member.""" + + session_id: str + route_epoch: int + token_id: int + seq_len: int + + +class Phase(str, Enum): + PENDING_PREFILL = "pending-prefill" + DECODING = "decoding" + DONE = "done" + + +class DoneReason(str, Enum): + COMPLETED = "completed" + CACHE_MISS = "cache-miss" + # DGR-013: a session can also leave the scheduler because the client cancelled + # it or because it failed (deadline/heartbeat loss, worker death, stream reset). + # These are distinguished so billing/work records never bill uncompleted work. + CANCELLED = "cancelled" + FAILED = "failed" + + +@dataclass +class SessionState: + """Live scheduler state for one admitted session (isolated per session).""" + + request: GenerationRequest + phase: Phase = Phase.PENDING_PREFILL + generated: list[int] = field(default_factory=list) + done_reason: DoneReason | None = None + cache_miss: CacheMiss | None = None + + @property + def session_id(self) -> str: + return self.request.session_id + + @property + def route_epoch(self) -> int: + return self.request.route_epoch + + @property + def remaining(self) -> int: + return self.request.max_new_tokens - len(self.generated) + + @property + def last_token(self) -> int: + return self.generated[-1] + + +class KvBatchEngine: + """Adapt a DGR-007 :class:`KvBoundaryAdapter` to the :class:`BatchEngine` contract. + + The adapter must wrap a *full* (head **and** tail) shard so a decode step + samples a token — a middle/head-only range emits a boundary bundle, which the + node-local scheduler does not turn into an output token. Multi-range routes + batch at the head node, whose adapter owns the final head. + + ``decode_batch`` runs each member through the adapter's cached decode. Each + session attends only over its own KV context, exactly as an independent + sequence would inside one native ``llama_decode`` batch; the pure-numpy engine + runs the members sequentially, while the pinned llama.cpp worker fuses them + into a single graph. The scheduling semantics — one batch per tick, isolated + positions and outputs — are identical, so this stands in for the native path + without a download or GPU. + """ + + def __init__(self, adapter: KvBoundaryAdapter) -> None: + if not (adapter.is_head and adapter.is_tail): + raise SchedulerError( + "KvBatchEngine requires a full (head+tail) shard so decode steps " + "sample tokens; got a partial range (head=%s tail=%s)" + % (adapter.is_head, adapter.is_tail) + ) + self._adapter = adapter + self._manager: HotKvStateManager = adapter.manager + + def recipe_fingerprint(self) -> str: + return self._adapter.recipe.fingerprint() + + def prefill( + self, session_id: str, route_epoch: int, token_ids: Sequence[int] + ) -> StepResult: + out = self._adapter.prefill(session_id, route_epoch, token_ids=list(token_ids)) + seq_len = self._manager.get(session_id, route_epoch).seq_len + return StepResult(session_id, route_epoch, int(out.token_id), seq_len) + + def decode_batch( + self, items: Sequence[DecodeItem] + ) -> list[StepResult | CacheMiss]: + results: list[StepResult | CacheMiss] = [] + for item in items: + out = self._adapter.decode( + item.session_id, item.route_epoch, token_ids=[item.token_id] + ) + if isinstance(out, CacheMiss): + results.append(out) + continue + seq_len = self._manager.get(item.session_id, item.route_epoch).seq_len + results.append( + StepResult(item.session_id, item.route_epoch, int(out.token_id), seq_len) + ) + return results + + def release(self, session_id: str, route_epoch: int) -> None: + self._manager.release(session_id, route_epoch) + + +# --------------------------------------------------------------------------- # +# Telemetry. +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class SchedulerTelemetry: + """A bounded, JSON-safe snapshot of node scheduling pressure. + + These are the capability signals a node advertises: enough to decide whether + it can take more work, and to spot saturation, without exposing session + contents. + """ + + active_sessions: int + queue_depth: int + batch_occupancy_last: int + batch_occupancy_avg: float + batch_occupancy_max: int + weight_bytes: int + kv_total_bytes: int + kv_budget_bytes: int + kv_pressure: float + scratch_used_bytes: int + scratch_budget_bytes: int + scratch_pressure: float + prefill_tokens_total: int + decode_tokens_total: int + prefill_tokens_per_sec: float + decode_tokens_per_sec: float + rejected_admissions_total: int + rejected_by_reason: Mapping[str, int] + completed_sessions: int + cancelled_sessions: int + failed_sessions: int + ticks: int + + def to_dict(self) -> dict: + return { + "active_sessions": self.active_sessions, + "queue_depth": self.queue_depth, + "batch_occupancy_last": self.batch_occupancy_last, + "batch_occupancy_avg": round(self.batch_occupancy_avg, 4), + "batch_occupancy_max": self.batch_occupancy_max, + "weight_bytes": self.weight_bytes, + "kv_total_bytes": self.kv_total_bytes, + "kv_budget_bytes": self.kv_budget_bytes, + "kv_pressure": round(self.kv_pressure, 4), + "scratch_used_bytes": self.scratch_used_bytes, + "scratch_budget_bytes": self.scratch_budget_bytes, + "scratch_pressure": round(self.scratch_pressure, 4), + "prefill_tokens_total": self.prefill_tokens_total, + "decode_tokens_total": self.decode_tokens_total, + "prefill_tokens_per_sec": round(self.prefill_tokens_per_sec, 4), + "decode_tokens_per_sec": round(self.decode_tokens_per_sec, 4), + "rejected_admissions_total": self.rejected_admissions_total, + "rejected_by_reason": dict(self.rejected_by_reason), + "completed_sessions": self.completed_sessions, + "cancelled_sessions": self.cancelled_sessions, + "failed_sessions": self.failed_sessions, + "ticks": self.ticks, + } + + +# --------------------------------------------------------------------------- # +# The scheduler. +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class TickReport: + """What one :meth:`ContinuousBatchScheduler.run_tick` did (for observability).""" + + prefilled: tuple[str, ...] + decoded: tuple[str, ...] + batch_occupancy: int + completed: tuple[str, ...] + admitted_from_queue: tuple[str, ...] + + @property + def did_work(self) -> bool: + return bool(self.prefilled or self.decoded) + + +class ContinuousBatchScheduler: + """Node-local continuous-batching scheduler with bounded admission. + + Fixed scheduling policy per :meth:`run_tick`: + + 1. Promote queued sessions into free active slots (respecting the active and + scratch caps). + 2. **Decode first:** form one batch from every active decoding session (up to + ``max_batch_size``) and run it once. This is what guarantees prefill cannot + starve decode. + 3. **Then bounded prefill:** prefill pending sessions until the per-tick prefill + token budget is spent (always allowing at least one, so a single large + prompt still makes progress). + 4. Reap completed/lost sessions, releasing their KV so budget returns. + + The scheduler is thread-safe (an ``RLock`` guards all state) so a real server + can call :meth:`submit` from request threads while a worker thread drives + :meth:`run_tick`; the deterministic tests drive both from one thread. + """ + + def __init__( + self, + engine: Any, + budget: NodeBudget | None = None, + *, + clock: Callable[[], float] | None = None, + ) -> None: + self._engine = engine + self._budget = budget or NodeBudget() + self._clock = clock or time.monotonic + self._fingerprint = str(engine.recipe_fingerprint()) + + self._active: dict[str, SessionState] = {} + self._queue: "deque[GenerationRequest]" = deque() + self._queued_ids: set[str] = set() + self._done: dict[str, SessionState] = {} + + # Telemetry counters. + self._started = self._clock() + self._ticks = 0 + self._prefill_tokens = 0 + self._decode_tokens = 0 + self._batch_occupancy_last = 0 + self._batch_occupancy_max = 0 + self._batch_sum = 0 + self._batch_count = 0 + self._completed = 0 + self._cancelled = 0 + self._failed = 0 + self._rejected = 0 + self._rejected_by_reason: dict[str, int] = {} + + self._lock = threading.RLock() + + # -- admission ------------------------------------------------------------ + + def submit(self, request: GenerationRequest) -> AdmissionDecision: + """Admit, queue, or reject one generation request (bounded admission). + + Order of checks: identity (duplicate) → hard feasibility (KV, scratch) → + capacity (free active slot vs bounded queue vs full). A full queue yields + :attr:`AdmissionReason.REJECTED_QUEUE_FULL`, the explicit backpressure + signal. + """ + with self._lock: + sid = request.session_id + if sid in self._active or sid in self._queued_ids: + return self._reject( + request, AdmissionReason.REJECTED_DUPLICATE, "already scheduled" + ) + + # Hard feasibility: a single session must be able to fit KV + scratch + # on its own; otherwise it can never run and is rejected up front + # rather than wedging the queue. + kv_need = self._kv_bytes_for(request) + if kv_need > self._budget.kv_budget_bytes: + return self._reject( + request, + AdmissionReason.REJECTED_KV_BUDGET, + f"needs {kv_need} KV bytes > budget " + f"{self._budget.kv_budget_bytes}", + ) + if self._budget.scratch_bytes_per_session > self._budget.scratch_budget_bytes: + return self._reject( + request, + AdmissionReason.REJECTED_SCRATCH_BUDGET, + "per-session scratch exceeds the scratch budget", + ) + + if self._has_capacity_locked(): + self._activate_locked(request) + return AdmissionDecision(sid, AdmissionReason.ADMITTED) + + if len(self._queue) < self._budget.max_queue_depth: + self._queue.append(request) + self._queued_ids.add(sid) + return AdmissionDecision(sid, AdmissionReason.QUEUED) + + return self._reject( + request, + AdmissionReason.REJECTED_QUEUE_FULL, + f"queue full at depth {self._budget.max_queue_depth}", + ) + + # -- cancellation / failure (DGR-013) ------------------------------------- + + def cancel( + self, + session_id: str, + *, + reason: DoneReason = DoneReason.CANCELLED, + detail: str = "", + ) -> bool: + """Remove a session from the scheduler, releasing its KV and queue slot. + + Cancellation is bounded and explicit: if the session is *queued* it is + dropped from the bounded queue (its queued buffer is released); if it is + *active* its KV is released through the engine and it is moved to the done + set with a non-completed :class:`DoneReason` so billing/work records never + count it as completed work. Returns ``True`` if a live (queued or active) + session was found. Idempotent: cancelling an unknown or already-finished + session returns ``False`` and mutates nothing. + + ``reason`` must be a terminal non-completed reason (``CANCELLED`` for an + explicit client cancel, ``FAILED`` for deadline/heartbeat/worker loss). + """ + if reason not in (DoneReason.CANCELLED, DoneReason.FAILED): + raise SchedulerError( + "cancel reason must be CANCELLED or FAILED, not %r" % (reason,) + ) + with self._lock: + # Queued but not yet running: drop it from the bounded queue so the + # backpressure boundary recovers and no execution slot is ever taken. + if session_id in self._queued_ids: + self._queued_ids.discard(session_id) + dropped = next( + (r for r in self._queue if r.session_id == session_id), None + ) + self._queue = deque( + r for r in self._queue if r.session_id != session_id + ) + self._finalize_cancelled_locked(session_id, reason, dropped) + return True + + state = self._active.get(session_id) + if state is None: + return False + # Active: release the KV context on this shard, then record the + # terminal reason. release() is idempotent, so a concurrent reap or a + # prior cache-miss release cannot double-free. + self._engine.release(state.session_id, state.route_epoch) + del self._active[session_id] + state.phase = Phase.DONE + state.done_reason = reason + self._done[session_id] = state + self._count_terminal_locked(reason) + return True + + def _finalize_cancelled_locked( + self, + session_id: str, + reason: DoneReason, + request: GenerationRequest | None, + ) -> None: + # A queued session has no live KV and no committed tokens yet; record a + # terminal state (with its original request when known) so results() and + # telemetry account for it distinctly from completed work. + if request is not None: + state = SessionState( + request=request, phase=Phase.DONE, done_reason=reason + ) + self._done[session_id] = state + self._count_terminal_locked(reason) + + def _count_terminal_locked(self, reason: DoneReason) -> None: + if reason is DoneReason.CANCELLED: + self._cancelled += 1 + elif reason is DoneReason.FAILED: + self._failed += 1 + + # -- scheduling ----------------------------------------------------------- + + def run_tick(self) -> TickReport: + """Run one scheduling step: admit, decode-batch, bounded-prefill, reap.""" + with self._lock: + self._ticks += 1 + admitted = self._admit_from_queue_locked() + decoded, occupancy = self._run_decode_batch_locked() + prefilled = self._run_prefill_locked() + completed = self._reap_locked() + # A reap frees slots; pull more work forward so the next caller sees a + # full node rather than an artificially idle one. + admitted = admitted + self._admit_from_queue_locked() + return TickReport( + prefilled=tuple(prefilled), + decoded=tuple(decoded), + batch_occupancy=occupancy, + completed=tuple(completed), + admitted_from_queue=tuple(admitted), + ) + + def run_to_completion(self, *, max_ticks: int | None = None) -> dict[str, list[int]]: + """Drive ticks until every submitted session finishes; return outputs. + + Returns ``{session_id: generated_token_ids}`` for every session that ran. + ``max_ticks`` is a safety bound; exceeding it raises rather than looping + forever on a misconfiguration. + """ + limit = max_ticks if max_ticks is not None else self._default_tick_limit() + for _ in range(limit): + with self._lock: + if not self._active and not self._queue: + break + self.run_tick() + else: + with self._lock: + pending = len(self._active) + len(self._queue) + if pending: + raise SchedulerError( + f"run_to_completion exceeded {limit} ticks with {pending} " + "sessions still pending; check budgets and token counts" + ) + with self._lock: + return {sid: list(s.generated) for sid, s in self._done.items()} + + # -- results -------------------------------------------------------------- + + def outputs(self) -> dict[str, list[int]]: + """Generated tokens for every completed session so far.""" + with self._lock: + return {sid: list(s.generated) for sid, s in self._done.items()} + + def session_result(self, session_id: str) -> SessionState | None: + with self._lock: + return self._done.get(session_id) or self._active.get(session_id) + + # -- telemetry ------------------------------------------------------------ + + def telemetry(self, *, now: float | None = None) -> SchedulerTelemetry: + """Capability snapshot: sessions, queue, batch, KV/scratch pressure, rates.""" + with self._lock: + observed = self._clock() if now is None else now + elapsed = max(observed - self._started, 1e-9) + kv_total = self._engine_kv_bytes() + kv_budget = self._budget.kv_budget_bytes + scratch_used = len(self._active) * self._budget.scratch_bytes_per_session + scratch_budget = self._budget.scratch_budget_bytes + avg_occupancy = ( + self._batch_sum / self._batch_count if self._batch_count else 0.0 + ) + return SchedulerTelemetry( + active_sessions=len(self._active), + queue_depth=len(self._queue), + batch_occupancy_last=self._batch_occupancy_last, + batch_occupancy_avg=avg_occupancy, + batch_occupancy_max=self._batch_occupancy_max, + weight_bytes=self._budget.weight_bytes, + kv_total_bytes=kv_total, + kv_budget_bytes=kv_budget, + kv_pressure=kv_total / kv_budget if kv_budget else 0.0, + scratch_used_bytes=scratch_used, + scratch_budget_bytes=scratch_budget, + scratch_pressure=scratch_used / scratch_budget if scratch_budget else 0.0, + prefill_tokens_total=self._prefill_tokens, + decode_tokens_total=self._decode_tokens, + prefill_tokens_per_sec=self._prefill_tokens / elapsed, + decode_tokens_per_sec=self._decode_tokens / elapsed, + rejected_admissions_total=self._rejected, + rejected_by_reason=dict(self._rejected_by_reason), + completed_sessions=self._completed, + cancelled_sessions=self._cancelled, + failed_sessions=self._failed, + ticks=self._ticks, + ) + + # -- internals ------------------------------------------------------------ + + def _reject( + self, request: GenerationRequest, reason: AdmissionReason, detail: str + ) -> AdmissionDecision: + self._rejected += 1 + self._rejected_by_reason[reason.value] = ( + self._rejected_by_reason.get(reason.value, 0) + 1 + ) + return AdmissionDecision(request.session_id, reason, detail) + + def _kv_bytes_for(self, request: GenerationRequest) -> int: + # bytes_per_token is defined by the loaded shard's KV recipe; the whole + # generation occupies prompt + (new-1) positions at its peak. + per_token = self._manager().recipe.bytes_per_token() + return request.final_seq_len * per_token + + def _manager(self) -> HotKvStateManager: + manager = getattr(self._engine, "_manager", None) + if manager is None: + raise SchedulerError( + "engine does not expose a Hot KV State manager for budget accounting" + ) + return manager + + def _engine_kv_bytes(self) -> int: + manager = getattr(self._engine, "_manager", None) + return int(manager.total_bytes) if manager is not None else 0 + + def _has_capacity_locked(self) -> bool: + return len(self._active) < self._budget.effective_active_cap + + def _activate_locked(self, request: GenerationRequest) -> None: + if self._fingerprint != str(self._engine.recipe_fingerprint()): + # The loaded shard's recipe must not change under the scheduler. + raise SchedulerError("engine recipe fingerprint changed mid-flight") + self._active[request.session_id] = SessionState(request=request) + + def _admit_from_queue_locked(self) -> list[str]: + admitted: list[str] = [] + while self._queue and self._has_capacity_locked(): + request = self._queue.popleft() + self._queued_ids.discard(request.session_id) + self._activate_locked(request) + admitted.append(request.session_id) + return admitted + + def _run_decode_batch_locked(self) -> tuple[list[str], int]: + decoding = [ + s for s in self._active.values() if s.phase is Phase.DECODING + ] + if not decoding: + self._batch_occupancy_last = 0 + return [], 0 + batch = decoding[: self._budget.max_batch_size] + items = [ + DecodeItem(s.session_id, s.route_epoch, s.last_token) for s in batch + ] + results = self._engine.decode_batch(items) + if len(results) != len(batch): + raise SchedulerError( + "engine returned %d results for a batch of %d" + % (len(results), len(batch)) + ) + decoded: list[str] = [] + for state, result in zip(batch, results): + if isinstance(result, CacheMiss): + state.phase = Phase.DONE + state.done_reason = DoneReason.CACHE_MISS + state.cache_miss = result + continue + state.generated.append(result.token_id) + self._decode_tokens += 1 + decoded.append(state.session_id) + if state.remaining <= 0: + state.phase = Phase.DONE + state.done_reason = DoneReason.COMPLETED + occupancy = len(batch) + self._batch_occupancy_last = occupancy + self._batch_occupancy_max = max(self._batch_occupancy_max, occupancy) + self._batch_sum += occupancy + self._batch_count += 1 + return decoded, occupancy + + def _run_prefill_locked(self) -> list[str]: + pending = [ + s for s in self._active.values() if s.phase is Phase.PENDING_PREFILL + ] + prefilled: list[str] = [] + spent = 0 + for state in pending: + # Always allow the first prefill of the tick (progress guarantee), + # then honour the per-tick token budget so prefill can't monopolise. + if prefilled and spent + state.request.prompt_len > self._budget.max_prefill_tokens_per_tick: + break + result = self._engine.prefill( + state.session_id, + state.route_epoch, + state.request.prompt_token_ids, + ) + state.generated.append(result.token_id) + self._prefill_tokens += state.request.prompt_len + spent += state.request.prompt_len + prefilled.append(state.session_id) + if state.remaining <= 0: + state.phase = Phase.DONE + state.done_reason = DoneReason.COMPLETED + else: + state.phase = Phase.DECODING + return prefilled + + def _reap_locked(self) -> list[str]: + completed: list[str] = [] + for sid, state in list(self._active.items()): + if state.phase is not Phase.DONE: + continue + self._engine.release(state.session_id, state.route_epoch) + del self._active[sid] + self._done[sid] = state + if state.done_reason is DoneReason.COMPLETED: + self._completed += 1 + completed.append(sid) + return completed + + def _default_tick_limit(self) -> int: + # Generous upper bound: worst case is fully serialized (one session at a + # time, one token per tick) plus slack for admission ticks. + pending_tokens = sum( + s.request.max_new_tokens for s in self._active.values() + ) + sum(r.max_new_tokens for r in self._queue) + return 8 * (pending_tokens + len(self._active) + len(self._queue) + 1) + + +# --------------------------------------------------------------------------- # +# Concurrency 1/2/4/8 sweep (deterministic saturation report). +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class ConcurrencyResult: + """One concurrency level's deterministic scheduling result.""" + + concurrency: int + ticks: int + decode_batches: int + decode_tokens: int + prefill_tokens: int + avg_batch_occupancy: float + max_batch_occupancy: int + tokens_per_tick: float + peak_kv_bytes: int + rejected_admissions: int + cache_misses: int + + def to_dict(self) -> dict: + return { + "concurrency": self.concurrency, + "ticks": self.ticks, + "decode_batches": self.decode_batches, + "decode_tokens": self.decode_tokens, + "prefill_tokens": self.prefill_tokens, + "avg_batch_occupancy": round(self.avg_batch_occupancy, 4), + "max_batch_occupancy": self.max_batch_occupancy, + "tokens_per_tick": round(self.tokens_per_tick, 4), + "peak_kv_bytes": self.peak_kv_bytes, + "rejected_admissions": self.rejected_admissions, + "cache_misses": self.cache_misses, + } + + +@dataclass(frozen=True) +class ConcurrencySweep: + """The full 1/2/4/8 report plus the derived saturation point.""" + + results: tuple[ConcurrencyResult, ...] + saturation_concurrency: int + corruption_free: bool + reference_outputs: Mapping[str, tuple[int, ...]] + + def to_dict(self) -> dict: + return { + "schema_version": 1, + "results": [r.to_dict() for r in self.results], + "saturation_concurrency": self.saturation_concurrency, + "corruption_free": self.corruption_free, + "reference_outputs": { + sid: list(tokens) for sid, tokens in self.reference_outputs.items() + }, + } + + +def run_concurrency_sweep( + engine_factory: Callable[[], Any], + requests: Iterable[GenerationRequest], + *, + concurrency_levels: Sequence[int] = (1, 2, 4, 8), + budget_factory: Callable[[int], NodeBudget] | None = None, + saturation_tolerance: float = 1e-9, +) -> ConcurrencySweep: + """Run the same jobs at each concurrency level and report saturation. + + For every level, a fresh engine (fresh KV manager) runs all ``requests`` with + ``max_active_sessions`` and ``max_batch_size`` capped to that level. The + concurrency-1 run is the serialized reference; every higher level must produce + the **byte-identical** per-session token stream (greedy sampling over isolated + KV is order-independent), which is the "no cross-session corruption" proof. + + Saturation is the smallest level at which average batch occupancy stops rising + (more slots no longer pack more sessions per batch) — i.e. the node is fully + utilised and adding concurrency yields no further batching gain for this load. + """ + requests = list(requests) + if not requests: + raise SchedulerError("run_concurrency_sweep needs at least one request") + levels = sorted({int(level) for level in concurrency_levels}) + if any(level < 1 for level in levels): + raise SchedulerError("concurrency levels must be >= 1") + + def default_budget(level: int) -> NodeBudget: + # Budgets sized so the load never evicts: correctness of the sweep must not + # depend on eviction. KV holds every session's whole generation at once. + engine = engine_factory() + per_token = getattr(engine, "_manager").recipe.bytes_per_token() + total_kv = sum(r.final_seq_len for r in requests) * per_token + return NodeBudget( + kv_budget_bytes=max(total_kv, per_token), + scratch_bytes_per_session=1, + scratch_budget_bytes=max(1, level), + max_active_sessions=level, + max_queue_depth=len(requests), + max_batch_size=level, + max_prefill_tokens_per_tick=max(r.prompt_len for r in requests), + ) + + budget_for = budget_factory or default_budget + results: list[ConcurrencyResult] = [] + reference: dict[str, tuple[int, ...]] | None = None + corruption_free = True + + for level in levels: + engine = engine_factory() + scheduler = ContinuousBatchScheduler(engine, budget_for(level)) + cache_misses = 0 + peak_kv = 0 + decode_batches = 0 + for request in requests: + decision = scheduler.submit(request) + if not decision.accepted: + raise SchedulerError( + f"sweep request {request.session_id} was rejected at " + f"concurrency {level}: {decision}" + ) + # Drive ticks manually so we can sample peak KV and count decode batches. + limit = scheduler._default_tick_limit() + for _ in range(limit): + if not scheduler._active and not scheduler._queue: + break + report = scheduler.run_tick() + if report.batch_occupancy > 0: + decode_batches += 1 + peak_kv = max(peak_kv, scheduler.telemetry().kv_total_bytes) + outputs = {sid: tuple(tokens) for sid, tokens in scheduler.outputs().items()} + for state in ( + scheduler.session_result(r.session_id) for r in requests + ): + if state is not None and state.done_reason is DoneReason.CACHE_MISS: + cache_misses += 1 + + if reference is None: + reference = outputs + elif outputs != reference: + corruption_free = False + + telem = scheduler.telemetry() + results.append( + ConcurrencyResult( + concurrency=level, + ticks=telem.ticks, + decode_batches=decode_batches, + decode_tokens=telem.decode_tokens_total, + prefill_tokens=telem.prefill_tokens_total, + avg_batch_occupancy=telem.batch_occupancy_avg, + max_batch_occupancy=telem.batch_occupancy_max, + tokens_per_tick=(telem.decode_tokens_total + telem.prefill_tokens_total) + / max(1, telem.ticks), + peak_kv_bytes=peak_kv, + rejected_admissions=telem.rejected_admissions_total, + cache_misses=cache_misses, + ) + ) + + saturation = _saturation_point(results, saturation_tolerance) + assert reference is not None + return ConcurrencySweep( + results=tuple(results), + saturation_concurrency=saturation, + corruption_free=corruption_free, + reference_outputs=reference, + ) + + +def _saturation_point( + results: Sequence[ConcurrencyResult], tolerance: float +) -> int: + """Smallest concurrency where average batch occupancy stops increasing.""" + if not results: + return 0 + best = results[0] + for current in results[1:]: + if current.avg_batch_occupancy <= best.avg_batch_occupancy + tolerance: + return best.concurrency + best = current + return results[-1].concurrency diff --git a/packages/node/meshnet_node/failure_semantics.py b/packages/node/meshnet_node/failure_semantics.py new file mode 100644 index 0000000..8419441 --- /dev/null +++ b/packages/node/meshnet_node/failure_semantics.py @@ -0,0 +1,893 @@ +"""Bounded failure, cancellation, and restart semantics for Shard streams (DGR-013). + +Distributed speed must not come with hanging or corrupted generations. This module +hardens the per-Route-Session decode stream that runs over the DGR-007 Hot KV State +manager (isolated ``(session, epoch)`` KV) and the DGR-012 continuous-batch +scheduler. It is deliberately backend-agnostic and pure-Python: it drives the same +``KvBoundaryAdapter`` the default deterministic gate uses, so the whole matrix stays +download-free, GPU-free, and API-credit-free while exercising the *real* KV +isolation path (the pinned llama.cpp worker, DGR-008, implements the identical +adapter contract). + +The guarantees, mapped to the story's acceptance criteria: + +* **Deadlines and heartbeat/health loss terminate blocked stream operations.** + :class:`DeadlineGuard` bounds every step against an absolute deadline and a + heartbeat-timeout; when either is breached it raises :class:`StreamTerminated` + so a blocked stream never hangs. +* **Cancellation propagates across every Shard and releases local KV and queued + buffers.** :class:`ShardCancellationGroup` fans a single cancel across every + node-local KV manager serving a Route Session and releases queued activation + buffers; the DGR-012 scheduler's :meth:`~meshnet_node.batch_scheduler. + ContinuousBatchScheduler.cancel` drops queued/active work on this node. +* **Duplicate steps are idempotent; uncertain mutations are never replayed + silently.** :class:`IdempotencyLedger` records each committed + ``(session, epoch, step)`` and returns the recorded token for a duplicate + delivery instead of re-running it. A step whose outcome is *uncertain* (the + worker died mid-mutation) is marked uncertain and can never be silently + replayed — a replay attempt raises :class:`UncertainMutationError`, forcing an + explicit verify-or-restart. +* **Alpha failover restarts from token zero on a newly compatible route rather + than importing unverified KV.** :class:`RestartController` opens a *new* route + epoch, releases every shard's prior-epoch KV, and the restart re-prefills the + whole prompt from token zero. The old epoch becomes stale (rejected by the KV + manager); unverified KV is never migrated (RALPH runtime decision #14). +* **Billing/work records distinguish completed, cancelled, failed, and unverified + work.** :class:`WorkLedger` records a typed :class:`WorkRecord` per attempt; + only :attr:`WorkStatus.COMPLETED` records are billable, so cancelled, failed, + and uncertain (unverified) work is accounted but never charged. + +:class:`HardenedSessionRunner` composes these into one drivable stream: it runs a +single session's prefill+decode through the adapter under a deadline/heartbeat +guard and a cancellation token, records the typed work outcome, and — via +:meth:`HardenedSessionRunner.run_with_failover` — restarts a transient failure +from token zero on a fresh epoch. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Any, Callable, Mapping, Sequence + +from meshnet_node.batch_scheduler import DoneReason, GenerationRequest +from meshnet_node.boundary_adapter import BoundaryContractError, TailOutput +from meshnet_node.hot_kv_state import ( + CacheMiss, + HotKvStateManager, + IncompatibleCacheRecipeError, + KvBoundaryAdapter, + KvCacheMissError, + StaleRouteEpochError, +) + + +class FailureSemanticsError(RuntimeError): + """Base class for failure/cancellation/restart errors.""" + + +# --------------------------------------------------------------------------- # +# Typed outcomes: failure kinds and billing/work statuses. +# --------------------------------------------------------------------------- # + + +class FailureKind(str, Enum): + """Why a stream step failed. Stable strings for the protocol's structured status.""" + + # Bounded termination of a blocked op. + DEADLINE_EXCEEDED = "deadline-exceeded" + HEARTBEAT_LOST = "heartbeat-lost" + # Transport / worker loss (transient — a restart from token zero may succeed). + WORKER_DEATH = "worker-death" + STREAM_RESET = "stream-reset" + # Protocol violations (deterministic — a restart would fail identically). + MALFORMED_BUNDLE = "malformed-bundle" + STALE_EPOCH = "stale-epoch" + INCOMPATIBLE_RECIPE = "incompatible-recipe" + # KV state expected by the caller is gone; re-prefill from token zero. + CACHE_MISS = "cache-miss" + # Explicit client cancellation. + CANCELLED = "cancelled" + + +# Failure kinds that a from-token-zero restart on a fresh route may recover from. +# A protocol violation or an explicit bound (deadline/cancel) is NOT restartable — +# retrying it would hang or fail identically, so we surface it instead. +_RESTARTABLE = frozenset( + { + FailureKind.WORKER_DEATH, + FailureKind.STREAM_RESET, + FailureKind.CACHE_MISS, + } +) + +# Failure kinds whose mutation outcome is *uncertain* — the KV may or may not have +# advanced, so the confirmed work is billed as UNVERIFIED and never replayed +# silently. Only an *unexpected* error raised while a step was executing is +# uncertain (mapped to WORKER_DEATH). A stream reset, deadline, or cache miss +# detected at a step boundary is certain: nothing committed for that step. +_UNCERTAIN = frozenset({FailureKind.WORKER_DEATH}) + + +class WorkStatus(str, Enum): + """The billing-relevant outcome class of a unit of work (AC: billing records). + + Only :attr:`COMPLETED` work is billable. Cancelled, failed, and unverified + work is recorded distinctly so a client is never charged for a generation that + hung, was cancelled, or whose mutations could not be verified. + """ + + COMPLETED = "completed" + CANCELLED = "cancelled" + FAILED = "failed" + UNVERIFIED = "unverified" + + +def work_status_for(kind: FailureKind) -> WorkStatus: + """Map a terminal failure kind to its billing/work status.""" + if kind is FailureKind.CANCELLED: + return WorkStatus.CANCELLED + if kind in _UNCERTAIN: + return WorkStatus.UNVERIFIED + return WorkStatus.FAILED + + +def classify_exception(exc: BaseException) -> FailureKind: + """Classify a raised error into a :class:`FailureKind`. + + Protocol violations map to their specific kind; a :class:`StreamTerminated` + carries its own kind; any *unexpected* error is treated as worker death + (an uncertain, transient loss), never silently ignored. + """ + if isinstance(exc, StreamTerminated): + return exc.kind + if isinstance(exc, OperationCancelled): + return FailureKind.CANCELLED + if isinstance(exc, StaleRouteEpochError): + return FailureKind.STALE_EPOCH + if isinstance(exc, IncompatibleCacheRecipeError): + return FailureKind.INCOMPATIBLE_RECIPE + if isinstance(exc, BoundaryContractError): + return FailureKind.MALFORMED_BUNDLE + if isinstance(exc, KvCacheMissError): + return FailureKind.CACHE_MISS + return FailureKind.WORKER_DEATH + + +# --------------------------------------------------------------------------- # +# Deadlines and heartbeat/health loss. +# --------------------------------------------------------------------------- # + + +class StreamTerminated(FailureSemanticsError): + """A blocked stream op was terminated by a deadline or heartbeat/health loss.""" + + def __init__(self, kind: FailureKind, detail: str = "") -> None: + self.kind = kind + self.detail = detail + suffix = f": {detail}" if detail else "" + super().__init__(f"stream terminated ({kind.value}){suffix}") + + +class OperationCancelled(FailureSemanticsError): + """Raised when a step observes its :class:`CancellationToken` is cancelled.""" + + def __init__(self, reason: str = "client-cancel") -> None: + self.reason = reason + super().__init__(f"operation cancelled: {reason}") + + +@dataclass +class DeadlineGuard: + """Bounds a blocked stream op against an absolute deadline and heartbeat loss. + + ``deadline`` is an absolute time on ``clock``'s scale (``None`` disables it). + ``heartbeat_timeout`` is the maximum tolerated gap since the last observed + heartbeat; when the peer stops sending heartbeats (its health is lost) the gap + grows past the timeout and :meth:`check` raises rather than blocking forever. + Both bounds are checked with an injected ``clock`` so the matrix is + deterministic. + """ + + deadline: float | None = None + heartbeat_timeout: float | None = None + clock: Callable[[], float] = time.monotonic + _last_heartbeat: float = field(default=0.0, init=False) + _started: bool = field(default=False, init=False) + + def __post_init__(self) -> None: + if self.heartbeat_timeout is not None and self.heartbeat_timeout <= 0: + raise FailureSemanticsError("heartbeat_timeout must be positive") + + def start(self) -> None: + self._last_heartbeat = self.clock() + self._started = True + + def heartbeat(self) -> None: + """Record that the peer is alive (resets the heartbeat gap).""" + self._last_heartbeat = self.clock() + + def check(self) -> None: + """Raise :class:`StreamTerminated` if the deadline or heartbeat lapsed.""" + if not self._started: + self.start() + now = self.clock() + if self.deadline is not None and now >= self.deadline: + raise StreamTerminated( + FailureKind.DEADLINE_EXCEEDED, + f"deadline {self.deadline} reached at {now}", + ) + if self.heartbeat_timeout is not None: + gap = now - self._last_heartbeat + if gap > self.heartbeat_timeout: + raise StreamTerminated( + FailureKind.HEARTBEAT_LOST, + f"no heartbeat for {gap} > {self.heartbeat_timeout}", + ) + + def remaining(self) -> float | None: + if self.deadline is None: + return None + return self.deadline - self.clock() + + +# --------------------------------------------------------------------------- # +# Cancellation that propagates across shards and releases KV + queued buffers. +# --------------------------------------------------------------------------- # + + +class CancellationToken: + """A thread-safe one-shot cancellation flag shared by a Route Session's steps.""" + + def __init__(self) -> None: + self._cancelled = False + self._reason = "" + self._lock = threading.Lock() + + def cancel(self, reason: str = "client-cancel") -> None: + with self._lock: + if not self._cancelled: + self._cancelled = True + self._reason = reason + + @property + def cancelled(self) -> bool: + with self._lock: + return self._cancelled + + @property + def reason(self) -> str: + with self._lock: + return self._reason + + def raise_if_cancelled(self) -> None: + with self._lock: + if self._cancelled: + raise OperationCancelled(self._reason) + + +@dataclass(frozen=True) +class CancellationOutcome: + """What a :meth:`ShardCancellationGroup.cancel` released (for observability).""" + + session_id: str + route_epoch: int + shards_released: int + buffers_released: int + + def to_dict(self) -> dict: + return { + "session_id": self.session_id, + "route_epoch": self.route_epoch, + "shards_released": self.shards_released, + "buffers_released": self.buffers_released, + } + + +class ShardCancellationGroup: + """Fan one cancellation across every node-local Shard of a Route Session. + + A Route Session spans a chain of Shards, each with its own local Hot KV State + manager (KV is never migrated between nodes). Cancelling the session must free + *all* of that state: this group releases the ``(session, epoch)`` KV on every + registered manager and invokes every registered queued-buffer release callback + (the pending activation bundles a node holds for the session). Release is + idempotent, so cancelling twice is safe. + """ + + def __init__(self, session_id: str, route_epoch: int) -> None: + if not isinstance(session_id, str) or not session_id.strip(): + raise FailureSemanticsError("session_id must be a non-empty string") + self.session_id = session_id + self.route_epoch = int(route_epoch) + self._managers: list[HotKvStateManager] = [] + self._buffers: list[Callable[[], None]] = [] + self._lock = threading.Lock() + self._cancelled = False + + def add_shard(self, manager: HotKvStateManager) -> "ShardCancellationGroup": + with self._lock: + self._managers.append(manager) + return self + + def add_queued_buffer( + self, release: Callable[[], None] + ) -> "ShardCancellationGroup": + """Register a queued activation buffer's release callback.""" + with self._lock: + self._buffers.append(release) + return self + + @property + def cancelled(self) -> bool: + with self._lock: + return self._cancelled + + def cancel(self) -> CancellationOutcome: + """Release every shard's KV and every queued buffer for this session.""" + with self._lock: + managers = list(self._managers) + buffers = list(self._buffers) + self._buffers.clear() + self._cancelled = True + shards_released = 0 + for manager in managers: + if manager.release(self.session_id, self.route_epoch): + shards_released += 1 + buffers_released = 0 + for release in buffers: + release() + buffers_released += 1 + return CancellationOutcome( + session_id=self.session_id, + route_epoch=self.route_epoch, + shards_released=shards_released, + buffers_released=buffers_released, + ) + + +# --------------------------------------------------------------------------- # +# Idempotency: duplicate steps are no-ops; uncertain mutations never replay. +# --------------------------------------------------------------------------- # + + +class StepPhase(str, Enum): + IN_FLIGHT = "in-flight" + COMMITTED = "committed" + UNCERTAIN = "uncertain" + + +class UncertainMutationError(FailureSemanticsError): + """Raised when a caller tries to replay a step whose outcome is uncertain. + + A step is uncertain when its mutation may or may not have been applied (worker + death / stream reset mid-append). Replaying it silently could double-apply KV + or bill unverified work, so the ledger refuses: the caller must verify against + the actual KV length or restart from token zero on a fresh epoch instead. + """ + + +@dataclass(frozen=True) +class StepKey: + """Identity of one idempotent stream step within a route epoch.""" + + session_id: str + route_epoch: int + step_index: int + + +@dataclass(frozen=True) +class StepDisposition: + """What :meth:`IdempotencyLedger.begin` decided for a step.""" + + fresh: bool + token: int | None = None + + @property + def duplicate(self) -> bool: + return not self.fresh + + +class IdempotencyLedger: + """Records committed/uncertain stream steps so duplicates never re-mutate. + + Keyed by ``(session, epoch, step_index)`` — the protocol's idempotency step. + + * :meth:`begin` on a *fresh* key marks it in-flight and returns "execute". + * :meth:`begin` on a *committed* key returns the recorded token so a duplicate + delivery is a no-op (idempotent replay). + * :meth:`begin` on an *in-flight* or *uncertain* key raises + :class:`UncertainMutationError` — a concurrent duplicate or a replay of an + unverified mutation is never silently applied. + """ + + def __init__(self) -> None: + self._phase: dict[StepKey, StepPhase] = {} + self._token: dict[StepKey, int] = {} + self._lock = threading.Lock() + + def begin(self, key: StepKey) -> StepDisposition: + with self._lock: + phase = self._phase.get(key) + if phase is None: + self._phase[key] = StepPhase.IN_FLIGHT + return StepDisposition(fresh=True) + if phase is StepPhase.COMMITTED: + return StepDisposition(fresh=False, token=self._token[key]) + # IN_FLIGHT (concurrent duplicate) or UNCERTAIN (post-crash replay): + # both are unverified and must not be silently re-applied. + raise UncertainMutationError( + f"step {key.step_index} for session {key.session_id[:8]} epoch " + f"{key.route_epoch} is {phase.value}; refusing silent replay" + ) + + def commit(self, key: StepKey, token: int) -> None: + with self._lock: + self._phase[key] = StepPhase.COMMITTED + self._token[key] = int(token) + + def mark_uncertain(self, key: StepKey, detail: str = "") -> None: + with self._lock: + # A committed step is verified; never downgrade it. + if self._phase.get(key) is StepPhase.COMMITTED: + return + self._phase[key] = StepPhase.UNCERTAIN + + def phase_of(self, key: StepKey) -> StepPhase | None: + with self._lock: + return self._phase.get(key) + + def committed_token(self, key: StepKey) -> int | None: + with self._lock: + return self._token.get(key) + + def has_uncertain(self) -> bool: + with self._lock: + return any(p is StepPhase.UNCERTAIN for p in self._phase.values()) + + +# --------------------------------------------------------------------------- # +# Restart / alpha failover: from token zero on a fresh compatible route. +# --------------------------------------------------------------------------- # + + +class RestartController: + """Alpha failover that restarts from token zero, never importing prior KV. + + RALPH runtime decision #14: when the alpha (the head owning embedding + final + head) fails, the route retries from token zero; unverified KV is never + migrated. :meth:`failover` opens the *next* route epoch and releases every + node-local shard's prior-epoch KV, so the restart begins with empty caches. The + KV manager then treats the failed epoch as stale (a later reference to it is + rejected), which is what keeps a half-computed cache from being reused. + """ + + def __init__(self, managers: Sequence[HotKvStateManager]) -> None: + self._managers = list(managers) + + def failover(self, session_id: str, failed_epoch: int) -> int: + """Advance to a fresh epoch and drop the failed epoch's KV on every shard.""" + new_epoch = int(failed_epoch) + 1 + for manager in self._managers: + manager.release(session_id, failed_epoch) + return new_epoch + + def assert_fresh_start(self, session_id: str, new_epoch: int) -> None: + """Verify no shard carries KV for the new epoch (a true token-zero restart). + + Any residual KV under the new epoch would be unverified imported state; + this fails closed so a restart can never silently attend over it. + """ + for manager in self._managers: + result = manager.resolve(session_id, new_epoch) + if not isinstance(result, CacheMiss): + raise FailureSemanticsError( + f"restart epoch {new_epoch} for session {session_id[:8]} is not " + "empty; refusing to import unverified KV" + ) + + +# --------------------------------------------------------------------------- # +# Billing / work records. +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class WorkRecord: + """A typed unit of served work, distinguishing what may be billed. + + ``tokens`` counts only *committed* generated tokens. Only a + :attr:`WorkStatus.COMPLETED` record is billable; cancelled/failed/unverified + records carry their confirmed token count for observability but are excluded + from billing so uncompleted or unverified work is never charged. + """ + + session_id: str + route_epoch: int + status: WorkStatus + tokens: int + failure_kind: FailureKind | None = None + detail: str = "" + + @property + def billable(self) -> bool: + return self.status is WorkStatus.COMPLETED + + def to_dict(self) -> dict: + return { + "session_id": self.session_id, + "route_epoch": self.route_epoch, + "status": self.status.value, + "tokens": self.tokens, + "failure_kind": self.failure_kind.value if self.failure_kind else None, + "detail": self.detail, + "billable": self.billable, + } + + +class WorkLedger: + """Append-only ledger of :class:`WorkRecord`, split by billing status.""" + + def __init__(self) -> None: + self._records: list[WorkRecord] = [] + self._lock = threading.Lock() + + def record(self, record: WorkRecord) -> WorkRecord: + with self._lock: + self._records.append(record) + return record + + def records(self) -> list[WorkRecord]: + with self._lock: + return list(self._records) + + def records_for(self, session_id: str) -> list[WorkRecord]: + with self._lock: + return [r for r in self._records if r.session_id == session_id] + + def billable_records(self) -> list[WorkRecord]: + with self._lock: + return [r for r in self._records if r.billable] + + def billable_tokens(self) -> int: + """Total tokens that may be charged (completed work only).""" + with self._lock: + return sum(r.tokens for r in self._records if r.billable) + + def counts_by_status(self) -> dict[str, int]: + counts: dict[str, int] = {s.value: 0 for s in WorkStatus} + with self._lock: + for record in self._records: + counts[record.status.value] += 1 + return counts + + def to_dict(self) -> dict: + with self._lock: + records = [r.to_dict() for r in self._records] + counts: dict[str, int] = {s.value: 0 for s in WorkStatus} + for record in records: + counts[record["status"]] += 1 + return { + "schema_version": 1, + "records": records, + "counts_by_status": counts, + "billable_tokens": sum(r["tokens"] for r in records if r["billable"]), + } + + +# --------------------------------------------------------------------------- # +# The hardened single-session stream runner. +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class RunOutcome: + """The typed result of one hardened generation attempt.""" + + session_id: str + route_epoch: int + status: WorkStatus + tokens: tuple[int, ...] + failure_kind: FailureKind | None + detail: str + + @property + def completed(self) -> bool: + return self.status is WorkStatus.COMPLETED + + @property + def token_count(self) -> int: + return len(self.tokens) + + @property + def restartable(self) -> bool: + return self.failure_kind in _RESTARTABLE + + def work_record(self) -> WorkRecord: + return WorkRecord( + session_id=self.session_id, + route_epoch=self.route_epoch, + status=self.status, + tokens=len(self.tokens), + failure_kind=self.failure_kind, + detail=self.detail, + ) + + +@dataclass(frozen=True) +class FailoverResult: + """The result of a run that may have restarted from token zero after a failure.""" + + outcome: RunOutcome + attempts: tuple[RunOutcome, ...] + restarts: int + + @property + def completed(self) -> bool: + return self.outcome.completed + + def to_dict(self) -> dict: + return { + "final_status": self.outcome.status.value, + "final_epoch": self.outcome.route_epoch, + "restarts": self.restarts, + "attempts": [ + { + "route_epoch": a.route_epoch, + "status": a.status.value, + "failure_kind": a.failure_kind.value if a.failure_kind else None, + "tokens": a.token_count, + } + for a in self.attempts + ], + } + + +class HardenedSessionRunner: + """Drive one Route Session's decode stream with bounded failure semantics. + + The runner owns a single full-shard :class:`KvBoundaryAdapter` (head **and** + tail, so a step samples a token) and threads every DGR-013 guarantee through a + step loop: + + * every step is bounded by a :class:`DeadlineGuard` and can observe a + :class:`CancellationToken`; + * every step is idempotent through an :class:`IdempotencyLedger` (a duplicate + returns the recorded token; an uncertain mutation is never replayed); + * any failure releases this session's KV (cancellation) and is recorded as a + typed :class:`WorkRecord` in the :class:`WorkLedger`; + * :meth:`run_with_failover` restarts a transient failure from token zero on a + fresh epoch via a :class:`RestartController`. + """ + + def __init__( + self, + adapter: KvBoundaryAdapter, + *, + clock: Callable[[], float] | None = None, + work_ledger: WorkLedger | None = None, + idempotency: IdempotencyLedger | None = None, + ) -> None: + if not (adapter.is_head and adapter.is_tail): + raise FailureSemanticsError( + "HardenedSessionRunner needs a full (head+tail) shard so decode " + "steps sample tokens; got a partial range " + f"(head={adapter.is_head} tail={adapter.is_tail})" + ) + self._adapter = adapter + self._manager: HotKvStateManager = adapter.manager + self._clock = clock or time.monotonic + self.work_ledger = work_ledger or WorkLedger() + self.idempotency = idempotency or IdempotencyLedger() + + # -- single attempt ------------------------------------------------------- + + def run( + self, + request: GenerationRequest, + *, + deadline: float | None = None, + heartbeat_timeout: float | None = None, + cancel_token: CancellationToken | None = None, + heartbeat: Callable[[int], bool] | None = None, + before_step: Callable[[int], None] | None = None, + ) -> RunOutcome: + """Run one attempt of ``request``; record and return a typed outcome. + + ``deadline`` (absolute, on the injected clock) and ``heartbeat_timeout`` + bound blocked steps. ``cancel_token`` lets a client cancel mid-stream. + ``heartbeat(step)`` returns ``True`` when a heartbeat was heard before that + step (resetting the health timer); ``before_step(step)`` is a fault- + injection / clock-advance hook run before each step and may raise + :class:`StreamTerminated` (e.g. a stream reset) or + :class:`OperationCancelled`. + """ + sid = request.session_id + epoch = request.route_epoch + guard = DeadlineGuard( + deadline=deadline, + heartbeat_timeout=heartbeat_timeout, + clock=self._clock, + ) + guard.start() + tokens: list[int] = [] + current_key: StepKey | None = None + try: + # step 0 is the prefill (emits the first token); steps 1..N are decodes. + for step_index in range(request.max_new_tokens): + # before_step is the fault-injection / clock-advance hook and may + # itself terminate the step (stream reset, cancel); run it first so + # a fault it raises takes effect on this step, then re-check the + # bounds it may have advanced (deadline / heartbeat / cancel). + if before_step is not None: + before_step(step_index) + if cancel_token is not None: + cancel_token.raise_if_cancelled() + if heartbeat is not None and heartbeat(step_index): + guard.heartbeat() + guard.check() + + current_key = StepKey(sid, epoch, step_index) + disposition = self.idempotency.begin(current_key) + if disposition.duplicate: + # Idempotent replay: reuse the recorded token, do not re-mutate. + assert disposition.token is not None + tokens.append(disposition.token) + continue + + token = self._execute_step(request, step_index, tokens) + if isinstance(token, CacheMiss): + # The expected KV was gone; the append never started, so this is + # a certain (not uncertain) miss — restartable from token zero. + return self._finish_failure( + request, + tokens, + FailureKind.CACHE_MISS, + str(token), + cancel_token, + ) + self.idempotency.commit(current_key, token) + tokens.append(token) + except (StreamTerminated, OperationCancelled) as exc: + return self._finish_failure( + request, tokens, classify_exception(exc), str(exc), cancel_token + ) + except ( + BoundaryContractError, + StaleRouteEpochError, + IncompatibleCacheRecipeError, + KvCacheMissError, + ) as exc: + # Deterministic protocol/state errors, all validated before any KV + # append committed — certain, not uncertain. + return self._finish_failure( + request, tokens, classify_exception(exc), str(exc), cancel_token + ) + except UncertainMutationError as exc: + # A replay of an unverified step reached the ledger — never silent. + return self._finish_failure( + request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token + ) + except Exception as exc: # noqa: BLE001 - unexpected == worker death + # An unexpected error mid-step may have left the KV half-mutated; mark + # the step uncertain so it can never be silently replayed, then fail + # closed as unverified work. + if current_key is not None: + self.idempotency.mark_uncertain(current_key, str(exc)) + return self._finish_failure( + request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token + ) + + return self._finish_completed(request, tokens) + + def _execute_step( + self, request: GenerationRequest, step_index: int, tokens: list[int] + ) -> int | CacheMiss: + sid = request.session_id + epoch = request.route_epoch + if step_index == 0: + out = self._adapter.prefill( + sid, epoch, token_ids=list(request.prompt_token_ids) + ) + else: + # expected_seq_len defends the KV layer against a desynchronised decode: + # prompt positions plus the tokens already committed this run. + expected = request.prompt_len + (step_index - 1) + out = self._adapter.decode( + sid, + epoch, + token_ids=[tokens[-1]], + expected_seq_len=expected, + ) + if isinstance(out, CacheMiss): + return out + if not isinstance(out, TailOutput): + raise FailureSemanticsError( + "full-shard step did not yield a sampled token; got " + f"{type(out).__name__}" + ) + return int(out.token_id) + + # -- failover across restarts -------------------------------------------- + + def run_with_failover( + self, + request: GenerationRequest, + controller: RestartController, + *, + max_restarts: int = 3, + **run_kwargs: Any, + ) -> FailoverResult: + """Run ``request``, restarting a transient failure from token zero. + + On a restartable failure (worker death, stream reset, cache miss) the + controller advances to a fresh epoch and drops the failed epoch's KV; the + next attempt re-prefills the whole prompt from token zero. A deterministic + failure (deadline, cancel, malformed bundle, stale epoch) is returned as-is + — retrying it would hang or fail identically. Per-attempt fault-injection + hooks (``before_step`` / ``heartbeat``) are only applied to the *first* + attempt so a restart runs clean. + """ + if max_restarts < 0: + raise FailureSemanticsError("max_restarts must be >= 0") + epoch = request.route_epoch + attempts: list[RunOutcome] = [] + first_kwargs = run_kwargs + for attempt in range(max_restarts + 1): + attempt_request = replace(request, route_epoch=epoch) + kwargs = first_kwargs if attempt == 0 else {} + outcome = self.run(attempt_request, **kwargs) + attempts.append(outcome) + if outcome.completed or not outcome.restartable or attempt == max_restarts: + return FailoverResult( + outcome=outcome, attempts=tuple(attempts), restarts=attempt + ) + # Alpha failover: fresh epoch, drop prior-epoch KV on every shard, and + # verify the new epoch starts empty (no unverified KV import). + epoch = controller.failover(request.session_id, epoch) + controller.assert_fresh_start(request.session_id, epoch) + # Unreachable: the loop always returns, but keep the type-checker happy. + raise FailureSemanticsError("run_with_failover exhausted without returning") + + # -- outcome bookkeeping -------------------------------------------------- + + def _finish_completed( + self, request: GenerationRequest, tokens: list[int] + ) -> RunOutcome: + outcome = RunOutcome( + session_id=request.session_id, + route_epoch=request.route_epoch, + status=WorkStatus.COMPLETED, + tokens=tuple(tokens), + failure_kind=None, + detail="", + ) + self.work_ledger.record(outcome.work_record()) + return outcome + + def _finish_failure( + self, + request: GenerationRequest, + tokens: list[int], + kind: FailureKind, + detail: str, + cancel_token: CancellationToken | None, + ) -> RunOutcome: + # Cancellation semantics: release this session's local KV so a failed or + # cancelled stream never leaks its cache. release() is idempotent. + self._manager.release(request.session_id, request.route_epoch) + if cancel_token is not None and kind is not FailureKind.CANCELLED: + # Ensure downstream shards sharing the token also stop. + cancel_token.cancel(kind.value) + outcome = RunOutcome( + session_id=request.session_id, + route_epoch=request.route_epoch, + status=work_status_for(kind), + tokens=tuple(tokens), + failure_kind=kind, + detail=detail, + ) + self.work_ledger.record(outcome.work_record()) + return outcome diff --git a/tests/test_batch_scheduler.py b/tests/test_batch_scheduler.py new file mode 100644 index 0000000..9893bb2 --- /dev/null +++ b/tests/test_batch_scheduler.py @@ -0,0 +1,472 @@ +"""Continuous batching and bounded admission (DGR-012). + +These tests drive the node-local continuous-batching scheduler with the *same* +pure-numpy KV-cached dense-Llama reference the Hot KV State manager uses +(DGR-007), imported from ``test_hot_kv_state``. That keeps the whole gate +deterministic, download-free, GPU-free, and API-credit-free while exercising the +real KV isolation path (``KvBoundaryAdapter`` + ``HotKvStateManager``) rather than +a mock. + +Coverage maps to the story's acceptance criteria: + +* bounded admission against weight/KV/scratch/queue budgets, +* compatible decode steps batched with per-session positions/outputs preserved, +* prefill never starving in-flight decode (explicit decode-first policy), +* backpressure when the bounded queue is full, +* capability telemetry reporting every required signal, +* a deterministic 1/2/4/8 concurrency sweep showing saturation and no + cross-session corruption. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from meshnet_node.hot_kv_state import ( + HotKvStateConfig, + HotKvStateManager, + KvBoundaryAdapter, + kv_recipe_for, +) +from meshnet_node.batch_scheduler import ( + AdmissionReason, + ContinuousBatchScheduler, + GenerationRequest, + KvBatchEngine, + NodeBudget, + Phase, + run_concurrency_sweep, +) + +# Reuse the certified numpy dense-Llama reference and shard from the DGR-007 gate. +from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard + + +# --------------------------------------------------------------------------- # +# Helpers. +# --------------------------------------------------------------------------- # + + +class _FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, delta: float) -> None: + self.now += delta + + +def _make_engine( + model: _KvDenseLlama | None = None, + *, + config: HotKvStateConfig | None = None, +) -> KvBatchEngine: + """A full-shard KV batch engine over the deterministic numpy dense-Llama.""" + model = model or _KvDenseLlama() + shard = _KvReferenceShard(model, 0, model.n_layers - 1) + manager = HotKvStateManager(kv_recipe_for(shard), config=config) + adapter = KvBoundaryAdapter(shard, manager) + return KvBatchEngine(adapter) + + +def _reference_tokens(model: _KvDenseLlama, prompt, n_new: int) -> list[int]: + return model.stateless_greedy(list(prompt), n_new) + + +def _generation(session_id: str, prompt, n_new: int, epoch: int = 0) -> GenerationRequest: + return GenerationRequest( + session_id=session_id, + route_epoch=epoch, + prompt_token_ids=tuple(prompt), + max_new_tokens=n_new, + ) + + +# --------------------------------------------------------------------------- # +# Bounded admission (weight / KV / scratch / queue budgets). +# --------------------------------------------------------------------------- # + + +def test_admission_respects_active_scratch_and_queue_budgets(): + "Admission fills active slots, queues the overflow, then rejects a full queue.\n\nTags: node, scheduler, admission" + engine = _make_engine() + budget = NodeBudget( + max_active_sessions=2, + scratch_bytes_per_session=1, + scratch_budget_bytes=2, # scratch also caps at 2 concurrent + max_queue_depth=1, + max_batch_size=2, + ) + scheduler = ContinuousBatchScheduler(engine, budget) + + a = scheduler.submit(_generation("a", [1, 2, 3], 4)) + b = scheduler.submit(_generation("b", [4, 5, 6], 4)) + assert a.reason is AdmissionReason.ADMITTED + assert b.reason is AdmissionReason.ADMITTED + + # Two active slots full -> the next goes to the bounded queue. + c = scheduler.submit(_generation("c", [7, 8, 9], 4)) + assert c.reason is AdmissionReason.QUEUED + + # Queue depth 1 is now full -> backpressure rejection. + d = scheduler.submit(_generation("d", [1, 1, 1], 4)) + assert d.reason is AdmissionReason.REJECTED_QUEUE_FULL + assert d.rejected + + telem = scheduler.telemetry() + assert telem.active_sessions == 2 + assert telem.queue_depth == 1 + assert telem.rejected_admissions_total == 1 + assert telem.rejected_by_reason[AdmissionReason.REJECTED_QUEUE_FULL.value] == 1 + + +def test_admission_rejects_a_session_that_cannot_fit_the_kv_budget(): + "A generation whose whole KV cannot fit the node budget is rejected up front.\n\nTags: node, scheduler, admission" + engine = _make_engine() + per_token = engine._manager.recipe.bytes_per_token() + # Budget holds only 3 positions; a prompt(4)+7 new = 10 final positions cannot fit. + budget = NodeBudget(kv_budget_bytes=per_token * 3) + scheduler = ContinuousBatchScheduler(engine, budget) + decision = scheduler.submit(_generation("big", [1, 2, 3, 4], 7)) + assert decision.reason is AdmissionReason.REJECTED_KV_BUDGET + assert scheduler.telemetry().rejected_admissions_total == 1 + + +def test_admission_rejects_when_per_session_scratch_exceeds_budget(): + "A per-session scratch larger than the whole scratch envelope is rejected.\n\nTags: node, scheduler, admission" + engine = _make_engine() + budget = NodeBudget(scratch_bytes_per_session=1024, scratch_budget_bytes=512) + scheduler = ContinuousBatchScheduler(engine, budget) + decision = scheduler.submit(_generation("s", [1, 2], 2)) + assert decision.reason is AdmissionReason.REJECTED_SCRATCH_BUDGET + + +def test_duplicate_submission_is_rejected(): + "Submitting a session id that is already scheduled is rejected as a duplicate.\n\nTags: node, scheduler, admission" + engine = _make_engine() + scheduler = ContinuousBatchScheduler(engine, NodeBudget(max_active_sessions=4)) + assert scheduler.submit(_generation("dup", [1, 2], 3)).reason is AdmissionReason.ADMITTED + assert scheduler.submit(_generation("dup", [3, 4], 3)).reason is AdmissionReason.REJECTED_DUPLICATE + + +def test_weight_budget_is_reported_in_telemetry(): + "The resident weight footprint is surfaced as a capability signal.\n\nTags: node, scheduler, telemetry" + engine = _make_engine() + budget = NodeBudget(weight_bytes=123_456) + scheduler = ContinuousBatchScheduler(engine, budget) + assert scheduler.telemetry().weight_bytes == 123_456 + + +# --------------------------------------------------------------------------- # +# Continuous batching preserves per-session positions and outputs. +# --------------------------------------------------------------------------- # + + +def test_batched_decode_preserves_per_session_positions_and_outputs(): + "Four sessions batched together each reproduce their own stateless tokens.\n\nTags: node, scheduler, batching" + model = _KvDenseLlama() + engine = _make_engine(model) + budget = NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4) + scheduler = ContinuousBatchScheduler(engine, budget) + + prompts = { + "alpha": [1, 2, 3, 4], + "bravo": [40, 39, 2, 15], + "charlie": [7, 7, 7, 7], + "delta": [31, 5, 18, 22], + } + n_new = 10 + references = {sid: _reference_tokens(model, p, n_new) for sid, p in prompts.items()} + # The four references must diverge, else "no cross-talk" would be vacuous. + assert len({tuple(v) for v in references.values()}) == 4 + + for sid, prompt in prompts.items(): + assert scheduler.submit(_generation(sid, prompt, n_new)).running + + outputs = scheduler.run_to_completion() + for sid in prompts: + assert outputs[sid] == references[sid], sid + + telem = scheduler.telemetry() + # A genuine batch formed: at least one decode tick carried all four sessions. + assert telem.batch_occupancy_max == 4 + assert telem.completed_sessions == 4 + assert telem.active_sessions == 0 + + +def test_positions_are_isolated_across_different_prompt_lengths(): + "Sessions with different prompt lengths keep independent positions when batched.\n\nTags: node, scheduler, batching" + model = _KvDenseLlama() + engine = _make_engine(model) + scheduler = ContinuousBatchScheduler( + engine, NodeBudget(max_active_sessions=3, max_batch_size=3, max_queue_depth=3) + ) + jobs = { + "short": ([5], 6), + "medium": ([2, 9, 14], 6), + "long": ([1, 2, 3, 4, 5, 6, 7], 6), + } + refs = {sid: _reference_tokens(model, p, n) for sid, (p, n) in jobs.items()} + for sid, (prompt, n) in jobs.items(): + scheduler.submit(_generation(sid, prompt, n)) + outputs = scheduler.run_to_completion() + for sid in jobs: + assert outputs[sid] == refs[sid], sid + + +# --------------------------------------------------------------------------- # +# Prefill does not starve decode. +# --------------------------------------------------------------------------- # + + +def test_prefill_does_not_starve_in_flight_decode(): + "A burst of new prefills never stalls an already-decoding session.\n\nTags: node, scheduler, fairness" + model = _KvDenseLlama() + engine = _make_engine(model) + # One prefill per tick (budget == a single prompt) so prefill is throttled and + # we can observe that decode still advances every tick. + budget = NodeBudget( + max_active_sessions=8, + max_batch_size=8, + max_queue_depth=8, + scratch_bytes_per_session=1, + scratch_budget_bytes=8, + max_prefill_tokens_per_tick=4, + ) + scheduler = ContinuousBatchScheduler(engine, budget) + + # Session A starts and prefills on tick 1. + scheduler.submit(_generation("A", [3, 14, 1, 5], 12)) + scheduler.run_tick() + a_state = scheduler.session_result("A") + assert a_state.phase is Phase.DECODING + a_len = len(a_state.generated) + assert a_len == 1 + + # Burst of new work arrives while A is decoding. + for sid in ("B", "C", "D", "E"): + scheduler.submit(_generation(sid, [2, 27, 18, 4], 12)) + + # Over the next few ticks A must decode on *every* tick (never starved), + # while at most one new session prefills per tick (prefill is bounded). + prefill_counts = [] + for _ in range(4): + report = scheduler.run_tick() + new_a_len = len(scheduler.session_result("A").generated) + assert new_a_len == a_len + 1, "decode of A stalled while prefills were pending" + a_len = new_a_len + assert "A" in report.decoded + prefill_counts.append(len(report.prefilled)) + + assert max(prefill_counts) <= 1, "prefill was not bounded per tick" + + +def test_decode_first_policy_is_explicit_in_a_single_tick(): + "In one tick decode of active sessions precedes prefill of new ones.\n\nTags: node, scheduler, fairness" + model = _KvDenseLlama() + engine = _make_engine(model) + scheduler = ContinuousBatchScheduler( + engine, + NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4, + scratch_bytes_per_session=1, scratch_budget_bytes=4), + ) + scheduler.submit(_generation("live", [1, 2, 3], 8)) + scheduler.run_tick() # 'live' prefills, now decoding + scheduler.submit(_generation("fresh", [9, 8, 7], 8)) + report = scheduler.run_tick() + assert "live" in report.decoded + assert "fresh" in report.prefilled + + +# --------------------------------------------------------------------------- # +# Backpressure and bounded memory. +# --------------------------------------------------------------------------- # + + +def test_backpressure_signals_when_queue_full_then_recovers(): + "A full queue rejects new work; a completed session frees a slot for the queue.\n\nTags: node, scheduler, backpressure" + engine = _make_engine() + budget = NodeBudget( + max_active_sessions=1, + max_batch_size=1, + max_queue_depth=1, + scratch_bytes_per_session=1, + scratch_budget_bytes=1, + ) + scheduler = ContinuousBatchScheduler(engine, budget) + assert scheduler.submit(_generation("first", [1, 2], 2)).running + assert scheduler.submit(_generation("second", [3, 4], 2)).reason is AdmissionReason.QUEUED + # Both a slot and the queue are full now. + assert scheduler.submit(_generation("third", [5, 6], 2)).reason is AdmissionReason.REJECTED_QUEUE_FULL + + # Drain 'first'; the queued 'second' must be pulled into the freed slot. + scheduler.run_to_completion() + outputs = scheduler.outputs() + assert set(outputs) == {"first", "second"} + + +def test_completed_sessions_release_kv_so_growth_is_bounded(): + "Finished sessions release their KV, so total KV returns to zero.\n\nTags: node, scheduler, backpressure" + engine = _make_engine() + scheduler = ContinuousBatchScheduler( + engine, NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=8) + ) + for sid in ("a", "b", "c", "d"): + scheduler.submit(_generation(sid, [1, 2, 3], 4)) + scheduler.run_to_completion() + telem = scheduler.telemetry() + assert telem.kv_total_bytes == 0, "KV not released after completion" + assert telem.active_sessions == 0 + assert telem.completed_sessions == 4 + + +# --------------------------------------------------------------------------- # +# Telemetry. +# --------------------------------------------------------------------------- # + + +def test_telemetry_reports_every_required_signal(): + "The capability snapshot reports sessions, queue, batch, KV, rates, rejections.\n\nTags: node, scheduler, telemetry" + model = _KvDenseLlama() + engine = _make_engine(model) + clock = _FakeClock() + budget = NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=1) + scheduler = ContinuousBatchScheduler(engine, budget, clock=clock) + + scheduler.submit(_generation("x", [1, 2, 3], 4)) + scheduler.submit(_generation("y", [4, 5, 6], 4)) + scheduler.submit(_generation("z", [7, 8, 9], 4)) # queued + rejected = scheduler.submit(_generation("w", [1, 1, 1], 4)) # queue full + assert rejected.rejected + + clock.advance(1.0) + scheduler.run_tick() # both prefill + clock.advance(1.0) + scheduler.run_tick() # both decode as a batch of 2 + + clock.advance(2.0) + telem = scheduler.telemetry() + snap = telem.to_dict() + for key in ( + "active_sessions", "queue_depth", "batch_occupancy_last", + "batch_occupancy_avg", "batch_occupancy_max", "weight_bytes", + "kv_total_bytes", "kv_budget_bytes", "kv_pressure", + "scratch_used_bytes", "scratch_budget_bytes", "scratch_pressure", + "prefill_tokens_total", "decode_tokens_total", + "prefill_tokens_per_sec", "decode_tokens_per_sec", + "rejected_admissions_total", "rejected_by_reason", + "completed_sessions", "ticks", + ): + assert key in snap, key + + assert telem.batch_occupancy_max == 2 + assert telem.prefill_tokens_total == 6 # two prompts of length 3 + assert telem.decode_tokens_total == 2 # one batched decode step, two sessions + assert telem.rejected_admissions_total == 1 + # Rates are deterministic under the injected clock: 4 seconds elapsed. + assert telem.decode_tokens_per_sec == pytest.approx(2 / 4.0) + assert telem.prefill_tokens_per_sec == pytest.approx(6 / 4.0) + assert 0.0 < telem.kv_pressure <= 1.0 + + +# --------------------------------------------------------------------------- # +# Concurrency 1/2/4/8 sweep: saturation and no corruption. +# --------------------------------------------------------------------------- # + + +def test_concurrency_sweep_identifies_saturation_without_corruption(): + "A 1/2/4/8 sweep raises batch occupancy, cuts ticks, and never corrupts output.\n\nTags: node, scheduler, benchmark" + model = _KvDenseLlama() + prompts = { + "s0": [1, 2, 3, 4], "s1": [5, 6, 7, 8], "s2": [9, 10, 11, 12], + "s3": [13, 14, 15, 16], "s4": [17, 18, 19, 20], "s5": [21, 22, 23, 24], + "s6": [25, 26, 27, 28], "s7": [29, 30, 31, 32], + } + n_new = 8 + requests = [_generation(sid, p, n_new) for sid, p in prompts.items()] + + sweep = run_concurrency_sweep( + lambda: _make_engine(model), + requests, + concurrency_levels=(1, 2, 4, 8), + ) + + assert sweep.corruption_free + assert [r.concurrency for r in sweep.results] == [1, 2, 4, 8] + + # No session hit a cache miss (budgets are sized to never evict here). + assert all(r.cache_misses == 0 for r in sweep.results) + assert all(r.rejected_admissions == 0 for r in sweep.results) + + # Each per-session stream matches the serialized (concurrency-1) reference. + for sid, prompt in prompts.items(): + assert list(sweep.reference_outputs[sid]) == _reference_tokens(model, prompt, n_new) + + occupancies = [r.avg_batch_occupancy for r in sweep.results] + ticks = [r.ticks for r in sweep.results] + tokens_per_tick = [r.tokens_per_tick for r in sweep.results] + + # Batching packs more sessions per decode step as concurrency rises, so + # average occupancy strictly increases and total ticks strictly decrease. + assert occupancies == sorted(occupancies) and len(set(occupancies)) == 4 + assert ticks == sorted(ticks, reverse=True) and len(set(ticks)) == 4 + # Aggregate work per tick rises with concurrency (the throughput win). + assert tokens_per_tick == sorted(tokens_per_tick) + + # For eight equal-length jobs the node keeps saturating up to the top level. + assert sweep.saturation_concurrency == 8 + + # The report is JSON-safe for durable evidence. + import json + + json.dumps(sweep.to_dict()) + + +def test_concurrency_sweep_saturates_below_max_when_load_is_small(): + "With fewer concurrent jobs than slots, saturation is found below the top level.\n\nTags: node, scheduler, benchmark" + model = _KvDenseLlama() + # Only three jobs: at concurrency 4 and 8 the batch can never exceed 3, so + # occupancy stops rising past the load and saturation is detected early. + requests = [ + _generation("j0", [1, 2, 3], 6), + _generation("j1", [4, 5, 6], 6), + _generation("j2", [7, 8, 9], 6), + ] + sweep = run_concurrency_sweep( + lambda: _make_engine(model), requests, concurrency_levels=(1, 2, 4, 8) + ) + assert sweep.corruption_free + assert sweep.saturation_concurrency <= 4 + # Levels at or above the load size share the same occupancy/tick profile. + top = [r for r in sweep.results if r.concurrency >= 4] + assert len({r.ticks for r in top}) == 1 + + +# --------------------------------------------------------------------------- # +# Engine contract guards. +# --------------------------------------------------------------------------- # + + +def test_kv_batch_engine_requires_a_full_shard(): + "The batch engine rejects a partial (non head+tail) shard.\n\nTags: node, scheduler" + model = _KvDenseLlama() + head = _KvReferenceShard(model, 0, 2) # head only, not tail + manager = HotKvStateManager(kv_recipe_for(head)) + adapter = KvBoundaryAdapter(head, manager) + with pytest.raises(Exception): + KvBatchEngine(adapter) + + +def test_run_to_completion_is_bounded_against_misconfiguration(): + "run_to_completion raises rather than looping forever when work cannot drain.\n\nTags: node, scheduler" + engine = _make_engine() + scheduler = ContinuousBatchScheduler( + engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4) + ) + scheduler.submit(_generation("only", [1, 2], 3)) + # A tiny explicit tick ceiling is exceeded deterministically. + with pytest.raises(Exception): + scheduler.run_to_completion(max_ticks=1) diff --git a/tests/test_failure_semantics.py b/tests/test_failure_semantics.py new file mode 100644 index 0000000..8059402 --- /dev/null +++ b/tests/test_failure_semantics.py @@ -0,0 +1,611 @@ +"""Bounded failure, cancellation, and restart semantics (DGR-013). + +These tests drive the hardened per-session decode stream with the *same* +pure-numpy KV-cached dense-Llama reference the Hot KV State manager (DGR-007) and +the continuous-batch scheduler (DGR-012) use, imported from ``test_hot_kv_state``. +The whole matrix stays deterministic, download-free, GPU-free, and API-credit-free +while exercising the real KV isolation path (``KvBoundaryAdapter`` + +``HotKvStateManager``) rather than a mock. + +Coverage maps to the story's acceptance criteria: + +* deadlines and heartbeat/health loss terminate blocked stream operations, +* cancellation propagates across every Shard and releases KV + queued buffers, +* duplicate steps are idempotent; uncertain mutations are never replayed silently, +* alpha failover restarts from token zero rather than importing unverified KV, +* worker death / stream reset / malformed bundle / stale epoch / cache miss, +* billing/work records distinguish completed, cancelled, failed, and unverified. +""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +from meshnet_node.batch_scheduler import ( + ContinuousBatchScheduler, + DoneReason, + GenerationRequest, + KvBatchEngine, + NodeBudget, +) +from meshnet_node.boundary_adapter import BoundaryBundle, BoundaryContractError +from meshnet_node.hot_kv_state import ( + CacheMiss, + CacheMissReason, + HotKvStateConfig, + HotKvStateManager, + KvBoundaryAdapter, + StaleRouteEpochError, + kv_recipe_for, +) +from meshnet_node.failure_semantics import ( + CancellationToken, + DeadlineGuard, + FailureKind, + HardenedSessionRunner, + IdempotencyLedger, + OperationCancelled, + RestartController, + ShardCancellationGroup, + StepKey, + StreamTerminated, + UncertainMutationError, + WorkLedger, + WorkRecord, + WorkStatus, + classify_exception, + work_status_for, +) + +# Reuse the certified numpy dense-Llama reference and shard from the DGR-007 gate. +from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard + + +# --------------------------------------------------------------------------- # +# Helpers. +# --------------------------------------------------------------------------- # + + +class _FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, delta: float) -> None: + self.now += delta + + +class _FaultyShard(_KvReferenceShard): + """A full-shard reference that raises on the Nth ``run_layers_cached`` call. + + ``run_layers_cached`` is invoked once per stream step, so ``fail_at_call=k`` + simulates a worker dying at step ``k-1`` (calls are 1-indexed). The call + counter persists across attempts, so a restart on a fresh epoch keeps counting + and does not re-trip the same fault. + """ + + def __init__(self, model, start, end, *, fail_at_call=None, error=None): + super().__init__(model, start, end) + self._fail_at_call = fail_at_call + self._error = error or RuntimeError("worker died mid-step") + self.calls = 0 + + def run_layers_cached(self, hidden, *, positions, past_kv): + self.calls += 1 + if self._fail_at_call is not None and self.calls == self._fail_at_call: + raise self._error + return super().run_layers_cached(hidden, positions=positions, past_kv=past_kv) + + +def _make_adapter(model=None, *, config=None, shard=None): + """A full-shard KV boundary adapter over the deterministic numpy dense-Llama.""" + model = model or _KvDenseLlama() + shard = shard or _KvReferenceShard(model, 0, model.n_layers - 1) + manager = HotKvStateManager(kv_recipe_for(shard), config=config) + adapter = KvBoundaryAdapter(shard, manager) + return adapter + + +def _generation(session_id, prompt, n_new, epoch=0): + return GenerationRequest( + session_id=session_id, + route_epoch=epoch, + prompt_token_ids=tuple(prompt), + max_new_tokens=n_new, + ) + + +# --------------------------------------------------------------------------- # +# Happy path (the baseline the failure paths deviate from). +# --------------------------------------------------------------------------- # + + +def test_clean_run_matches_stateless_reference_and_is_billable(): + "A clean stream reproduces the stateless tokens and records completed work.\n\nTags: node, failure, billing" + model = _KvDenseLlama() + adapter = _make_adapter(model) + runner = HardenedSessionRunner(adapter) + prompt = [1, 2, 3, 4] + n_new = 8 + outcome = runner.run(_generation("clean", prompt, n_new)) + assert outcome.status is WorkStatus.COMPLETED + assert list(outcome.tokens) == model.stateless_greedy(prompt, n_new) + record = runner.work_ledger.records_for("clean")[0] + assert record.billable + assert record.tokens == n_new + assert runner.work_ledger.billable_tokens() == n_new + + +# --------------------------------------------------------------------------- # +# Deadlines and heartbeat/health loss terminate blocked operations. +# --------------------------------------------------------------------------- # + + +def test_deadline_terminates_a_blocked_stream_and_releases_kv(): + "A deadline reached mid-stream terminates the run and frees its KV.\n\nTags: node, failure, deadline" + clock = _FakeClock() + adapter = _make_adapter() + manager = adapter.manager + runner = HardenedSessionRunner(adapter, clock=clock) + + # Each step advances the clock by 1.0; the deadline fires at t=3. + def before_step(_step): + clock.advance(1.0) + + outcome = runner.run( + _generation("slow", [5, 6, 7], 20), + deadline=3.0, + before_step=before_step, + ) + assert outcome.status is WorkStatus.FAILED + assert outcome.failure_kind is FailureKind.DEADLINE_EXCEEDED + # The stream did not hang and did not finish: only the steps before the + # deadline committed, and the session's KV was released. + assert outcome.token_count < 20 + assert isinstance(manager.resolve("slow", 0), CacheMiss) + + +def test_heartbeat_loss_terminates_a_blocked_stream(): + "Losing the peer heartbeat past the timeout terminates the stream.\n\nTags: node, failure, heartbeat" + clock = _FakeClock() + adapter = _make_adapter() + runner = HardenedSessionRunner(adapter, clock=clock) + + def before_step(_step): + clock.advance(1.0) + + # Heartbeats stop arriving after step 2; with a timeout of 1.5 the gap grows + # past the bound and the stream is terminated (health loss). + def heartbeat(step): + return step < 2 + + outcome = runner.run( + _generation("hb", [9, 8, 7], 20), + heartbeat_timeout=1.5, + heartbeat=heartbeat, + before_step=before_step, + ) + assert outcome.status is WorkStatus.FAILED + assert outcome.failure_kind is FailureKind.HEARTBEAT_LOST + assert outcome.token_count < 20 + + +def test_deadline_guard_reports_remaining_and_resets_on_heartbeat(): + "The guard exposes remaining time and a heartbeat resets the health timer.\n\nTags: node, failure, deadline" + clock = _FakeClock() + guard = DeadlineGuard(deadline=10.0, heartbeat_timeout=2.0, clock=clock) + guard.start() + guard.check() + assert guard.remaining() == 10.0 + clock.advance(1.5) + guard.heartbeat() # health refreshed at t=1.5 + clock.advance(1.0) # gap since heartbeat is 1.0 < 2.0 + guard.check() + clock.advance(2.5) # gap since heartbeat is now 3.5 > 2.0 + with pytest.raises(StreamTerminated) as exc: + guard.check() + assert exc.value.kind is FailureKind.HEARTBEAT_LOST + + +# --------------------------------------------------------------------------- # +# Cancellation propagates across shards and releases KV + queued buffers. +# --------------------------------------------------------------------------- # + + +def test_cancellation_token_terminates_stream_and_releases_kv(): + "A client cancel mid-stream stops the run and releases the session KV.\n\nTags: node, failure, cancel" + adapter = _make_adapter() + manager = adapter.manager + token = CancellationToken() + runner = HardenedSessionRunner(adapter) + + # Cancel after two steps have run. + def before_step(step): + if step == 2: + token.cancel("client-hangup") + + outcome = runner.run( + _generation("cancelme", [1, 2, 3], 20), + cancel_token=token, + before_step=before_step, + ) + assert outcome.status is WorkStatus.CANCELLED + assert outcome.failure_kind is FailureKind.CANCELLED + assert outcome.token_count == 2 # steps 0 and 1 committed before the cancel + assert isinstance(manager.resolve("cancelme", 0), CacheMiss) + + +def test_shard_cancellation_group_releases_every_shard_and_queued_buffers(): + "One cancel frees KV on every node-local shard and releases queued buffers.\n\nTags: node, failure, cancel" + model = _KvDenseLlama() + # Three node-local shards of the same route, each with its own KV manager. + managers = [] + for start, end in ((0, 1), (2, 3), (4, 5)): + shard = _KvReferenceShard(model, start, end) + mgr = HotKvStateManager(kv_recipe_for(shard)) + mgr.open("route", 0) # each holds live state for the session + managers.append(mgr) + + released_buffers = [] + group = ShardCancellationGroup("route", 0) + for mgr in managers: + group.add_shard(mgr) + group.add_queued_buffer(lambda: released_buffers.append("bundle-a")) + group.add_queued_buffer(lambda: released_buffers.append("bundle-b")) + + outcome = group.cancel() + assert outcome.shards_released == 3 + assert outcome.buffers_released == 2 + assert released_buffers == ["bundle-a", "bundle-b"] + # Every shard's KV is gone: a lookup now yields an explicit released miss. + for mgr in managers: + miss = mgr.resolve("route", 0) + assert isinstance(miss, CacheMiss) + assert miss.reason is CacheMissReason.RELEASED + # Cancellation is idempotent. + again = group.cancel() + assert again.shards_released == 0 + assert again.buffers_released == 0 + + +def test_scheduler_cancel_drains_queue_and_releases_active_kv(): + "The scheduler cancel drops queued work and frees an active session's KV.\n\nTags: node, scheduler, cancel" + model = _KvDenseLlama() + shard = _KvReferenceShard(model, 0, model.n_layers - 1) + manager = HotKvStateManager(kv_recipe_for(shard)) + engine = KvBatchEngine(KvBoundaryAdapter(shard, manager)) + scheduler = ContinuousBatchScheduler( + engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4) + ) + assert scheduler.submit(_generation("active", [1, 2, 3], 8)).running + assert scheduler.submit(_generation("waiting", [4, 5, 6], 8)).reason.value == "queued" + scheduler.run_tick() # 'active' prefills and starts decoding, holding KV + + # Cancel the queued one: it leaves the queue without ever taking a slot. + assert scheduler.cancel("waiting") is True + # Cancel the active one: its KV is released and it is recorded as cancelled. + assert scheduler.cancel("active") is True + assert manager.total_bytes == 0 + + telem = scheduler.telemetry() + assert telem.cancelled_sessions == 2 + assert telem.completed_sessions == 0 + assert telem.active_sessions == 0 + assert telem.queue_depth == 0 + # Cancelling an unknown / already-finished session is a no-op. + assert scheduler.cancel("active") is False + assert scheduler.cancel("never-seen") is False + + +def test_scheduler_cancel_rejects_a_completed_reason(): + "cancel() refuses a non-terminal reason so completed work is never faked.\n\nTags: node, scheduler, cancel" + model = _KvDenseLlama() + shard = _KvReferenceShard(model, 0, model.n_layers - 1) + manager = HotKvStateManager(kv_recipe_for(shard)) + engine = KvBatchEngine(KvBoundaryAdapter(shard, manager)) + scheduler = ContinuousBatchScheduler(engine) + scheduler.submit(_generation("x", [1, 2], 4)) + with pytest.raises(Exception): + scheduler.cancel("x", reason=DoneReason.COMPLETED) + + +# --------------------------------------------------------------------------- # +# Idempotency: duplicate steps are no-ops; uncertain mutations never replay. +# --------------------------------------------------------------------------- # + + +def test_duplicate_step_delivery_is_idempotent_no_remutation(): + "Replaying a committed step returns the recorded token without re-mutating KV.\n\nTags: node, failure, idempotency" + ledger = IdempotencyLedger() + key = StepKey("s", 0, 5) + disposition = ledger.begin(key) + assert disposition.fresh + ledger.commit(key, 42) + # A duplicate delivery of the same step returns the recorded token and is a + # no-op — the caller must not re-run the mutation. + replay = ledger.begin(key) + assert replay.duplicate + assert replay.token == 42 + + +def test_idempotent_run_replays_tokens_without_advancing_kv(): + "Re-running a completed stream on the same ledger/epoch re-mutates nothing.\n\nTags: node, failure, idempotency" + model = _KvDenseLlama() + adapter = _make_adapter(model) + ledger = IdempotencyLedger() + runner = HardenedSessionRunner(adapter, idempotency=ledger) + request = _generation("idem", [3, 1, 4], 6) + + first = runner.run(request) + assert first.status is WorkStatus.COMPLETED + kv_len_after_first = adapter.manager.get("idem", 0).seq_len + + # A duplicate delivery of the entire stream: every step is a committed + # duplicate, so the runner replays the identical tokens and the KV length is + # unchanged (no double-append). + second = runner.run(request) + assert second.status is WorkStatus.COMPLETED + assert list(second.tokens) == list(first.tokens) + assert adapter.manager.get("idem", 0).seq_len == kv_len_after_first + + +def test_uncertain_mutation_is_never_replayed_silently(): + "A step marked uncertain refuses a silent replay; it must be verified/restarted.\n\nTags: node, failure, idempotency" + ledger = IdempotencyLedger() + key = StepKey("s", 0, 3) + ledger.begin(key) + ledger.mark_uncertain(key, "worker died before ack") + # Replaying an uncertain mutation is refused rather than silently re-applied. + with pytest.raises(UncertainMutationError): + ledger.begin(key) + assert ledger.has_uncertain() + + +def test_in_flight_duplicate_is_treated_as_uncertain(): + "A second begin before commit is refused (concurrent duplicate is unverified).\n\nTags: node, failure, idempotency" + ledger = IdempotencyLedger() + key = StepKey("s", 0, 1) + ledger.begin(key) # in-flight, not yet committed + with pytest.raises(UncertainMutationError): + ledger.begin(key) + + +# --------------------------------------------------------------------------- # +# Worker death, stream reset, malformed bundle, stale epoch, cache miss. +# --------------------------------------------------------------------------- # + + +def test_worker_death_midstream_is_unverified_and_marks_step_uncertain(): + "A worker dying mid-step yields unverified work and an unreplayable step.\n\nTags: node, failure, worker-death" + model = _KvDenseLlama() + # Fail on the 3rd step call (step index 2), after two tokens committed. + shard = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3) + adapter = _make_adapter(model, shard=shard) + ledger = IdempotencyLedger() + runner = HardenedSessionRunner(adapter, idempotency=ledger) + + outcome = runner.run(_generation("dead", [1, 2, 3], 8)) + assert outcome.status is WorkStatus.UNVERIFIED + assert outcome.failure_kind is FailureKind.WORKER_DEATH + assert outcome.token_count == 2 # the two committed steps + assert not outcome.completed + # The failed step is uncertain and can never be silently replayed. + assert ledger.has_uncertain() + with pytest.raises(UncertainMutationError): + ledger.begin(StepKey("dead", 0, 2)) + # KV was released on failure. + assert isinstance(adapter.manager.resolve("dead", 0), CacheMiss) + + +def test_stream_reset_is_restartable_failure(): + "A stream reset injected mid-stream fails the run as a restartable transport loss.\n\nTags: node, failure, stream-reset" + adapter = _make_adapter() + runner = HardenedSessionRunner(adapter) + + def before_step(step): + if step == 2: + raise StreamTerminated(FailureKind.STREAM_RESET, "peer reset the stream") + + outcome = runner.run(_generation("reset", [1, 2, 3], 8), before_step=before_step) + assert outcome.status is WorkStatus.FAILED + assert outcome.failure_kind is FailureKind.STREAM_RESET + assert outcome.restartable + + +def test_malformed_bundle_is_classified_and_does_not_corrupt_kv(): + "A malformed activation bundle is rejected and leaves the KV context empty.\n\nTags: node, failure, malformed-bundle" + model = _KvDenseLlama() + mid = _KvReferenceShard(model, 2, 3) # middle range: not head, not tail + manager = HotKvStateManager(kv_recipe_for(mid)) + adapter = KvBoundaryAdapter(mid, manager) + assert not adapter.is_head and not adapter.is_tail + + # A bundle that hands over at the wrong layer is malformed. + bad = BoundaryBundle( + architecture_adapter=adapter.architecture.adapter, + schema_version=adapter.architecture.boundary_schema_version, + tensor_name=adapter.architecture.boundary_tensor_name, + residual=np.zeros((1, 3, model.hidden), dtype=np.float32), + positions=np.arange(3, dtype=np.int64)[None, :], + next_layer=adapter.start_layer + 5, # wrong handover layer + normalized=False, + ) + with pytest.raises(BoundaryContractError) as exc: + adapter.prefill("mal", 0, boundary=bad) + assert classify_exception(exc.value) is FailureKind.MALFORMED_BUNDLE + # The malformed step never appended KV: the context is empty, not corrupted. + assert manager.get("mal", 0).seq_len == 0 + + +def test_stale_epoch_reference_is_rejected_and_classified(): + "A reference to a superseded epoch is rejected as stale, never silently reused.\n\nTags: node, failure, stale-epoch" + model = _KvDenseLlama() + adapter = _make_adapter(model) + manager = adapter.manager + manager.open("sess", 5) # current epoch is now 5 + with pytest.raises(StaleRouteEpochError) as exc: + manager.resolve("sess", 4) # epoch 4 is stale + assert classify_exception(exc.value) is FailureKind.STALE_EPOCH + + # Driving the hardened runner on the stale epoch fails closed as STALE_EPOCH. + runner = HardenedSessionRunner(adapter) + outcome = runner.run(_generation("sess", [1, 2, 3], 4, epoch=3)) + assert outcome.status is WorkStatus.FAILED + assert outcome.failure_kind is FailureKind.STALE_EPOCH + + +def test_cache_miss_midstream_is_restartable(): + "A KV eviction mid-stream surfaces an explicit cache miss the head can restart.\n\nTags: node, failure, cache-miss" + adapter = _make_adapter() + manager = adapter.manager + runner = HardenedSessionRunner(adapter) + + # Evict the session's KV just before step 3's decode. + def before_step(step): + if step == 3: + manager.release("evict", 0) + + outcome = runner.run(_generation("evict", [1, 2, 3], 10), before_step=before_step) + assert outcome.failure_kind is FailureKind.CACHE_MISS + assert outcome.restartable + assert outcome.token_count == 3 # steps 0..2 committed before the eviction + + +# --------------------------------------------------------------------------- # +# Alpha failover: restart from token zero, never import unverified KV. +# --------------------------------------------------------------------------- # + + +def test_alpha_failover_restarts_from_token_zero_and_completes(): + "A transient worker death fails over to a fresh epoch and reproduces the tokens.\n\nTags: node, failure, failover" + model = _KvDenseLlama() + # Die on the 3rd step of the first attempt; the persistent call counter means + # the restart (which keeps counting) does not re-trip the fault. + shard = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3) + adapter = _make_adapter(model, shard=shard) + manager = adapter.manager + runner = HardenedSessionRunner(adapter) + controller = RestartController([manager]) + + prompt = [7, 3, 9, 1] + n_new = 6 + result = runner.run_with_failover( + _generation("alpha", prompt, n_new, epoch=0), controller, max_restarts=2 + ) + assert result.completed + assert result.restarts == 1 + # The restart began on a fresh epoch and reproduced the full stateless stream + # from token zero — no half-computed KV was imported. + assert result.outcome.route_epoch == 1 + assert list(result.outcome.tokens) == model.stateless_greedy(prompt, n_new) + # The failed epoch's KV is gone and the epoch is now stale. + with pytest.raises(StaleRouteEpochError): + manager.resolve("alpha", 0) + # First attempt was unverified, the restart completed: only the restart bills. + statuses = [a.status for a in result.attempts] + assert statuses == [WorkStatus.UNVERIFIED, WorkStatus.COMPLETED] + assert runner.work_ledger.billable_tokens() == n_new + + +def test_failover_refuses_to_import_unverified_kv(): + "assert_fresh_start fails closed if any shard still holds new-epoch KV.\n\nTags: node, failure, failover" + model = _KvDenseLlama() + adapter = _make_adapter(model) + manager = adapter.manager + controller = RestartController([manager]) + + new_epoch = controller.failover("s", 0) + assert new_epoch == 1 + # A clean fresh start passes. + controller.assert_fresh_start("s", new_epoch) + # If unverified KV were present under the new epoch, the guard refuses it. + manager.open("s", new_epoch) + manager.append( + "s", + new_epoch, + {i: (np.zeros((1, model.n_heads, model.head_dim), dtype=np.float32), + np.zeros((1, model.n_heads, model.head_dim), dtype=np.float32)) + for i in range(model.n_layers)}, + ) + with pytest.raises(Exception): + controller.assert_fresh_start("s", new_epoch) + + +def test_non_restartable_failure_is_not_retried(): + "A deterministic failure (deadline) returns immediately without a restart.\n\nTags: node, failure, failover" + clock = _FakeClock() + adapter = _make_adapter() + runner = HardenedSessionRunner(adapter, clock=clock) + controller = RestartController([adapter.manager]) + + def before_step(_step): + clock.advance(1.0) + + result = runner.run_with_failover( + _generation("bounded", [1, 2, 3], 20), + controller, + max_restarts=3, + deadline=2.0, + before_step=before_step, + ) + assert not result.completed + assert result.restarts == 0 + assert result.outcome.failure_kind is FailureKind.DEADLINE_EXCEEDED + + +# --------------------------------------------------------------------------- # +# Billing / work records distinguish completed, cancelled, failed, unverified. +# --------------------------------------------------------------------------- # + + +def test_work_ledger_distinguishes_all_four_statuses(): + "The work ledger keeps completed/cancelled/failed/unverified distinct.\n\nTags: node, failure, billing" + ledger = WorkLedger() + ledger.record(WorkRecord("a", 0, WorkStatus.COMPLETED, tokens=8)) + ledger.record(WorkRecord("b", 0, WorkStatus.CANCELLED, tokens=3, + failure_kind=FailureKind.CANCELLED)) + ledger.record(WorkRecord("c", 0, WorkStatus.FAILED, tokens=1, + failure_kind=FailureKind.DEADLINE_EXCEEDED)) + ledger.record(WorkRecord("d", 0, WorkStatus.UNVERIFIED, tokens=2, + failure_kind=FailureKind.WORKER_DEATH)) + + counts = ledger.counts_by_status() + assert counts == { + "completed": 1, "cancelled": 1, "failed": 1, "unverified": 1, + } + # Only completed work is billable — cancelled/failed/unverified tokens are + # recorded for observability but never charged. + assert ledger.billable_tokens() == 8 + assert [r.session_id for r in ledger.billable_records()] == ["a"] + # JSON-safe for durable evidence. + payload = ledger.to_dict() + assert payload["billable_tokens"] == 8 + assert payload["counts_by_status"]["unverified"] == 1 + json.dumps(payload) + + +def test_work_status_and_classification_mapping(): + "Failure kinds map to the right billing status and exception classes.\n\nTags: node, failure, billing" + assert work_status_for(FailureKind.CANCELLED) is WorkStatus.CANCELLED + assert work_status_for(FailureKind.WORKER_DEATH) is WorkStatus.UNVERIFIED + # A stream reset detected at a step boundary is a certain failure (nothing + # committed for that step) — only an unexpected mid-step error is unverified. + assert work_status_for(FailureKind.STREAM_RESET) is WorkStatus.FAILED + assert work_status_for(FailureKind.DEADLINE_EXCEEDED) is WorkStatus.FAILED + assert work_status_for(FailureKind.MALFORMED_BUNDLE) is WorkStatus.FAILED + assert work_status_for(FailureKind.STALE_EPOCH) is WorkStatus.FAILED + assert work_status_for(FailureKind.CACHE_MISS) is WorkStatus.FAILED + + assert classify_exception(OperationCancelled()) is FailureKind.CANCELLED + assert classify_exception(StaleRouteEpochError("x")) is FailureKind.STALE_EPOCH + assert classify_exception(BoundaryContractError("x")) is FailureKind.MALFORMED_BUNDLE + assert classify_exception(RuntimeError("boom")) is FailureKind.WORKER_DEATH + assert ( + classify_exception(StreamTerminated(FailureKind.HEARTBEAT_LOST)) + is FailureKind.HEARTBEAT_LOST + )