From 3191b58705d1a99632ee8e7fcc820eb3323c230b Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 22 Sep 2026 04:13:10 -0400 Subject: [PATCH 020/103] sched: scheduler allocation log for KV capacity probes Env-gated per-step log of blocks allocated and freed and scheduled tokens per request; prefix-cache hits are read behind a schema check so an upstream field rename fails visibly. --- vllm/v1/core/sched/e2_alloc_log.py | 111 +++++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 3 + 2 files changed, 114 insertions(+) create mode 100644 vllm/v1/core/sched/e2_alloc_log.py diff --git a/vllm/v1/core/sched/e2_alloc_log.py b/vllm/v1/core/sched/e2_alloc_log.py new file mode 100644 index 0000000..0efcbd7 --- /dev/null +++ b/vllm/v1/core/sched/e2_alloc_log.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Debug-only scheduler allocation log for KV capacity probes (off unless VLLM_E2_SCHED_LOG=). + +One JSON line per scheduler update while any request exists, and one "fin" line per request as it is freed: + {"t": wall, "m": monotonic, "step": n, "total": pool blocks, "free": free blocks, "running": r, "waiting": w, + "tot": tokens scheduled this step, + "reqs": [[id, prompt, computed, output, cached, preemptions, spec_scheduled, status, [blocks per group], + tokens scheduled for this request this step]]} + {"t": wall, "m": monotonic, "fin": [id, prompt, computed, output, cached, preemptions, status, [blocks per group]]} +Blocks are counted per KV-cache group, null blocks excluded. A logging failure writes an "error" line and never +raises into the scheduler. Unset = the scheduler only pays one env check per update. +""" +import json +import os +import time + +_PATH = os.environ.get("VLLM_E2_SCHED_LOG", "") +_state = {"f": None, "step": 0, "flush": 0.0, "wrapped": False, "errors": 0, "checked": False} +REQ_FIELDS = ("request_id", "num_prompt_tokens", "num_computed_tokens", "num_output_tokens", "num_preemptions", + "status", "prefill_stats") + + +def enabled() -> bool: + return bool(_PATH) + + +def _blocks(sched, rid): + return [sum(1 for b in m.req_to_blocks.get(rid, ()) if not b.is_null) + for m in sched.kv_cache_manager.coordinator.single_type_managers] + + +def _status(r): + return getattr(r.status, "name", str(r.status)) + + +def _cached(r): + """Prefix-cache hit for this request. vLLM keeps it in `request.prefill_stats` (PrefillStats.num_cached_tokens = + local + external); the output processor takes those stats after the first prefill (`take_prefill_stats` sets the + attribute to None), so -1 means "no longer knowable from the scheduler", never a silent zero.""" + ps = getattr(r, "prefill_stats", None) + return int(getattr(ps, "num_cached_tokens", -1)) if ps is not None else -1 + + +def _contract(r): + """Once per process: every field this log reads must exist on the real Request object. A missing field is written + once as a schema_error record — the evaluator fails closed on it — instead of turning every row into an error.""" + if _state["checked"]: + return + _state["checked"] = True + missing = [f for f in REQ_FIELDS if not hasattr(r, f)] + if missing: + _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), "schema_error": missing}, force=True) + + +def _write(rec, force=False): + f = _state["f"] + if f is None: + f = _state["f"] = open(_PATH, "a", buffering=1 << 16) + f.write(json.dumps(rec, separators=(",", ":")) + "\n") + if force or rec["t"] - _state["flush"] > 0.5: + f.flush() + _state["flush"] = rec["t"] + + +def _error(e): + _state["errors"] += 1 + if _state["errors"] <= 20: + try: + _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), + "error": f"{type(e).__name__}: {e}"[:300]}, force=True) + except Exception: # noqa: BLE001 + pass + + +def _wrap_free(sched): + orig = sched._free_request + + def _free_request(request, *args, **kwargs): + try: + _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), + "fin": [request.request_id, request.num_prompt_tokens, request.num_computed_tokens, + request.num_output_tokens, _cached(request), request.num_preemptions, + _status(request), _blocks(sched, request.request_id)]}, force=True) + except Exception as e: # noqa: BLE001 - the debug log must never break the scheduler + _error(e) + return orig(request, *args, **kwargs) + + sched._free_request = _free_request + _state["wrapped"] = True + + +def log_step(sched, scheduler_output) -> None: + try: + if not _state["wrapped"]: + _wrap_free(sched) + if not sched.requests: + return + _state["step"] += 1 + spec = getattr(scheduler_output, "scheduled_spec_decode_tokens", None) or {} + pool = sched.kv_cache_manager.block_pool + nsched = getattr(scheduler_output, "num_scheduled_tokens", None) or {} + _contract(next(iter(sched.requests.values()))) + rows = [[rid, r.num_prompt_tokens, r.num_computed_tokens, r.num_output_tokens, _cached(r), + r.num_preemptions, len(spec.get(rid, ())), _status(r), _blocks(sched, rid), int(nsched.get(rid, 0))] + for rid, r in sched.requests.items()] + _write({"t": round(time.time(), 4), "m": round(time.monotonic(), 4), "step": _state["step"], + "total": pool.num_gpu_blocks, "free": pool.get_num_free_blocks(), + "running": len(sched.running), "waiting": len(sched.waiting), + "tot": int(getattr(scheduler_output, "total_num_scheduled_tokens", 0) or 0), "reqs": rows}) + except Exception as e: # noqa: BLE001 + _error(e) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f1c9b79..25af1b6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,6 +70,7 @@ from vllm.v1.metrics.stats import ( SchedulerStats, ) from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats @@ -2359,6 +2360,8 @@ class Scheduler(SchedulerInterface): engine_core_outputs[0] = eco = EngineCoreOutputs() eco.scheduler_stats = stats + if _e2_alloc_log.enabled(): + _e2_alloc_log.log_step(self, scheduler_output) return engine_core_outputs def _ec_transfer_pending(self, request: Request, num_computed_tokens: int) -> bool: -- 2.43.0