Reverb/open3dforge

1"""2ZeroGPU quota tracking.3 4HF Pro gives 1500s (25 min) of H200 time per day. There's no official Python5API to query remaining quota directly — we track it locally per-call.6 7This module:8- Records GPU time consumed via the `@spaces.GPU` decorated functions9- Provides estimates for upcoming operations10- Resets daily (per-day file on disk)11 12Note: this is approximate. The authoritative source is HF's quota error if13you go over. Our tracking is for UX (showing "~18/25 min used today").14"""15 16from __future__ import annotations17 18import json19import time20from dataclasses import dataclass21from pathlib import Path22 23from .workspace import WORKSPACE24 25QUOTA_FILE = WORKSPACE / "quota_log.json"26DAILY_QUOTA_SECONDS = 1500   # 25 minutes for HF Pro27OVERAGE_RATE_PER_SECOND = 1.0 / 600.0   # $1 per 600s (10 min)28 29 30@dataclass31class QuotaState:32    date: str                # YYYY-MM-DD (UTC)33    used_seconds: float34    operations: list[dict]   # log of recent operations35 36    def remaining_seconds(self) -> float:37        return max(0.0, DAILY_QUOTA_SECONDS - self.used_seconds)38 39    def usage_fraction(self) -> float:40        return min(1.0, self.used_seconds / DAILY_QUOTA_SECONDS)41 42    def overage_cost_usd(self) -> float:43        overage = max(0.0, self.used_seconds - DAILY_QUOTA_SECONDS)44        return overage * OVERAGE_RATE_PER_SECOND45 46 47def _today_utc() -> str:48    return time.strftime("%Y-%m-%d", time.gmtime())49 50 51def _load_or_new() -> QuotaState:52    today = _today_utc()53    if QUOTA_FILE.exists():54        try:55            with QUOTA_FILE.open("r") as f:56                data = json.load(f)57            if data.get("date") == today:58                return QuotaState(59                    date=data["date"],60                    used_seconds=float(data["used_seconds"]),61                    operations=data.get("operations", []),62                )63        except (json.JSONDecodeError, KeyError, ValueError):64            pass65    # Fresh day or corrupted file66    return QuotaState(date=today, used_seconds=0.0, operations=[])67 68 69def _save(state: QuotaState) -> None:70    with QUOTA_FILE.open("w") as f:71        json.dump(72            {73                "date": state.date,74                "used_seconds": state.used_seconds,75                "operations": state.operations[-50:],  # keep last 5076            },77            f,78            indent=2,79        )80 81 82def get_state() -> QuotaState:83    """Get current quota state, refreshed for today."""84    return _load_or_new()85 86 87def record_usage(operation: str, seconds: float) -> QuotaState:88    """Record a completed GPU operation. Returns updated state."""89    state = _load_or_new()90    state.used_seconds += seconds91    state.operations.append({92        "op": operation,93        "seconds": round(seconds, 2),94        "timestamp": time.time(),95    })96    _save(state)97    return state98 99 100# ---------------------------------------------------------------------------101# Per-operation estimates (used by UI to warn before expensive operations)102# ---------------------------------------------------------------------------103 104ESTIMATES = {105    # Stage 1106    "generate_trellis2_fast": 30,107    "generate_trellis2_balanced": 60,108    "generate_trellis2_hero": 90,109    "generate_hunyuan3d": 60,110    # Stage 2 — baking (nvdiffrast is fast)111    "bake_normal_2k": 5,112    "bake_normal_4k": 12,113    "bake_albedo": 2,114    "bake_materials": 3,115    "bake_ao_fast": 3,116    "bake_ao_standard": 10,117    "bake_ao_high": 30,118    # Stage 2 — optional119    "inpaint_sdxl": 30,120    # Stage 3121    "auto_rig": 40,122}123 124 125def estimate(operation: str) -> int:126    """Get the typical GPU duration for an operation, in seconds.127 128    Used to:129      - set `@spaces.GPU(duration=N)` correctly130      - show cost warnings in the UI before triggering an operation131    """132    return ESTIMATES.get(operation, 60)133 134 135def format_status() -> str:136    """One-line quota summary for the UI status bar."""137    s = get_state()138    used_min = s.used_seconds / 60139    total_min = DAILY_QUOTA_SECONDS / 60140    remaining_min = s.remaining_seconds() / 60141 142    if s.used_seconds >= DAILY_QUOTA_SECONDS:143        cost = s.overage_cost_usd()144        return f"⚠️ Quota: {used_min:.1f}/{total_min:.0f} min (overage: ${cost:.2f})"145    elif s.usage_fraction() > 0.8:146        return f"⚠️ Quota: {used_min:.1f}/{total_min:.0f} min ({remaining_min:.1f} min left)"147    else:148        return f"Quota: {used_min:.1f}/{total_min:.0f} min H200 today"149