src/core/request_stats.py

"""Per-request pipeline stat stashes — lease / CPU-subprocess / worker-dispatch.

One RequestStats object per request: pipeline entries call
``reset_request_stats()`` once, the GPU/CPU dispatch layers each fill their
slot, and the log-row builder reads the whole object at build time
(``log_blocks.build_timing``). Slot dict shapes are telemetry wire shapes —
field names inside them are load-bearing.
"""
from contextvars import ContextVar
from dataclasses import dataclass


@dataclass
class RequestStats:
    lease: dict | None = None      # gpu_runtime duration estimators
    cpu: dict | None = None        # zero_gpu local-subprocess dispatch
    dispatch: dict | None = None   # worker_pool remote-worker dispatch


_STATS: ContextVar[RequestStats | None] = ContextVar("request_stats", default=None)


def get_request_stats() -> RequestStats:
    """This request's stats object, created on first access."""
    stats = _STATS.get()
    if stats is None:
        stats = RequestStats()
        _STATS.set(stats)
    return stats


def reset_request_stats() -> None:
    """Drop all per-request stats. Called once at every pipeline entry."""
    _STATS.set(RequestStats())