perf(serverless): skip torch fitness checks unless loaded

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 (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")
         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_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):