recitation_engine/gpu_profiler.py

"""
GPU Lease Profiler for tracking ZeroGPU usage.

Tracks all GPU lease calls and provides summary statistics.
"""
from dataclasses import dataclass, field
from typing import List
import time


@dataclass
class LeaseRecord:
    """Record of a single GPU lease."""
    name: str
    requested_duration: float  # Lease duration requested (seconds)
    actual_duration: float     # Actual runtime (seconds)
    timestamp: float           # When the lease started
    details: str = ""          # Optional details (e.g., segment count)


class GPUProfiler:
    """Singleton profiler for tracking GPU lease usage."""

    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._records: List[LeaseRecord] = []
            cls._instance._start_time: float = 0
        return cls._instance

    def reset(self):
        """Reset profiler for a new processing run."""
        self._records = []
        self._start_time = time.time()

    def record_lease(
        self,
        name: str,
        requested_duration: float,
        actual_duration: float,
        details: str = ""
    ):
        """Record a GPU lease usage."""
        self._records.append(LeaseRecord(
            name=name,
            requested_duration=requested_duration,
            actual_duration=actual_duration,
            timestamp=time.time(),
            details=details,
        ))

    def get_summary(self) -> str:
        """Generate a summary of GPU lease usage."""
        if not self._records:
            return "[GPU PROFILE] No GPU leases recorded"

        total_requested = sum(r.requested_duration for r in self._records)
        total_actual = sum(r.actual_duration for r in self._records)
        total_utilization = (total_actual / total_requested * 100) if total_requested > 0 else 0

        lines = [
            "",
            "=" * 70,
            "GPU LEASE PROFILING SUMMARY",
            "=" * 70,
            f"{'#':<3} {'Function':<25} {'Requested':>10} {'Actual':>10} {'Util %':>8} {'Details':<15}",
            "-" * 70,
        ]

        for i, record in enumerate(self._records, 1):
            util = (record.actual_duration / record.requested_duration * 100) if record.requested_duration > 0 else 0
            lines.append(
                f"{i:<3} {record.name:<25} {record.requested_duration:>9.1f}s {record.actual_duration:>9.2f}s {util:>7.1f}% {record.details:<15}"
            )

        lines.extend([
            "-" * 70,
            f"{'TOTAL':<3} {'':<25} {total_requested:>9.1f}s {total_actual:>9.2f}s {total_utilization:>7.1f}%",
            "",
            f"Total GPU leases: {len(self._records)}",
            f"Total lease time requested: {total_requested:.1f}s",
            f"Total actual GPU time: {total_actual:.2f}s",
            f"Overall utilization: {total_utilization:.1f}%",
            "=" * 70,
            "",
        ])

        return "\n".join(lines)

    def print_summary(self):
        """Print the summary to stdout."""
        print(self.get_summary())


# Module-level convenience functions
_profiler = GPUProfiler()


def reset_profiler():
    """Reset the profiler for a new run."""
    _profiler.reset()


def record_gpu_lease(name: str, requested_duration: float, actual_duration: float, details: str = ""):
    """Record a GPU lease usage."""
    _profiler.record_lease(name, requested_duration, actual_duration, details)


def get_gpu_summary() -> str:
    """Get the profiling summary string."""
    return _profiler.get_summary()


def print_gpu_summary():
    """Print the profiling summary."""
    _profiler.print_summary()