From 99fc475397251a7162a311a93092fb88412c38c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Mon, 19 Oct 2026 01:11:29 -0700 Subject: [PATCH 1/2] perf(serverless): skip torch fitness checks unless loaded The worker-start CUDA init and GPU benchmark checks imported torch themselves. On a cold worker that import is ~1.5s, about 90% of the check, and it ran before the first job-take even for handlers that never use torch or import it lazily (every Flash handler). - Use torch or CuPy only if the handler already imported it (sys.modules); never import either from the check - Skip both checks before calling nvidia-smi when neither is loaded; the gpu_test binary still verifies device init and allocation - Unchanged behavior when torch or CuPy is loaded Real GPU worker (RTX 4090, Flash handler): worker-start fitness pass went from 1339/2147ms to 31ms. Linear: SLS-815 --- runpod/_health/system.py | 58 +++++++++++++--- .../test_fitness/test_system_checks.py | 69 +++++++++++++++++++ 2 files changed, 119 insertions(+), 8 deletions(-) diff --git a/runpod/_health/system.py b/runpod/_health/system.py index 8a80327cd..60c68ed7d 100644 --- a/runpod/_health/system.py +++ b/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,11 +373,19 @@ 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 (gpu_test binary already verified device init)" + ) + return + # Skip on CPU-only workers if not gpu_available(): log.debug("No GPU detected, skipping CUDA initialization check") @@ -359,7 +393,7 @@ async def _check_cuda_initialization() -> None: # 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,11 +475,18 @@ 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") @@ -453,7 +494,7 @@ async def _check_gpu_compute_benchmark() -> None: # 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") diff --git a/tests/test_serverless/test_modules/test_fitness/test_system_checks.py b/tests/test_serverless/test_modules/test_fitness/test_system_checks.py index 2ac370b15..824b9393d 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_system_checks.py +++ b/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): From fa99717551790fccf159673f04594e021d625fcb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Mon, 19 Oct 2026 14:13:29 -0700 Subject: [PATCH 2/2] fix(health): fail GPU check on reported alloc errors The nvidia-smi fallback only lists GPUs, so it passed workers whose gpu_test run reported a failed device allocation. With the torch CUDA init check now skipped for handlers that never import torch, nothing else caught that case. - Raise GpuAllocationError when gpu_test enumerated GPUs and one failed its allocation test; _check_gpu_health re-raises it without falling back - Keep the fallback when the binary is missing, not executable, or never reached the GPUs (NVML or loader failure) - Correct the skip log that claimed gpu_test had verified device init --- runpod/_health/gpu.py | 22 ++++++++--- runpod/_health/system.py | 2 +- .../test_fitness/test_gpu_checks.py | 39 +++++++++++++++++++ 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/runpod/_health/gpu.py b/runpod/_health/gpu.py index cae54ffed..930dba966 100644 --- a/runpod/_health/gpu.py +++ b/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 diff --git a/runpod/_health/system.py b/runpod/_health/system.py index 60c68ed7d..578396f4d 100644 --- a/runpod/_health/system.py +++ b/runpod/_health/system.py @@ -382,7 +382,7 @@ async def _check_cuda_initialization() -> None: if not _handler_loaded_gpu_framework(): log.debug( "Handler has not loaded PyTorch or CuPy, skipping CUDA initialization " - "check (gpu_test binary already verified device init)" + "check (the gpu_test binary covers device allocation)" ) return diff --git a/tests/test_serverless/test_modules/test_fitness/test_gpu_checks.py b/tests/test_serverless/test_modules/test_fitness/test_gpu_checks.py index b7aaa92e9..fc1c6ffb9 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_gpu_checks.py +++ b/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"] + # ============================================================================