| from __future__ import annotations |
|
|
| import contextvars |
| import os |
| import time |
| from collections import deque |
| from contextlib import asynccontextmanager, contextmanager |
| from dataclasses import dataclass, field |
| from functools import wraps |
| from typing import Any |
|
|
| PROFILING_ENABLED = os.environ.get("GRADIO_PROFILING", "").strip() in ("1", "true") |
|
|
|
|
| @dataclass |
| class RequestTrace: |
| event_id: str | None = None |
| fn_name: str | None = None |
| session_hash: str | None = None |
| timestamp: float = field(default_factory=time.time) |
|
|
| queue_wait_ms: float = 0.0 |
| preprocess_ms: float = 0.0 |
| fn_call_ms: float = 0.0 |
| postprocess_ms: float = 0.0 |
| streaming_diff_ms: float = 0.0 |
| total_ms: float = 0.0 |
| n_iterations: int = 0 |
| upload_ms: float = 0.0 |
| preprocess_move_to_cache_ms: float = 0.0 |
| preprocess_format_image_ms: float = 0.0 |
| postprocess_save_img_array_to_cache_ms: float = 0.0 |
| preprocess_audio_from_file_ms: float = 0.0 |
| postprocess_save_audio_to_cache_ms: float = 0.0 |
| preprocess_video_ms: float = 0.0 |
| postprocess_video_convert_video_to_playable_mp4_ms: float = 0.0 |
| postprocess_update_state_in_config_ms: float = 0.0 |
| postprocess_move_to_cache_ms: float = 0.0 |
| postprocess_video_ms: float = 0.0 |
| postprocess_save_pil_to_cache_ms: float = 0.0 |
| postprocess_save_bytes_to_cache_ms: float = 0.0 |
| save_file_to_cache_ms: float = 0.0 |
|
|
| def set_phase(self, name: str, duration_ms: float): |
| attr = f"{name}_ms" |
| if hasattr(self, attr): |
| |
| setattr(self, attr, getattr(self, attr) + duration_ms) |
| if name == "total": |
| self.n_iterations += 1 |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "event_id": self.event_id, |
| "fn_name": self.fn_name, |
| "session_hash": self.session_hash, |
| "timestamp": self.timestamp, |
| "queue_wait_ms": self.queue_wait_ms, |
| "preprocess_ms": self.preprocess_ms, |
| "fn_call_ms": self.fn_call_ms, |
| "postprocess_ms": self.postprocess_ms, |
| "streaming_diff_ms": self.streaming_diff_ms, |
| "total_ms": self.total_ms, |
| "n_iterations": self.n_iterations, |
| "preprocess_move_to_cache_ms": self.preprocess_move_to_cache_ms, |
| "preprocess_format_image_ms": self.preprocess_format_image_ms, |
| "postprocess_save_img_array_to_cache_ms": self.postprocess_save_img_array_to_cache_ms, |
| "preprocess_audio_from_file_ms": self.preprocess_audio_from_file_ms, |
| "postprocess_save_audio_to_cache_ms": self.postprocess_save_audio_to_cache_ms, |
| "preprocess_video_ms": self.preprocess_video_ms, |
| "postprocess_video_convert_video_to_playable_mp4_ms": self.postprocess_video_convert_video_to_playable_mp4_ms, |
| "postprocess_update_state_in_config_ms": self.postprocess_update_state_in_config_ms, |
| "postprocess_move_to_cache_ms": self.postprocess_move_to_cache_ms, |
| "postprocess_video_ms": self.postprocess_video_ms, |
| "postprocess_save_pil_to_cache_ms": self.postprocess_save_pil_to_cache_ms, |
| "postprocess_save_bytes_to_cache_ms": self.postprocess_save_bytes_to_cache_ms, |
| "save_file_to_cache_ms": self.save_file_to_cache_ms, |
| } |
|
|
|
|
| _current_trace: contextvars.ContextVar[RequestTrace | None] = contextvars.ContextVar( |
| "_current_trace", default=None |
| ) |
|
|
|
|
| def get_current_trace() -> RequestTrace | None: |
| return _current_trace.get() |
|
|
|
|
| def set_current_trace(trace: RequestTrace) -> contextvars.Token: |
| return _current_trace.set(trace) |
|
|
|
|
| @asynccontextmanager |
| async def trace_phase(name: str): |
| """Async context manager that records timing for a named phase into the current trace.""" |
| trace = _current_trace.get() |
| if trace is None: |
| yield |
| return |
| start = time.monotonic() |
| try: |
| yield |
| finally: |
| duration_ms = (time.monotonic() - start) * 1000 |
| trace.set_phase(name, duration_ms) |
|
|
|
|
| @contextmanager |
| def trace_phase_sync(name: str): |
| """Context manager that records timing for a named phase into the current trace.""" |
| trace = _current_trace.get() |
| if trace is None: |
| yield |
| return |
| start = time.monotonic() |
| try: |
| yield |
| finally: |
| duration_ms = (time.monotonic() - start) * 1000 |
| trace.set_phase(name, duration_ms) |
|
|
|
|
| def traced(phase): |
| if not PROFILING_ENABLED: |
| return lambda f: f |
|
|
| def _factory(f): |
| @wraps(f) |
| async def wrapper(*args, **kwargs): |
| async with trace_phase(phase): |
| return await f(*args, **kwargs) |
|
|
| return wrapper |
|
|
| return _factory |
|
|
|
|
| def traced_sync(phase): |
| if not PROFILING_ENABLED: |
| return lambda f: f |
|
|
| def _factory(f): |
| @wraps(f) |
| def wrapper(*args, **kwargs): |
| with trace_phase_sync(phase): |
| return f(*args, **kwargs) |
|
|
| return wrapper |
|
|
| return _factory |
|
|
|
|
| class TraceCollector: |
| def __init__(self, maxlen: int = 100_000): |
| self._traces: deque[RequestTrace] = deque(maxlen=maxlen) |
|
|
| def add(self, trace: RequestTrace): |
| self._traces.append(trace) |
|
|
| def get_all(self, last_n: int | None = None) -> list[dict[str, Any]]: |
| traces = list(self._traces) |
| if last_n is not None: |
| traces = traces[-last_n:] |
| return [t.to_dict() for t in traces] |
|
|
| def get_summary(self) -> dict[str, Any]: |
| if not self._traces: |
| return {"count": 0, "phases": {}} |
|
|
| import numpy as np |
|
|
| prediction_traces = [ |
| t for t in self._traces if t.fn_name != "gradio_file_upload" |
| ] |
| upload_traces = [t for t in self._traces if t.fn_name == "gradio_file_upload"] |
|
|
| phases = [ |
| "queue_wait", |
| "preprocess", |
| "fn_call", |
| "postprocess", |
| "streaming_diff", |
| "total", |
| ] |
|
|
| def _percentiles(arr): |
| return { |
| "p50": float(np.percentile(arr, 50)), |
| "p90": float(np.percentile(arr, 90)), |
| "p95": float(np.percentile(arr, 95)), |
| "p99": float(np.percentile(arr, 99)), |
| "mean": float(np.mean(arr)), |
| "min": float(np.min(arr)), |
| "max": float(np.max(arr)), |
| } |
|
|
| result: dict[str, Any] = { |
| "count": len(prediction_traces), |
| "phases": {}, |
| } |
| for phase in phases: |
| values = [getattr(t, f"{phase}_ms") for t in prediction_traces] |
| if values: |
| result["phases"][phase] = _percentiles(np.array(values)) |
| else: |
| result["phases"][phase] = { |
| "p50": 0.0, |
| "p90": 0.0, |
| "p95": 0.0, |
| "p99": 0.0, |
| "mean": 0.0, |
| "min": 0.0, |
| "max": 0.0, |
| } |
|
|
| if upload_traces: |
| upload_values = [t.upload_ms for t in upload_traces] |
| result["upload"] = { |
| "count": len(upload_traces), |
| **_percentiles(np.array(upload_values)), |
| } |
|
|
| return result |
|
|
| def clear(self): |
| self._traces.clear() |
|
|
|
|
| |
| collector = TraceCollector() |
|
|
|
|
| if not PROFILING_ENABLED: |
| |
|
|
| @asynccontextmanager |
| async def trace_phase(name: str): |
| yield |
|
|
| @contextmanager |
| def trace_phase_sync(name: str): |
| yield |
|
|
