perf(serverless): skip torch fitness checks unless loaded

runpod/_health/gpu.py

@@ -28,6 +28,13 @@
 MAX_ERROR_MESSAGES = 10
 
 
+class GpuAllocationError(RuntimeError):
+    """gpu_test reached the GPUs and a device failed its allocation test.
+
+    The nvidia-smi fallback only lists GPUs, so it must not override this.
+    """
+
+
 def configure() -> None:
     """Read current fitness settings; setup errors are handled by the runner."""
     global TIMEOUT_SECONDS, MAX_ERROR_MESSAGES
@@ -161,6 +168,10 @@ async def _run_gpu_test_binary() -> dict[str, Any]:
             error_msg = "GPU memory allocation test failed"
             if result["errors"]:
                 error_msg += f": {'; '.join(result['errors'][:MAX_ERROR_MESSAGES])}"
+            # Without "Found N GPUs" the binary never reached the devices
+            # (NVML or loader failure), so nvidia-smi may still decide.
+            if result["found_gpus"]:
+                raise GpuAllocationError(error_msg)
             raise RuntimeError(error_msg)
 
         log.info(
@@ -176,9 +187,7 @@ async def _run_gpu_test_binary() -> dict[str, Any]:
         raise RuntimeError(
             f"GPU test binary timed out after {TIMEOUT_SECONDS}s"
         ) from None
-    except FileNotFoundError:
-        raise
-    except PermissionError:
+    except (FileNotFoundError, PermissionError, GpuAllocationError):
         raise
     except Exception as exc:
         raise RuntimeError(f"GPU test binary execution failed: {exc}") from exc
@@ -238,8 +247,9 @@ async def _check_gpu_health() -> None:
 
     Execution strategy:
     1. Try binary test if available
-    2. Fall back to Python check if binary fails/missing
-    3. Raise RuntimeError if all methods fail
+    2. Fail if the binary reports a GPU allocation failure
+    3. Otherwise fall back to Python check if binary fails/missing
+    4. Raise RuntimeError if all methods fail
 
     Raises:
         RuntimeError: If GPU health check fails
@@ -251,6 +261,8 @@ async def _check_gpu_health() -> None:
     try:
         await _run_gpu_test_binary()
         return  # Success!
+    except GpuAllocationError:
+        raise
     except FileNotFoundError as exc:
         log.debug(f"GPU binary not found: {exc}")
         binary_error = exc

runpod/_health/system.py

@@ -17,6 +17,7 @@
 import os
 import re
 import shutil
+import sys
 import time
 from urllib.parse import urlsplit
 
@@ -39,6 +40,31 @@
 CUDA_VERSION_PROBE_TIMEOUT = 5
 
 
+GPU_FRAMEWORKS = ("torch", "cupy")
+
+
+def _loaded_module(name: str):
+    """
+    Return a module the handler has already imported, without importing it.
+
+    Importing torch costs ~1.5s on a cold worker. The torch and CuPy checks
+    reuse whatever the handler loaded; a framework it has not loaded is
+    treated as unavailable, and the gpu_test binary still verifies the GPU.
+
+    Raises:
+        ImportError: If the handler has not imported the module.
+    """
+    module = sys.modules.get(name)
+    if module is None:
+        raise ImportError(f"{name} not loaded by the handler")
+    return module
+
+
+def _handler_loaded_gpu_framework() -> bool:
+    """True if the handler has imported torch or CuPy."""
+    return any(sys.modules.get(name) is not None for name in GPU_FRAMEWORKS)
+
+
 def configure() -> None:
     """Read current fitness settings; setup errors are handled by the runner."""
     global \
@@ -347,19 +373,27 @@ async def _check_cuda_initialization() -> None:
 
     Tests actual device initialization, memory access, and device properties.
     This catches issues where CUDA appears available but fails at runtime.
-    Skips silently on CPU-only workers.
+    Skips silently on CPU-only workers and when the handler has not
+    imported PyTorch or CuPy.
 
     Raises:
         RuntimeError: If CUDA initialization or device access fails
     """
+    if not _handler_loaded_gpu_framework():
+        log.debug(
+            "Handler has not loaded PyTorch or CuPy, skipping CUDA initialization "
+            "check (the gpu_test binary covers device allocation)"
+        )
+        return
+
     # Skip on CPU-only workers
     if not gpu_available():
         log.debug("No GPU detected, skipping CUDA initialization check")
         return
 
     # Try PyTorch first (most common)
     try:
-        import torch
+        torch = _loaded_module("torch")
 
         if not torch.cuda.is_available():
             log.debug("CUDA not available in PyTorch, skipping initialization check")
@@ -403,7 +437,7 @@ async def _check_cuda_initialization() -> None:
 
     # Fallback: try CuPy
     try:
-        import cupy as cp
+        cp = _loaded_module("cupy")
 
         # Reset CuPy state
         cp.cuda.Device().synchronize()
@@ -441,19 +475,26 @@ async def _check_gpu_compute_benchmark() -> None:
     Quick GPU compute benchmark using matrix multiplication.
 
     Tests basic tensor operations to ensure GPU is functional and responsive.
-    Skips silently on CPU-only workers.
+    Skips silently on CPU-only workers and when the handler has not
+    imported PyTorch or CuPy.
 
     Raises:
         RuntimeError: If GPU compute fails or is too slow
     """
+    if not _handler_loaded_gpu_framework():
+        log.debug(
+            "Handler has not loaded PyTorch or CuPy, skipping GPU compute benchmark"
+        )
+        return
+
     # Skip on CPU-only workers
     if not gpu_available():
         log.debug("No GPU detected, skipping GPU compute benchmark")
         return
 
     # Try PyTorch first
     try:
-        import torch
+        torch = _loaded_module("torch")
 
         if not torch.cuda.is_available():
             log.debug("CUDA not available in PyTorch, skipping benchmark")
@@ -492,7 +533,7 @@ async def _check_gpu_compute_benchmark() -> None:
 
     # Fallback: try CuPy
     try:
-        import cupy as cp
+        cp = _loaded_module("cupy")
 
         size = 1024
         start_time = time.perf_counter()
@@ -536,9 +577,10 @@ def auto_register_system_checks() -> None:
     Registers memory, disk, and network checks for all workers.
     Registers CUDA version, initialization, and GPU benchmark checks only if GPU is detected.
 
-    The two checks that import torch and allocate on the device are marked
+    The two checks that use torch and allocate on the device are marked
     @defer_to_worker_start so the import-time pass cannot create a CUDA context
-    before the handler module runs.
+    before the handler module runs. They only use a framework the handler has
+    already imported, so they never add a torch import to cold start.
     """
     configure()
     log.debug("Registering system resource fitness checks")

tests/test_serverless/test_modules/test_fitness/test_gpu_checks.py

@@ -288,6 +288,45 @@ async def test_health_check_binary_success(self):
             # Should not raise
             await _check_gpu_health()
 
+    @staticmethod
+    def _run_health_check(binary_output: str, nvidia_smi_lists_gpu: bool = True):
+        """Run _check_gpu_health with gpu_test printing binary_output."""
+        process = AsyncMock()
+        process.communicate = AsyncMock(return_value=(binary_output.encode(), b""))
+        nvidia_smi = MagicMock(
+            returncode=0 if nvidia_smi_lists_gpu else 1,
+            stdout="GPU 0: NVIDIA GeForce RTX 4090 (UUID: GPU-1)\n",
+            stderr="",
+        )
+
+        with patch(
+            "runpod.serverless.modules.rp_gpu_fitness._get_gpu_test_binary_path",
+            return_value=Path("/fake/gpu_test"),
+        ), patch("os.access", return_value=True), patch(
+            "asyncio.create_subprocess_exec", return_value=process
+        ), patch("subprocess.run", return_value=nvidia_smi) as mock_run:
+            return mock_run, asyncio.run(_check_gpu_health())
+
+    def test_health_check_rejects_allocation_failure_reported_by_binary(self):
+        """nvidia-smi still listing the GPU must not override a failed allocation."""
+        output = (
+            "Found 1 GPUs:\n"
+            "GPU 0: NVIDIA GeForce RTX 4090 (UUID: GPU-1)\n"
+            "GPU 0 memory allocation test failed. Error code: 2 (out of memory)\n"
+        )
+
+        with pytest.raises(RuntimeError, match="GPU memory allocation test failed"):
+            self._run_health_check(output)
+
+    def test_health_check_falls_back_when_binary_cannot_enumerate_gpus(self):
+        """A binary that never reached the GPUs (e.g. NVML init) still falls back."""
+        mock_run, _ = self._run_health_check(
+            "Failed to initialize NVML: version mismatch\n"
+        )
+
+        mock_run.assert_called_once()
+        assert mock_run.call_args.args[0] == ["nvidia-smi", "--list-gpus"]
+
 
 
 # ============================================================================

tests/test_serverless/test_modules/test_fitness/test_system_checks.py

@@ -6,6 +6,9 @@
 """
 
 import asyncio
+import builtins
+import contextlib
+import sys
 from unittest.mock import patch, MagicMock, AsyncMock
 
 import pytest
@@ -35,6 +38,27 @@ def _make_async_process(returncode=0, stdout="", stderr=""):
     return mock_process
 
 
+@contextlib.contextmanager
+def _no_gpu_framework_loaded():
+    """Simulate a handler that has not imported torch or CuPy.
+
+    Yields the set of top-level module names the code under test tries to
+    import, so a test can assert the check did not import torch itself.
+    """
+    attempted_imports = set()
+    real_import = builtins.__import__
+
+    def recording_import(name, *args, **kwargs):
+        attempted_imports.add(name.split(".")[0])
+        return real_import(name, *args, **kwargs)
+
+    with patch.dict("sys.modules"):
+        sys.modules.pop("torch", None)
+        sys.modules.pop("cupy", None)
+        with patch("builtins.__import__", side_effect=recording_import):
+            yield attempted_imports
+
+
 # ============================================================================
 # Memory Check Tests
 # ============================================================================
@@ -444,6 +468,21 @@ async def test_cuda_init_cupy_fallback(self, mock_gpu_available):
             # Should not raise
             await _check_cuda_initialization()
 
+    @pytest.mark.asyncio
+    @patch("runpod.serverless.modules.rp_system_fitness.gpu_available")
+    async def test_cuda_init_skips_when_handler_has_not_loaded_torch(
+        self, mock_gpu_available
+    ):
+        """The check must not import torch itself; that import costs ~1.5s on a cold worker."""
+        mock_gpu_available.return_value = True
+
+        with _no_gpu_framework_loaded() as attempted_imports:
+            await _check_cuda_initialization()
+
+        assert "torch" not in attempted_imports
+        assert "cupy" not in attempted_imports
+        mock_gpu_available.assert_not_called()
+
     @pytest.mark.asyncio
     @patch("runpod.serverless.modules.rp_system_fitness.gpu_available")
     async def test_cuda_init_no_libraries(self, mock_gpu_available):
@@ -494,6 +533,36 @@ async def test_gpu_benchmark_with_torch_available(self, mock_gpu_available):
             # The function should complete without raising
             await _check_gpu_compute_benchmark()
 
+    @pytest.mark.asyncio
+    @patch("runpod.serverless.modules.rp_system_fitness.gpu_available")
+    async def test_gpu_benchmark_skips_when_handler_has_not_loaded_torch(
+        self, mock_gpu_available
+    ):
+        """The benchmark must not import torch itself either."""
+        mock_gpu_available.return_value = True
+
+        with _no_gpu_framework_loaded() as attempted_imports:
+            await _check_gpu_compute_benchmark()
+
+        assert "torch" not in attempted_imports
+        assert "cupy" not in attempted_imports
+        mock_gpu_available.assert_not_called()
+
+    @pytest.mark.asyncio
+    @patch("runpod.serverless.modules.rp_system_fitness.gpu_available")
+    async def test_gpu_benchmark_uses_cupy_when_handler_loaded_it(
+        self, mock_gpu_available
+    ):
+        """A handler that loaded CuPy but not torch still gets the benchmark."""
+        mock_gpu_available.return_value = True
+        mock_cupy = MagicMock()
+
+        with patch.dict("sys.modules", {"cupy": mock_cupy}):
+            sys.modules.pop("torch", None)
+            await _check_gpu_compute_benchmark()
+
+        mock_cupy.matmul.assert_called_once()
+
     @pytest.mark.asyncio
     @patch("runpod.serverless.modules.rp_system_fitness.gpu_available")
     async def test_gpu_benchmark_skips_no_libraries(self, mock_gpu_available):