src/gpu.py

"""ZeroGPU dispatch.

ZeroGPU attaches a GPU only for the duration of a decorated call, and it
**refuses to start a Space that declares no `@spaces.GPU` function at all** --
the failure is `RUNTIME_ERROR: No @spaces.GPU function detected during
startup`, which is fatal and happens before any request. So the decorated entry
point below is defined at import time, unconditionally, whether or not this
process will ever run on GPU hardware.

Off ZeroGPU -- locally, in CI, on cpu-basic -- the `spaces` package is absent
and `gpu()` is the identity, so the same code path runs inline on CPU. The
package is injected by the ZeroGPU builder and is deliberately absent from
requirements.txt: pinning it conflicts with the builder's own copy.
"""

from __future__ import annotations

import logging
import os

log = logging.getLogger("arena.gpu")

try:  # pragma: no cover - depends on the runtime, not on the code
    import spaces

    HAS_SPACES = True
except ImportError:
    spaces = None
    HAS_SPACES = False

ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))

# How long a GPU-tier forecast may hold the GPU.
#
# ZeroGPU enforces a ceiling and scales the request before checking it: asking
# for 180 was rejected as "the requested GPU duration (270s) is larger than the
# maximum allowed". 60 leaves comfortable headroom -- these models take seconds
# on an A10G, and the weights are already on local disk by the time the GPU
# call runs, because `warm()` pulls them outside it.
GPU_DURATION_S = 60


def gpu(duration: int = GPU_DURATION_S):
    """`@spaces.GPU` where that exists, the identity everywhere else."""
    def decorate(fn):
        if HAS_SPACES:
            return spaces.GPU(duration=duration)(fn)
        return fn
    return decorate


def available() -> bool:
    """Whether a GPU-tier model can actually run in this process."""
    return HAS_SPACES and ON_ZEROGPU