feat: checkpoint batching and release-gate stories

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)