feat: checkpoint batching and release-gate stories
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user