engines/zero_gpu.py

from __future__ import annotations

from typing import Any, Callable


try:
    import spaces as _spaces  # type: ignore[import-not-found]
except Exception:  # pragma: no cover - local fallback when spaces is not installed
    class _SpacesShim:
        @staticmethod
        def GPU(
            fn: Callable[..., Any] | None = None,
            *,
            duration: int | Callable[..., int] | None = None,
            size: str | None = None,
        ):
            del duration, size
            if fn is not None:
                return fn

            def _decorator(inner: Callable[..., Any]) -> Callable[..., Any]:
                return inner

            return _decorator

    _spaces = _SpacesShim()


spaces = _spaces


def live_request_duration(request_payload: dict[str, Any]) -> int:
    context_length = int(request_payload.get("context_length") or 0)
    mode = str(request_payload.get("mode") or "dense")
    custom_prompt = bool(str(request_payload.get("custom_prompt") or "").strip())

    if context_length <= 512:
        duration = 45
    elif context_length <= 1024:
        duration = 75
    elif context_length <= 4096:
        duration = 150
    else:
        duration = 240

    if mode != "dense":
        duration += 30
    if not custom_prompt:
        duration += 15

    return max(45, min(duration, 300))