1025 lines
41 KiB
Python
1025 lines
41 KiB
Python
"""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
|