src/gpu_timing.py

from __future__ import annotations

import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

import torch


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _round(value: float | None) -> float | None:
    return round(float(value), 6) if value is not None else None


@dataclass
class GpuCallbackTiming:
    """Server-side timing captured from the first line of a @spaces.GPU callback.

    Direct dispatch deliberately does not insert a browser timestamp event before
    the ZeroGPU function. Queue/allocation wait therefore remains unavailable,
    while callback-body phases are measured with one monotonic server clock.
    """

    callback_entry_perf_counter: float
    callback_entry_utc: str
    cuda_available_at_entry: bool
    cuda_initial_sync_seconds: float | None = None
    callback_entry_to_cuda_ready_seconds: float | None = None
    cuda_ready_utc: str | None = None
    cuda_ready_status: str = "not-attempted"
    cuda_ready_error: str | None = None
    dispatch_mode: str = "direct-button-to-spaces-gpu-callback"
    checkpoints: dict[str, float] = field(default_factory=dict)

    def elapsed(self) -> float:
        return max(0.0, time.perf_counter() - self.callback_entry_perf_counter)

    def checkpoint(self, name: str) -> float:
        value = self.elapsed()
        self.checkpoints[str(name)] = value
        return value

    def to_dict(self, *, body_seconds: float | None = None) -> dict[str, Any]:
        return {
            "measurement_schema": "sesa-zerogpu-callback-timing-v1",
            "dispatch_mode": self.dispatch_mode,
            "browser_click_to_callback_seconds": None,
            "browser_click_to_callback_status": "unavailable-by-design-no-frontend-timestamp",
            "zerogpu_queue_allocation_wait_seconds": None,
            "zerogpu_queue_allocation_wait_status": "unavailable-by-design-outside-decorated-callback",
            "callback_entry_utc": self.callback_entry_utc,
            "cuda_available_at_entry": self.cuda_available_at_entry,
            "cuda_initial_sync_seconds": _round(self.cuda_initial_sync_seconds),
            "callback_entry_to_cuda_ready_seconds": _round(
                self.callback_entry_to_cuda_ready_seconds
            ),
            "cuda_ready_utc": self.cuda_ready_utc,
            "cuda_ready_status": self.cuda_ready_status,
            "cuda_ready_error": self.cuda_ready_error,
            "gpu_callback_body_seconds": _round(
                self.elapsed() if body_seconds is None else body_seconds
            ),
            "measurement_boundary": "server callback entry through summary-finalization checkpoint",
            "checkpoints_seconds_from_callback_entry": {
                key: _round(value) for key, value in self.checkpoints.items()
            },
        }


def start_gpu_callback_timing(
    *, dispatch_mode: str = "direct-button-to-spaces-gpu-callback"
) -> GpuCallbackTiming:
    """Call as the first executable line inside a decorated GPU function."""

    entry = time.perf_counter()
    timing = GpuCallbackTiming(
        callback_entry_perf_counter=entry,
        callback_entry_utc=_utc_now(),
        cuda_available_at_entry=bool(torch.cuda.is_available()),
        dispatch_mode=str(dispatch_mode),
    )
    if not timing.cuda_available_at_entry:
        timing.cuda_ready_status = "cuda-unavailable"
        timing.callback_entry_to_cuda_ready_seconds = timing.elapsed()
        timing.cuda_ready_utc = _utc_now()
        return timing

    sync_started = time.perf_counter()
    try:
        torch.cuda.synchronize()
    except Exception as exc:  # diagnostics must survive a timing inspection failure
        timing.cuda_initial_sync_seconds = time.perf_counter() - sync_started
        timing.callback_entry_to_cuda_ready_seconds = timing.elapsed()
        timing.cuda_ready_utc = _utc_now()
        timing.cuda_ready_status = "initial-sync-failed"
        timing.cuda_ready_error = f"{type(exc).__name__}: {exc}"
    else:
        timing.cuda_initial_sync_seconds = time.perf_counter() - sync_started
        timing.callback_entry_to_cuda_ready_seconds = timing.elapsed()
        timing.cuda_ready_utc = _utc_now()
        timing.cuda_ready_status = "ready-after-synchronize"
    return timing


def synchronized_wall_time(operation, *, use_cuda: bool) -> tuple[Any, float, str | None]:
    """Measure wall time with CUDA synchronization around an operation.

    This is not pure kernel-active time. It is a synchronized wall window that
    includes Python/C++/ONNX Runtime work inside the operation and ensures queued
    CUDA work is complete before the end timestamp.
    """

    sync_error = None
    if use_cuda:
        try:
            torch.cuda.synchronize()
        except Exception as exc:
            sync_error = f"pre:{type(exc).__name__}: {exc}"
    started = time.perf_counter()
    result = operation()
    if use_cuda:
        try:
            torch.cuda.synchronize()
        except Exception as exc:
            suffix = f"post:{type(exc).__name__}: {exc}"
            sync_error = f"{sync_error}; {suffix}" if sync_error else suffix
    return result, max(0.0, time.perf_counter() - started), sync_error