"""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