hackathon_advisor/zerogpu.py

from __future__ import annotations

from collections.abc import Callable
from typing import ParamSpec, TypeVar

from hackathon_advisor.config import bool_env, int_env

P = ParamSpec("P")
R = TypeVar("R")


DEFAULT_GPU_DURATION_SECONDS = 60
MAX_GPU_DURATION_SECONDS = 120


def zero_gpu_enabled() -> bool:
    return bool_env("ADVISOR_ZERO_GPU")


def gpu_device() -> str:
    """torch device for the GPU path: 'cuda' under ZeroGPU, else 'local' (auto-resolved at load)."""
    return "cuda" if zero_gpu_enabled() else "local"


def zero_gpu_duration_seconds() -> int:
    return int_env(
        "ADVISOR_ZERO_GPU_DURATION",
        DEFAULT_GPU_DURATION_SECONDS,
        minimum=1,
        maximum=MAX_GPU_DURATION_SECONDS,
    )


def gpu_task(function: Callable[P, R]) -> Callable[P, R]:
    if not zero_gpu_enabled():
        return function
    try:
        import spaces
    except ImportError as error:
        raise RuntimeError(
            "ADVISOR_ZERO_GPU=1 requires the Hugging Face `spaces` package. "
            "Install runtime requirements before enabling ZeroGPU."
        ) from error
    return spaces.GPU(duration=zero_gpu_duration_seconds())(function)


QUOTA_ERROR_HINTS = ("quota", "gpu task aborted", "no gpu", "exceeded", "gpu is not available")


def is_gpu_quota_error(error: BaseException) -> bool:
    """Heuristically detect a ZeroGPU allocation/quota failure so the caller can fall back to
    a CPU run. ZeroGPU raises before the wrapped function body executes, so this is checked
    against the exception that surfaces from the first pull of the GPU generator."""
    name = type(error).__name__.lower()
    if "quota" in name or "gpu" in name:
        return True
    message = str(error).lower()
    return any(hint in message for hint in QUOTA_ERROR_HINTS)