telemetry: clock rule, every terminal path, real cache hits, schema + drop counts

README.md

@@ -22,8 +22,10 @@ referee, GPU VMs) lives in [BlameGraph](https://github.com/AdvayMonga/BlameGraph
 - FastAPI + SSE, `session_id` threaded end to end; OpenAI shim (`/v1/completions`,
   `/v1/chat/completions`, `/v1/models`). No built-in UI
 - Sliding-window p50/p95/p99 on `/scheduler/stats`; Prometheus `/metrics` + Grafana (`monitoring/`)
-- Per-request telemetry rows (`TELEMETRY_DIR`) and the engine event timeline (`TIMELINE_DIR`):
-  one JSONL event per scheduler decision and a profiler range per phase, keyed by step id
+- Per-request telemetry rows (`TELEMETRY_DIR`: one SQLite file per run, `requests` keyed by the
+  client's `X-Trace-Id` plus a `meta` table with schema version and drop counts) and the engine
+  event timeline (`TIMELINE_DIR`): one JSONL event per scheduler decision and a profiler range per
+  phase, keyed by step id
 
 ---
 

src/inference_server/backends/base.py

@@ -47,6 +47,7 @@ class InferenceBackend(ABC):
 
     cache_adapter: "CacheManager | None" = None
     last_cache_hit_tokens: int = 0  # tokens served from cache on the most recent request
+    last_batch_cache_hit_tokens: list[int] | None = None  # per row of the last prefill_batch
 
     def set_cache_adapter(self, adapter: "CacheManager") -> None:
         """Attach a CacheManager for prefix caching. Optional."""

src/inference_server/backends/custom_torch_backend.py

@@ -536,6 +536,7 @@ def prefill_batch(self, prompts, session_ids=None):
             matched.append(m)
             suffixes.append(p[m:])
         self.last_cache_hit_tokens = matched[-1] if matched else 0
+        self.last_batch_cache_hit_tokens = list(matched)
 
         Smax = max(len(s) for s in suffixes)
         Pmax = max(matched) if matched else 0
@@ -614,6 +615,7 @@ def _prefill_batch_kernel(self, prompts):
                     self._capture_all_prefill_graphs()
                 if any(k[0] == bucket for k in self._prefill_graphs):
                     self.last_cache_hit_tokens = matched0
+                    self.last_batch_cache_hit_tokens = [matched0]
                     return [self._replay_prefill(p, bucket, matched0, shared0)]
         from inference_server.models.paged_kv_cache import PagedKVCache
         dev = self.device
@@ -626,6 +628,7 @@ def _prefill_batch_kernel(self, prompts):
             matched.append(m)
             suffixes.append(p[m:])
         self.last_cache_hit_tokens = matched[-1] if matched else 0
+        self.last_batch_cache_hit_tokens = list(matched)
 
         Smax = max(len(s) for s in suffixes)
         input_ids = torch.zeros(len(prompts), Smax, dtype=torch.long, device=dev)

src/inference_server/openai_shim.py

@@ -111,10 +111,12 @@ async def _stream(req: ScheduledRequest, tokenizer, cid: str, created: int, mode
                   prompt_tokens: int, max_tokens: int) -> AsyncGenerator[str, None]:
     """Stream one OpenAI chunk per generated token, then finish_reason + usage + [DONE]."""
     n = 0
+    ended = False
     try:
         while True:
             tok_id = await req.token_queue.get()
             if tok_id is None:
+                ended = True
                 break
             text = tokenizer.decode_token(tok_id)
             n += 1
@@ -123,6 +125,7 @@ async def _stream(req: ScheduledRequest, tokenizer, cid: str, created: int, mode
         if req.future.done() and req.future.exception():
             raise req.future.exception()  # type: ignore[misc]
     finally:
+        req.cancelled = not ended      # client went away mid-stream: the scheduler stops the row
         finish = "length" if n >= max_tokens else "stop"
         yield _chunk(cid, created, model, "", finish)
         usage_payload = {
@@ -163,6 +166,7 @@ async def _score(request: Request, body: CompletionRequest, token_ids: list[int]
 @router.post("/v1/completions")
 async def completions(body: CompletionRequest, request: Request, response: Response):
     """Text completion over the continuous-batch scheduler — the open-loop bench surface."""
+    arrival = time.perf_counter()
     scheduler = request.app.state.scheduler
     tokenizer = request.app.state.tokenizer
     loop = asyncio.get_running_loop()
@@ -190,7 +194,7 @@ async def completions(body: CompletionRequest, request: Request, response: Respo
         req = ScheduledRequest(
             token_ids=token_ids, max_tokens=body.max_tokens, session_id=session_id,
             future=loop.create_future(), token_queue=asyncio.Queue(), sampling=sampling,
-            trace_id=trace_id, turn_index=turn_index,
+            trace_id=trace_id, turn_index=turn_index, arrival_ts=arrival,
         )
         try:
             scheduler.enqueue(req)
@@ -208,7 +212,7 @@ async def completions(body: CompletionRequest, request: Request, response: Respo
     req = ScheduledRequest(
         token_ids=token_ids, max_tokens=body.max_tokens, session_id=session_id,
         future=loop.create_future(), sampling=sampling,
-        trace_id=trace_id, turn_index=turn_index,
+        trace_id=trace_id, turn_index=turn_index, arrival_ts=arrival,
     )
     try:
         generated_ids = await scheduler.submit(req)
@@ -243,11 +247,13 @@ async def _chat_stream(req: ScheduledRequest, tokenizer, cid: str, created: int,
                        prompt_tokens: int, max_tokens: int) -> AsyncGenerator[str, None]:
     """Role delta first (clients key off it), then one content delta per token, then usage."""
     n = 0
+    ended = False
     try:
         yield _chat_chunk(cid, created, model, {"role": "assistant", "content": ""}, None)
         while True:
             tok_id = await req.token_queue.get()
             if tok_id is None:
+                ended = True
                 break
             text = tokenizer.decode_token(tok_id)
             n += 1
@@ -256,6 +262,7 @@ async def _chat_stream(req: ScheduledRequest, tokenizer, cid: str, created: int,
         if req.future.done() and req.future.exception():
             raise req.future.exception()  # type: ignore[misc]
     finally:
+        req.cancelled = not ended      # client went away mid-stream: the scheduler stops the row
         yield _chat_chunk(cid, created, model, {}, "length" if n >= max_tokens else "stop")
         usage_payload = {
             "id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
@@ -269,6 +276,7 @@ async def _chat_stream(req: ScheduledRequest, tokenizer, cid: str, created: int,
 @router.post("/v1/chat/completions")
 async def chat_completions(body: ChatCompletionRequest, request: Request, response: Response):
     """Chat completion over the continuous-batch scheduler — the frontend-facing surface."""
+    arrival = time.perf_counter()
     scheduler = request.app.state.scheduler
     tokenizer = request.app.state.tokenizer
     loop = asyncio.get_running_loop()
@@ -295,7 +303,7 @@ async def chat_completions(body: ChatCompletionRequest, request: Request, respon
         req = ScheduledRequest(
             token_ids=token_ids, max_tokens=max_tokens, session_id=session_id,
             future=loop.create_future(), token_queue=asyncio.Queue(), sampling=sampling,
-            trace_id=trace_id, turn_index=turn_index,
+            trace_id=trace_id, turn_index=turn_index, arrival_ts=arrival,
         )
         try:
             scheduler.enqueue(req)
@@ -313,7 +321,7 @@ async def chat_completions(body: ChatCompletionRequest, request: Request, respon
     req = ScheduledRequest(
         token_ids=token_ids, max_tokens=max_tokens, session_id=session_id,
         future=loop.create_future(), sampling=sampling,
-        trace_id=trace_id, turn_index=turn_index,
+        trace_id=trace_id, turn_index=turn_index, arrival_ts=arrival,
     )
     try:
         generated_ids = await scheduler.submit(req)

src/inference_server/scheduler.py

@@ -46,9 +46,10 @@ class ScheduledRequest:
     future: asyncio.Future
     token_queue: asyncio.Queue | None = None
     generated: list[int] = field(default_factory=list)
-    cache_hit_tokens: int = 0   # set by scheduler after prefill
+    cache_hit_tokens: int | None = None  # set by scheduler after prefill; None = not reported
     arrival_seq: int = 0        # set by scheduler at enqueue (monotonic)
     priority: int = 0           # higher = more important; tiebreak field for policies
+    arrival_ts: float = 0.0     # perf_counter at HTTP receipt (set by the handler); enqueue_ts if 0
     enqueue_ts: float = 0.0     # set by scheduler at enqueue (perf_counter)
     first_token_ts: float = 0.0 # set when first token is produced
     admit_ts: float = 0.0       # set when the scheduler picks it up — splits TTFT into
@@ -62,6 +63,7 @@ class ScheduledRequest:
     decode_width_sum: int = 0
     decode_width_steps: int = 0
     decode_width_max: int = 0
+    cancelled: bool = False     # set by the HTTP side when the client goes away
 
     def __post_init__(self):
         if not self.trace_id:
@@ -243,13 +245,18 @@ async def stop(self) -> None:
             self._pending_cv.notify_all()
         await asyncio.get_running_loop().run_in_executor(None, self._worker.join)
         self._worker = None
+        exc = RuntimeError("scheduler stopped")
+        self._fail_inflight(exc, state="aborted")   # worker has exited: safe from this thread
+        self._fail_pending(exc, state="aborted")
         if self._telemetry is not None:
             self._telemetry.close()
         self._timeline.close()
 
     def enqueue(self, request: ScheduledRequest) -> None:
         """Synchronously enqueue. Raises QueueFullError if no capacity."""
         request.enqueue_ts = time.perf_counter()
+        if not request.arrival_ts:
+            request.arrival_ts = request.enqueue_ts
         with self._pending_cv:
             self._timeline.event("enqueue", trace_id=request.trace_id, session_id=request.session_id,
                                  prompt_tokens=len(request.token_ids), max_tokens=request.max_tokens,
@@ -280,7 +287,11 @@ def enqueue(self, request: ScheduledRequest) -> None:
 
     async def submit(self, request: ScheduledRequest) -> list[int]:
         self.enqueue(request)
-        return await request.future
+        try:
+            return await request.future
+        except asyncio.CancelledError:
+            request.cancelled = True
+            raise
 
     def stats(self) -> dict:
         with self._pending_lock:
@@ -401,26 +412,30 @@ def _step_event(self, t0: float) -> None:
     def _evict_finished(self) -> None:
         """Process each active row's current_token, evict finished rows."""
         # Walk indices high→low so popping doesn't shift earlier indices
-        to_evict: list[int] = []
+        to_evict: list[tuple[int, str]] = []
         for i, row in enumerate(self._active):
             tok = row.current_token
 
+            if row.request.cancelled:
+                to_evict.append((i, "cancelled"))
+                continue
+
             if self.backend.is_eos(tok):
-                to_evict.append(i)
+                to_evict.append((i, "ok"))
                 continue
 
             row.request.generated.append(tok)
             self._stream_token(row.request, tok)
 
             if len(row.request.generated) >= row.request.max_tokens:
-                to_evict.append(i)
+                to_evict.append((i, "ok"))
 
-        for i in reversed(to_evict):
+        for i, state in reversed(to_evict):
             row = self._active[i]
             self.policy.on_request_finished(row.request)
-            self._evict_row(i)
+            self._evict_row(i, state=state)
 
-    def _evict_row(self, idx: int, resolve: bool = True) -> None:
+    def _evict_row(self, idx: int, resolve: bool = True, state: str = "ok") -> None:
         row = self._active.pop(idx)
         self._active_kv_reserved -= len(row.request.token_ids) + row.request.max_tokens
         self.backend.kv_release(len(row.request.token_ids), row.request.max_tokens)
@@ -437,7 +452,7 @@ def _evict_row(self, idx: int, resolve: bool = True) -> None:
                     self._attention_mask[idx + 1:],
                 ], dim=0)
         if resolve:
-            self._resolve(row.request)
+            self._resolve(row.request, state)
 
     def _consume(self, request: "ScheduledRequest", plan: list) -> None:
         """Remove `request` from the policy and from this wave's plan (caller holds the lock)."""
@@ -448,7 +463,7 @@ def _consume(self, request: "ScheduledRequest", plan: list) -> None:
     # --- Phase 2: admit new requests ---
 
     def _shed_expired(self) -> None:
-        """Reject every request past the admission deadline, wherever it sits in the queue.
+        """Reject every request past the admission deadline (or cancelled), wherever it sits in the queue.
 
         Expiry used to be checked only against peek_next(). That is equivalent to shedding all
         stale work only if the head is always the oldest request, and no policy here orders that
@@ -460,14 +475,19 @@ def _shed_expired(self) -> None:
         Runs once per admission pass, over the pending set. O(pending) against a decode step that
         costs milliseconds, so the sweep is free in comparison.
         """
-        if self.max_queue_wait_s <= 0:
-            return
         now = time.perf_counter()
         # Under the same lock enqueue() holds: it mutates the policy's pending set and
         # _pending_count, so sweeping without it races an arriving request.
         with self._pending_cv:
             for req in self.policy.pending():
-                if not req.enqueue_ts or now - req.enqueue_ts <= self.max_queue_wait_s:
+                if req.cancelled:                  # client gone before admission: never run it
+                    self.policy.pick(req)
+                    self._pending_count -= 1
+                    self.policy.on_request_finished(req)
+                    self._reject(req, RuntimeError("client disconnected"), state="cancelled")
+                    continue
+                if (self.max_queue_wait_s <= 0 or not req.enqueue_ts
+                        or now - req.enqueue_ts <= self.max_queue_wait_s):
                     continue
                 waited = now - req.enqueue_ts
                 self.policy.pick(req)
@@ -606,6 +626,7 @@ def _admit_batched(self, reqs: list[ScheduledRequest], device: str) -> None:
         wave (release each reservation) rather than risk partial/leaked state."""
         try:
             self._wave_sizes[len(reqs)] = self._wave_sizes.get(len(reqs), 0) + 1
+            self.backend.last_batch_cache_hit_tokens = None   # never read a previous wave's hits
             with self._timeline.phase("prefill_batch", wave=len(reqs),
                                       tokens=sum(len(r.token_ids) for r in reqs)):
                 results = self.backend.prefill_batch([r.token_ids for r in reqs])
@@ -618,8 +639,11 @@ def _admit_batched(self, reqs: list[ScheduledRequest], device: str) -> None:
                 self.policy.on_request_finished(req)
                 self._reject(req, e)
             return
-        for req, (kv, first_token, kv_len) in zip(reqs, results):
-            req.cache_hit_tokens = 0  # batched path does a full prefill (no prefix-cache lookup)
+        hits = self.backend.last_batch_cache_hit_tokens
+        if hits is None or len(hits) != len(reqs):
+            hits = [None] * len(reqs)                  # backend did not report per-row hits
+        for req, (kv, first_token, kv_len), hit in zip(reqs, results, hits):
+            req.cache_hit_tokens = hit
             self._promote_to_decode(req, kv, kv_len, first_token, device)
             self._total_admitted += 1
 
@@ -754,13 +778,14 @@ def _stream_token(self, request: ScheduledRequest, token: int) -> None:
             return
         self._loop.call_soon_threadsafe(request.token_queue.put_nowait, token)
 
-    def _resolve(self, request: ScheduledRequest) -> None:
+    def _resolve(self, request: ScheduledRequest, state: str = "ok") -> None:
         if self._loop is None:
             return
         result = list(request.generated)
-        self._total_completed += 1
-        self._finish_row(request, "ok")
-        if request.first_token_ts > 0 and request.enqueue_ts > 0:
+        self._finish_row(request, state)
+        if state == "ok":
+            self._total_completed += 1
+        if state == "ok" and request.first_token_ts > 0 and request.enqueue_ts > 0:
             now = time.perf_counter()
             ttft = request.first_token_ts - request.enqueue_ts
             total = now - request.enqueue_ts
@@ -802,7 +827,9 @@ def _arrival_record(self, request: ScheduledRequest) -> RequestRecord:
         free = cache.free_blocks if cache is not None else None
         return RequestRecord(
             trace_id=request.trace_id, session_id=request.session_id,
-            turn_index=request.turn_index, arrival_ts=time.time(),
+            turn_index=request.turn_index,
+            # One wall/monotonic pair, back-dated to HTTP receipt: no duration mixes the two clocks.
+            arrival_ts=time.time() - (time.perf_counter() - request.arrival_ts),
             pending_depth=self._pending_count,
             active_size=len(self._active) + len(self._prefilling),
             max_batch_size=self.max_batch_size,
@@ -813,7 +840,7 @@ def _arrival_record(self, request: ScheduledRequest) -> RequestRecord:
             # session's first turn sees only OTHER sessions.
             concurrent_sessions=len(self._session_load),
             prompt_tokens=len(request.token_ids),
-            replica_age_s=request.enqueue_ts - self._start_ts,
+            replica_age_s=request.arrival_ts - self._start_ts,
         )
 
     def _finish_row(self, request: ScheduledRequest, state: str) -> None:
@@ -835,7 +862,8 @@ def _finish_row(self, request: ScheduledRequest, state: str) -> None:
                    tokens_out=len(request.generated), cache_hit_tokens=request.cache_hit_tokens,
                    decode_width_sum=request.decode_width_sum,
                    decode_width_steps=request.decode_width_steps,
-                   decode_width_max=request.decode_width_max)
+                   decode_width_max=request.decode_width_max, arrival_mono=request.arrival_ts,
+                   prefill_mode=self.prefill_mode if request.admit_ts else None)
         self._telemetry.put(rec)
 
     def _preempt_newest(self, exc: BaseException) -> None:
@@ -852,7 +880,7 @@ def _preempt_newest(self, exc: BaseException) -> None:
             f"preempted under KV pressure after {len(row.request.generated)} tokens"),
             state="preempted")
 
-    def _fail_inflight(self, exc: BaseException) -> None:
+    def _fail_inflight(self, exc: BaseException, state: str | None = None) -> None:
         """Drop everything in flight and RELEASE ITS KV, leaving the worker able to continue.
 
         The blocks matter: the old _fail_all cleared _active without routing through
@@ -868,7 +896,7 @@ def _fail_inflight(self, exc: BaseException) -> None:
                 except Exception:
                     logger.exception("failed to release KV for a dropped row")
             self.policy.on_request_finished(row.request)
-            self._reject(row.request, exc)
+            self._reject(row.request, exc, state=state)
         for prow in self._prefilling:
             self.backend.kv_release(len(prow.request.token_ids), prow.request.max_tokens)
             free = getattr(prow.partial_kv, "free_all", None)
@@ -878,12 +906,22 @@ def _fail_inflight(self, exc: BaseException) -> None:
                 except Exception:
                     logger.exception("failed to release partial prefill KV")
             self.policy.on_request_finished(prow.request)
-            self._reject(prow.request, exc)
+            self._reject(prow.request, exc, state=state)
         self._active.clear()
         self._prefilling.clear()
         self._active_kv_reserved = 0
         self._batched_kv = None
         self._attention_mask = None
 
+    def _fail_pending(self, exc: BaseException, state: str | None = None) -> None:
+        """Reject everything still queued, so no caller waits forever and every request gets a row."""
+        with self._pending_cv:
+            for req in self.policy.pending():
+                self.policy.pick(req)
+                self._pending_count -= 1
+                self.policy.on_request_finished(req)
+                self._reject(req, exc, state=state)
+
     def _fail_all(self, exc: BaseException) -> None:
         self._fail_inflight(exc)
+        self._fail_pending(exc)