File size: 5,267 Bytes
81ba775 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | 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
|