diff --git a/openrag/services/workers/parsers/whisper_workers.py b/openrag/services/workers/parsers/whisper_workers.py index 990602d9a..249df7694 100644 --- a/openrag/services/workers/parsers/whisper_workers.py +++ b/openrag/services/workers/parsers/whisper_workers.py @@ -20,7 +20,23 @@ def _whisper_num_gpus(config) -> float: - return config.loader.local_whisper.whisper_num_gpus if torch.cuda.is_available() else 0 + """Return Whisper's Ray GPU reservation, falling back to CUDA detection. + + Mirrors Marker's GPU selection (see ``_marker_num_gpus``): the + ``WhisperPool`` actor runs with ``num_gpus=0``, so Ray hides CUDA in its + process and ``torch.cuda.is_available()`` returns False even on a GPU + node — which would silently pin the ``WhisperActor``s to CPU. Ask whether + the Ray cluster has GPU capacity instead, keeping the CUDA check only as a + fallback for when Ray cannot report cluster resources yet. + """ + requested_gpus = config.loader.local_whisper.whisper_num_gpus + if requested_gpus <= 0: + return 0 + try: + return requested_gpus if ray.cluster_resources().get("GPU", 0) > 0 else 0 + except Exception as exc: + logger.warning("Failed to query Ray cluster GPU resources; falling back to CUDA check", error=str(exc)) + return requested_gpus if torch.cuda.is_available() else 0 def whisper_actor_options(config) -> dict[str, float | int]: diff --git a/tests/unit/services/workers/parsers/test_whisper_workers.py b/tests/unit/services/workers/parsers/test_whisper_workers.py new file mode 100644 index 000000000..33bd6b78d --- /dev/null +++ b/tests/unit/services/workers/parsers/test_whisper_workers.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from services.workers.parsers import whisper_workers + + +def _config(whisper_num_gpus: float = 0.25): + """Build the minimal config shape consumed by Whisper GPU selection.""" + return SimpleNamespace(loader=SimpleNamespace(local_whisper=SimpleNamespace(whisper_num_gpus=whisper_num_gpus))) + + +def test_whisper_num_gpus_uses_ray_cluster_resources_when_cuda_is_hidden(monkeypatch): + """Whisper should request GPUs from Ray even when local CUDA is hidden. + + The ``WhisperPool`` actor runs with ``num_gpus=0``, so Ray hides CUDA in + its process; the GPU request must come from Ray cluster resources, not + from ``torch.cuda.is_available()``. + """ + monkeypatch.setattr(whisper_workers.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(whisper_workers.ray, "cluster_resources", lambda: {"GPU": 1.0}) + + assert whisper_workers._whisper_num_gpus(_config()) == 0.25 + + +def test_whisper_num_gpus_is_zero_when_disabled_by_config(monkeypatch): + """A non-positive configured request keeps Whisper off the GPU.""" + monkeypatch.setattr(whisper_workers.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(whisper_workers.ray, "cluster_resources", lambda: {"GPU": 1.0}) + + assert whisper_workers._whisper_num_gpus(_config(whisper_num_gpus=0)) == 0 + + +def test_whisper_num_gpus_falls_back_to_cuda_when_ray_lookup_fails(monkeypatch): + """If Ray cannot report resources, fall back to the local CUDA check.""" + + def _boom(): + raise RuntimeError("ray not ready") + + monkeypatch.setattr(whisper_workers.ray, "cluster_resources", _boom) + monkeypatch.setattr(whisper_workers.torch.cuda, "is_available", lambda: True) + assert whisper_workers._whisper_num_gpus(_config()) == 0.25 + + monkeypatch.setattr(whisper_workers.torch.cuda, "is_available", lambda: False) + assert whisper_workers._whisper_num_gpus(_config()) == 0