From f732e77dc8fa392aaeff9ddc2223b201e026fbfe Mon Sep 17 00:00:00 2001
From: andyne13 <aneugebauer@linagora.com>
Date: Thu, 25 Jun 2026 19:44:16 +0200
Subject: [PATCH] fix(docling): reserve the GPU from Ray cluster resources, not
 local CUDA (#580)

_docling_num_gpus ran inside the DoclingPool actor (scheduled with no GPU),
so torch.cuda.is_available() was False there, it reserved 0 GPUs for the
worker, and Docling silently parsed on CPU (~5-10x slower than marker; no
DoclingWorker on nvidia-smi). Mirror _marker_num_gpus: query the Ray
cluster for GPU capacity, with local CUDA only as a startup fallback.

Tests cover cluster-has-GPU/CUDA-hidden -> reserves, no cluster GPU -> 0,
and not-requested -> 0.
---
 .../workers/parsers/docling_workers.py        | 17 +++++++++-
 .../workers/parsers/test_docling_workers.py   | 32 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/services/workers/parsers/test_docling_workers.py

diff --git a/openrag/services/workers/parsers/docling_workers.py b/openrag/services/workers/parsers/docling_workers.py
index 1716516b4..595025444 100644
--- a/openrag/services/workers/parsers/docling_workers.py
+++ b/openrag/services/workers/parsers/docling_workers.py
@@ -37,7 +37,22 @@
 
 
 def _docling_num_gpus(config) -> float:
-    return config.loader.docling_num_gpus if torch.cuda.is_available() else 0
+    """Return Docling's Ray GPU reservation, falling back to CUDA detection.
+
+    Must mirror Marker: ``_docling_num_gpus`` runs inside the ``DoclingPool``
+    actor, which Ray schedules with no GPU — so ``torch.cuda.is_available()`` is
+    False there (``CUDA_VISIBLE_DEVICES=""``) and a naive local check reserves 0
+    GPUs for the worker, leaving Docling stuck on CPU. Query the Ray *cluster*
+    for GPU capacity instead, with local CUDA only as a startup fallback.
+    """
+    requested_gpus = config.loader.docling_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
 
 
 @ray.remote
diff --git a/tests/unit/services/workers/parsers/test_docling_workers.py b/tests/unit/services/workers/parsers/test_docling_workers.py
new file mode 100644
index 000000000..d4d7649ac
--- /dev/null
+++ b/tests/unit/services/workers/parsers/test_docling_workers.py
@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+from services.workers.parsers import docling_workers
+
+
+def _config(docling_num_gpus: float = 0.25):
+    """Build the minimal config shape consumed by Docling GPU selection."""
+    return SimpleNamespace(loader=SimpleNamespace(docling_num_gpus=docling_num_gpus))
+
+
+def test_docling_num_gpus_uses_ray_cluster_resources_when_cuda_is_hidden(monkeypatch):
+    """Docling must request a GPU from Ray even when local CUDA is hidden in the
+    pool process (CUDA_VISIBLE_DEVICES=""), otherwise it silently parses on CPU."""
+    monkeypatch.setattr(docling_workers.torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(docling_workers.ray, "cluster_resources", lambda: {"GPU": 1.0})
+
+    assert docling_workers._docling_num_gpus(_config()) == 0.25
+
+
+def test_docling_num_gpus_zero_when_cluster_has_no_gpu(monkeypatch):
+    monkeypatch.setattr(docling_workers.torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(docling_workers.ray, "cluster_resources", dict)
+
+    assert docling_workers._docling_num_gpus(_config()) == 0
+
+
+def test_docling_num_gpus_zero_when_not_requested(monkeypatch):
+    monkeypatch.setattr(docling_workers.ray, "cluster_resources", lambda: {"GPU": 1.0})
+
+    assert docling_workers._docling_num_gpus(_config(docling_num_gpus=0)) == 0