evaluation/runtime.py

"""Shared evaluation runtime: device selection, ZeroGPU bypass, CPU thread hints."""

from __future__ import annotations

import importlib.util
import os
import sys
import types

_cpu_thread_settings_applied = False
_torchcodec_patch_applied = False


def _force_transformers_torchcodec_unavailable() -> None:
    """Tell transformers to use librosa/soundfile for audio, not torchcodec."""
    try:
        import transformers.utils.import_utils as import_utils

        import_utils._torchcodec_available = False

        def _is_torchcodec_available() -> bool:
            return False

        import_utils.is_torchcodec_available = _is_torchcodec_available
    except Exception:
        pass


def disable_broken_torchcodec() -> None:
    """
    Transformers prefers torchcodec for ``load_audio`` when the package is installed,
    but Hub Job images often lack FFmpeg shared libraries (libavutil.so.*).

    If torchcodec cannot load, mark it unavailable so transformers falls back to
    librosa / soundfile.
    """
    global _torchcodec_patch_applied
    if _torchcodec_patch_applied:
        return
    _torchcodec_patch_applied = True

    force_off = os.environ.get("FFASR_DISABLE_TORCHCODEC", "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )

    if force_off or importlib.util.find_spec("torchcodec") is None:
        if force_off and importlib.util.find_spec("torchcodec") is not None:
            print(
                "[ffasr] FFASR_DISABLE_TORCHCODEC=1: using librosa/soundfile for audio "
                "(torchcodec disabled on Hub Jobs).",
                flush=True,
            )
        _force_transformers_torchcodec_unavailable()
        return

    try:
        from torchcodec.decoders import AudioDecoder  # noqa: F401
    except (OSError, RuntimeError, ImportError, AttributeError):
        print(
            "[ffasr] torchcodec native libraries unavailable; "
            "transformers will use librosa/soundfile for audio.",
            flush=True,
        )
        _force_transformers_torchcodec_unavailable()
        stub = types.ModuleType("torchcodec")
        stub.__dict__["decoders"] = types.ModuleType("torchcodec.decoders")
        sys.modules["torchcodec"] = stub
        sys.modules["torchcodec.decoders"] = stub.__dict__["decoders"]
    except Exception:
        _force_transformers_torchcodec_unavailable()


def patch_transformers_load_audio_for_paths() -> None:
    """
    Route ``transformers.audio_utils.load_audio`` for file paths through soundfile
    (no torchcodec / FFmpeg). Safe to call before loading user custom scripts.
    """
    try:
        import numpy as np
        import transformers.audio_utils as audio_utils
        from backends._audio_utils import load_wav_mono

        _orig = audio_utils.load_audio

        def _load_audio(audio, sampling_rate=16000, timeout=None):
            if isinstance(audio, str):
                return load_wav_mono(audio, sampling_rate=int(sampling_rate))
            if isinstance(audio, np.ndarray):
                return audio
            return _orig(audio, sampling_rate=sampling_rate, timeout=timeout)

        audio_utils.load_audio = _load_audio
    except Exception:
        pass


def zerogpu_disabled() -> bool:
    """When True, never wrap evaluation in ``spaces.GPU`` even if ``spaces`` is installed."""
    return os.environ.get("FFASR_DISABLE_ZEROGPU", "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )


def spaces_available() -> bool:
    return importlib.util.find_spec("spaces") is not None


def use_spaces_gpu_decorator() -> bool:
    """Use ZeroGPU decorator only when ``spaces`` is present and not explicitly disabled."""
    return spaces_available() and not zerogpu_disabled()


def resolve_eval_devices() -> tuple[str, int]:
    """
    Return ``(device_str, device_int)`` for backends / HF pipeline.

    ``FFASR_DEVICE`` (default ``auto``): ``cpu`` | ``cuda`` | ``auto``.
    ``auto`` uses CUDA when available, else CPU.
    """
    choice = (os.environ.get("FFASR_DEVICE", "auto") or "auto").strip().lower()
    if choice not in ("auto", "cpu", "cuda"):
        choice = "auto"

    import torch

    has_cuda = torch.cuda.is_available()
    if choice == "cpu":
        return "cpu", -1
    if choice == "cuda":
        if not has_cuda:
            raise RuntimeError("FFASR_DEVICE=cuda but torch.cuda.is_available() is False.")
        return "cuda", 0
    # auto
    if has_cuda:
        return "cuda", 0
    return "cpu", -1


def apply_cpu_thread_settings_once() -> None:
    """
    Apply optional torch thread limits from env (once per process).

    - ``FFASR_TORCH_NUM_THREADS`` — passed to ``torch.set_num_threads`` if a positive int.
    - ``FFASR_TORCH_NUM_INTEROP_THREADS`` — ``torch.set_num_interop_threads`` if set.

    For BLAS/OpenMP, operators typically set ``OMP_NUM_THREADS`` / ``MKL_NUM_THREADS``
    in the Space environment (see docs); those are read by the native libs at load time.
    """
    global _cpu_thread_settings_applied
    if _cpu_thread_settings_applied:
        return
    _cpu_thread_settings_applied = True

    import torch

    raw = os.environ.get("FFASR_TORCH_NUM_THREADS", "").strip()
    if raw.isdigit():
        n = int(raw)
        if n > 0:
            try:
                torch.set_num_threads(n)
            except Exception:
                pass

    raw_i = os.environ.get("FFASR_TORCH_NUM_INTEROP_THREADS", "").strip()
    if raw_i.isdigit():
        ni = int(raw_i)
        if ni > 0:
            try:
                torch.set_num_interop_threads(ni)
            except Exception:
                pass