diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py index 255d6fe4f..66f04f296 100644 --- a/openrag/services/workers/parsers/marker_workers.py +++ b/openrag/services/workers/parsers/marker_workers.py @@ -25,7 +25,15 @@ def _marker_num_gpus(config) -> float: - return config.loader.marker_num_gpus if torch.cuda.is_available() else 0 + """Return the configured Marker GPU request when Ray reports GPU capacity.""" + requested_gpus = config.loader.marker_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_marker_workers.py b/tests/unit/services/workers/parsers/test_marker_workers.py new file mode 100644 index 000000000..7ecba70f4 --- /dev/null +++ b/tests/unit/services/workers/parsers/test_marker_workers.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from services.workers.parsers import marker_workers + + +def _config(marker_num_gpus: float = 0.25): + """Build the minimal config shape consumed by Marker GPU selection.""" + return SimpleNamespace(loader=SimpleNamespace(marker_num_gpus=marker_num_gpus)) + + +def test_marker_num_gpus_uses_ray_cluster_resources_when_cuda_is_hidden(monkeypatch): + """Marker should request GPUs from Ray even when local CUDA is hidden.""" + monkeypatch.setattr(marker_workers.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(marker_workers.ray, "cluster_resources", lambda: {"GPU": 1.0}) + + assert marker_workers._marker_num_gpus(_config()) == 0.25 diff --git a/tests/unit/services/workers/test_dispatcher.py b/tests/unit/services/workers/test_dispatcher.py index 7da62c7bb..046fd4d19 100644 --- a/tests/unit/services/workers/test_dispatcher.py +++ b/tests/unit/services/workers/test_dispatcher.py @@ -254,9 +254,13 @@ async def test_delete_file_cleans_database_before_vector_store() -> None: ) call_order = [] - workspace_repo.remove_file_from_all_workspaces = AsyncMock(side_effect=lambda *a, **k: call_order.append("workspace")) + workspace_repo.remove_file_from_all_workspaces = AsyncMock( + side_effect=lambda *a, **k: call_order.append("workspace") + ) document_repo.remove_file_from_partition = AsyncMock(side_effect=lambda *a, **k: call_order.append("document")) - vector_store.query_ids_by_filter = AsyncMock(return_value=["1", "2"], side_effect=lambda *a, **k: call_order.append("query") or ["1", "2"]) + vector_store.query_ids_by_filter = AsyncMock( + return_value=["1", "2"], side_effect=lambda *a, **k: call_order.append("query") or ["1", "2"] + ) vector_store.delete = AsyncMock(side_effect=lambda *a, **k: call_order.append("delete") or None) await dispatcher.delete_file("file-1", "tenant-a")