feat(serverless): run fitness checks at startup, add skip env var

ARCHITECTURE.md

@@ -613,7 +613,7 @@ log.error(message, job_id=None)
 - `clear_fitness_checks()`: Clear registry (testing only)
 
 **Execution Flow**:
-1. Called from `worker.py:40` before heartbeat starts: `asyncio.run(run_fitness_checks())`
+1. Runs twice per worker: built-in checks at `import runpod.serverless` via `run_startup_fitness_checks()`, then user-registered and `@defer_to_worker_start` checks from `worker.py:40` before heartbeat starts: `asyncio.run(run_fitness_checks())`; completed checks are not repeated, and `RUNPOD_DEFER_FITNESS_CHECKS=true` collapses both passes into the `worker.py` one
 2. Runs only in production mode (skipped for local testing)
 3. Auto-detects sync vs async using `inspect.iscoroutinefunction()`
 4. Executes checks in registration order (list preserves order)

docs/serverless/worker_fitness_checks.md

@@ -41,6 +41,23 @@ if __name__ == "__main__":
     runpod.serverless.start({"handler": handler})
 ```
 
+## When Checks Run
+
+The built-in GPU and system checks run at **import time** — when your handler module runs `import runpod`, before it loads a model — so a broken GPU or full disk fails the worker in seconds instead of after a multi-minute load.
+
+Two checks stay at `runpod.serverless.start()`: the CUDA initialization check and the GPU compute benchmark. Both import `torch` and allocate on the device, which would leave a CUDA context in a process your handler may later fork — unsupported by CUDA, and something vLLM and DeepSpeed trip over. The remaining built-ins (memory, disk, network, CUDA version via `nvidia-smi`, and the native `gpu_test` binary) run at import.
+
+Your own `@register_fitness_check` functions are registered after that import, so they also run at `start()`. Checks that already passed are not repeated.
+
+Note that the memory check now measures a fresh container rather than one with your model loaded, so `RUNPOD_MIN_MEMORY_GB` validates the environment you were given, not the headroom left after loading.
+
+The import-time pass no-ops outside a real worker (no `RUNPOD_WEBHOOK_GET_JOB`), leaving local runs and tests unaffected — note that inside a worker container *any* `import runpod` triggers it. Set `RUNPOD_DEFER_FITNESS_CHECKS=true` to run everything at `start()` instead.
+
+Two details worth knowing:
+
+- Thresholds and skip flags (`RUNPOD_MIN_*`, `RUNPOD_SKIP_*`, `RUNPOD_GPU_*`) are read when the checks first run, so set them **before** `import runpod` — Dockerfile `ENV` recommended; setting them from Python in your handler is too late on the real platform.
+- Workers serving the realtime API (`--rp_serve_api`) never enter the worker loop, so only the import-time checks apply there; the two deferred CUDA checks do not run in that mode. Child processes created with multiprocessing's `spawn` start method re-import this module but inherit a marker and skip the checks.
+
 ## Async Fitness Checks
 
 Fitness checks support both synchronous and asynchronous functions:
@@ -371,13 +388,15 @@ ENV RUNPOD_NETWORK_CHECK_TIMEOUT=10
 ENV RUNPOD_GPU_BENCHMARK_TIMEOUT=2
 ```
 
-Or in Python:
+Or in Python, before `import runpod` (on the real platform the checks run at import):
 
 ```python
 import os
 
 os.environ["RUNPOD_MIN_MEMORY_GB"] = "8.0"
 os.environ["RUNPOD_MIN_DISK_PERCENT"] = "15.0"
+
+import runpod
 ```
 
 ### Disabling Built-in Checks
@@ -388,6 +407,8 @@ For testing or specialized deployments, built-in checks can be disabled via envi
 |---|---|
 | `RUNPOD_SKIP_AUTO_SYSTEM_CHECKS=true` | Skips auto-registration of memory, disk, network, CUDA version, CUDA init, and GPU benchmark checks |
 | `RUNPOD_SKIP_GPU_CHECK=true` | Skips auto-registration of the native GPU memory allocation test (`gpu_test` binary) |
+| `RUNPOD_SKIP_FITNESS_CHECKS=true` | Skips every fitness check, built-in **and** user-registered |
+| `RUNPOD_DEFER_FITNESS_CHECKS=true` | Keeps the checks but runs them only at `runpod.serverless.start()`, not at import |
 
 ```python
 import os
@@ -397,15 +418,19 @@ os.environ["RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"] = "true"
 
 # Disable the automatic GPU memory allocation test
 os.environ["RUNPOD_SKIP_GPU_CHECK"] = "true"
+
+import runpod
 ```
 
-User-registered checks via `@register_fitness_check` still run regardless of these flags.
+As with the thresholds, set these before `import runpod` on the real platform.
+
+User-registered checks via `@register_fitness_check` still run regardless of `RUNPOD_SKIP_AUTO_SYSTEM_CHECKS` and `RUNPOD_SKIP_GPU_CHECK`. Only `RUNPOD_SKIP_FITNESS_CHECKS` disables those too.
 
 ## Behavior
 
 ### Execution Timing
 
-- Fitness checks run **only once at worker startup**
+- Each check runs **once per worker**: built-ins at import, your registered checks and the deferred CUDA checks at `start()`; checks that passed are not repeated
 - They run **before the first job is processed**
 - They run **only on the actual Runpod serverless platform**
 - Local development and testing modes skip fitness checks
@@ -555,7 +580,7 @@ async def check_api_with_retry():
 
 ## Testing
 
-When developing locally, fitness checks don't run. To test them, you can manually invoke the runner:
+When developing locally, fitness checks don't run. To test them, you can manually invoke the runner. Note that each check runs once per process: a second `run_fitness_checks()` call skips checks that already passed, so call `clear_fitness_checks()` (as below) between runs:
 
 ```python
 import asyncio

runpod/serverless/__init__.py

@@ -16,7 +16,7 @@
 from . import worker
 from .modules.rp_logger import RunPodLogger
 from .modules.rp_progress import progress_update
-from .modules.rp_fitness import register_fitness_check
+from .modules.rp_fitness import register_fitness_check, run_startup_fitness_checks
 from .utils.rp_volume_cache import VolumeCache
 
 __all__ = [
@@ -29,6 +29,10 @@
 
 log = RunPodLogger()
 
+# Check the environment here rather than in start(), which a handler module
+# only reaches after loading its model. No-op outside a real worker.
+run_startup_fitness_checks()
+
 
 # ---------------------------------------------------------------------------- #
 #                              Run Time Arguments                              #

runpod/serverless/modules/rp_fitness.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import contextlib
 import inspect
 import os
@@ -50,6 +51,43 @@ def _terminate_unhealthy(code: int = 1) -> None:
 # Global registry for fitness check functions, preserves registration order
 _fitness_checks: list[Callable] = []
 
+# Checks that already passed. Checks run twice per worker -- at import and in
+# run_worker -- so the second pass only runs what was registered in between.
+_completed_checks: list[Callable] = []
+
+# Disables every check, built-in and user-registered.
+SKIP_FITNESS_CHECKS_ENV = "RUNPOD_SKIP_FITNESS_CHECKS"
+
+# Keeps the checks but runs them only in run_worker, as before.
+DEFER_FITNESS_CHECKS_ENV = "RUNPOD_DEFER_FITNESS_CHECKS"
+
+# Set once this process has claimed the startup pass. Child processes spawned
+# with multiprocessing 'spawn' (vLLM, DeepSpeed) re-import this module and
+# inherit the environment; the marker tells them to skip the checks.
+_CHECKS_DONE_ENV = "RUNPOD_FITNESS_CHECKS_DONE"
+
+
+def _env_flag(name: str) -> bool:
+    """True if the env var is set to a truthy value."""
+    return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on")
+
+
+def defer_to_worker_start(func: Callable) -> Callable:
+    """
+    Mark a check as unsafe to run at import.
+
+    The import-time pass skips these; they run in run_worker as before. Used
+    for checks that initialize CUDA in this process -- doing that before the
+    handler module runs would leave a CUDA context in a process the handler
+    may later fork (vLLM, DeepSpeed), which CUDA does not support.
+    """
+    func._runpod_defer_to_worker_start = True
+    return func
+
+
+def _is_deferred(func: Callable) -> bool:
+    return getattr(func, "_runpod_defer_to_worker_start", False)
+
 
 def register_fitness_check(func: Callable) -> Callable:
     """
@@ -92,6 +130,7 @@ def clear_fitness_checks() -> None:
     Not intended for production use.
     """
     _fitness_checks.clear()
+    _completed_checks.clear()
 
 
 _registration_state: dict[str, bool] = {
@@ -165,14 +204,18 @@ def _ensure_gpu_check_registered() -> None:
     if _registration_state["gpu_check"]:
         return
 
-    _registration_state["gpu_check"] = True
-
+    # Latch only on success: a registration failure (e.g. a malformed
+    # RUNPOD_GPU_TEST_TIMEOUT) must re-raise in run_worker, not silently
+    # disable the checks in both passes.
     try:
         from .rp_gpu_fitness import auto_register_gpu_check
-
-        auto_register_gpu_check()
     except ImportError:
         log.debug("GPU fitness check module not found, skipping auto-registration")
+        _registration_state["gpu_check"] = True
+        return
+
+    auto_register_gpu_check()
+    _registration_state["gpu_check"] = True
 
 
 def _ensure_system_checks_registered() -> None:
@@ -182,30 +225,30 @@ def _ensure_system_checks_registered() -> None:
     Deferred until first run to avoid circular import issues during module
     initialization. Called from run_fitness_checks() on first invocation.
     """
-    import os
-
     if _registration_state["system_checks"]:
         return
 
     # Allow disabling system checks for testing
-    if os.environ.get("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "").lower() == "true":
+    if _env_flag("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"):
         log.debug(
             "System fitness checks disabled via environment (RUNPOD_SKIP_AUTO_SYSTEM_CHECKS)"
         )
         _registration_state["system_checks"] = True
         return
 
-    _registration_state["system_checks"] = True
-
+    # Same latch-on-success rule as _ensure_gpu_check_registered.
     try:
         from .rp_system_fitness import auto_register_system_checks
-
-        auto_register_system_checks()
     except ImportError:
         log.debug("System fitness check module not found, skipping auto-registration")
+        _registration_state["system_checks"] = True
+        return
+
+    auto_register_system_checks()
+    _registration_state["system_checks"] = True
 
 
-async def run_fitness_checks() -> None:
+async def run_fitness_checks(include_deferred: bool = True) -> None:
     """
     Execute all registered fitness checks sequentially at startup.
 
@@ -227,6 +270,10 @@ async def run_fitness_checks() -> None:
     6. On successful completion of all checks:
        - Log completion message with total execution time
 
+    Each check runs once per process: completed checks are skipped on later
+    calls, and @defer_to_worker_start checks are skipped when include_deferred
+    is False (the import-time pass).
+
     Note:
         Checks run in registration order (list preserves order).
         Sequential execution (not parallel) ensures clear error reporting
@@ -237,22 +284,37 @@ async def run_fitness_checks() -> None:
         A failing check terminates the process via os._exit(1); this function
         does not return in that case and does not raise SystemExit.
     """
+    if _env_flag(SKIP_FITNESS_CHECKS_ENV):
+        log.info(f"Fitness checks disabled via {SKIP_FITNESS_CHECKS_ENV}, skipping.")
+        return
+
     # Defer GPU check auto-registration until fitness checks are about to run
     # This avoids circular import issues during module initialization
     _ensure_gpu_check_registered()
 
     # Defer system check auto-registration until fitness checks are about to run
     _ensure_system_checks_registered()
 
-    if not _fitness_checks:
-        log.debug("No fitness checks registered, skipping.")
+    # Identity, not equality: two distinct registrations may compare equal
+    # (e.g. fresh bound-method objects of one method), and `==` would skip one.
+    pending = [
+        check
+        for check in _fitness_checks
+        if not any(check is done for done in _completed_checks)
+    ]
+
+    if not include_deferred:
+        pending = [check for check in pending if not _is_deferred(check)]
+
+    if not pending:
+        log.debug("No pending fitness checks, skipping.")
         return
 
-    log.info(f"Running {len(_fitness_checks)} fitness check(s)...")
+    log.info(f"Running {len(pending)} fitness check(s)...")
 
     total_start_time = time.perf_counter()
 
-    for check_func in _fitness_checks:
+    for check_func in pending:
         check_name = check_func.__name__
 
         try:
@@ -266,6 +328,7 @@ async def run_fitness_checks() -> None:
                 check_func()
 
             check_elapsed_ms = (time.perf_counter() - check_start_time) * 1000
+            _completed_checks.append(check_func)
             log.debug(f"Fitness check passed: {check_name} ({check_elapsed_ms:.2f}ms)")
 
         except Exception as exc:
@@ -294,3 +357,54 @@ async def run_fitness_checks() -> None:
 
     total_elapsed_ms = (time.perf_counter() - total_start_time) * 1000
     log.info(f"All fitness checks passed. ({total_elapsed_ms:.2f}ms)")
+
+
+def _event_loop_running() -> bool:
+    """True if called from inside a running event loop."""
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return False
+    return True
+
+
+def run_startup_fitness_checks() -> None:
+    """
+    Run the built-in fitness checks at import, before the handler loads a model.
+
+    A user's @register_fitness_check functions are registered after this import,
+    so they still run in run_worker, which skips whatever passed here. Checks
+    marked with @defer_to_worker_start are also left to run_worker.
+
+    No-ops outside a real worker (no RUNPOD_WEBHOOK_GET_JOB), when the checks
+    are disabled or deferred, inside a running event loop, and in child
+    processes (multiprocessing 'spawn' re-imports this module; the worker marks
+    itself done via env so children skip). Ordinary exceptions from running the
+    checks are logged and swallowed: a failure to run the checks must not stop
+    a worker from booting. A failing check still force-exits, which is the point.
+    """
+    if _env_flag(SKIP_FITNESS_CHECKS_ENV) or _env_flag(DEFER_FITNESS_CHECKS_ENV):
+        return
+
+    if not os.environ.get("RUNPOD_WEBHOOK_GET_JOB"):
+        return
+
+    if os.environ.get(_CHECKS_DONE_ENV):
+        return
+    os.environ[_CHECKS_DONE_ENV] = "1"
+
+    if _event_loop_running():
+        log.debug("Event loop already running, deferring fitness checks to run_worker.")
+        return
+
+    try:
+        # Own loop rather than asyncio.run: run() resets the thread's loop
+        # policy state, after which asyncio.get_event_loop() in handler code
+        # raises RuntimeError on Python 3.10+.
+        loop = asyncio.new_event_loop()
+        try:
+            loop.run_until_complete(run_fitness_checks(include_deferred=False))
+        finally:
+            loop.close()
+    except Exception as exc:  # pragma: no cover - defensive
+        log.error(f"Startup fitness checks could not run: {exc}")

runpod/serverless/modules/rp_gpu_fitness.py

@@ -17,7 +17,7 @@
 from typing import Any
 
 from runpod._binary_helpers import get_binary_path
-from .rp_fitness import register_fitness_check
+from .rp_fitness import _env_flag, register_fitness_check
 from .rp_logger import RunPodLogger
 
 log = RunPodLogger()
@@ -286,7 +286,7 @@ def auto_register_gpu_check() -> None:
     - RUNPOD_SKIP_GPU_CHECK: Set to "true" to skip auto-registration
     """
     # Allow skipping during tests
-    if os.environ.get("RUNPOD_SKIP_GPU_CHECK", "").lower() == "true":
+    if _env_flag("RUNPOD_SKIP_GPU_CHECK"):
         log.debug("GPU fitness check auto-registration disabled via environment")
         return
 

runpod/serverless/modules/rp_system_fitness.py

@@ -19,7 +19,7 @@
 import shutil
 import time
 
-from .rp_fitness import register_fitness_check
+from .rp_fitness import defer_to_worker_start, register_fitness_check
 from .rp_logger import RunPodLogger
 from ..utils.rp_cuda import is_available as gpu_available
 
@@ -470,6 +470,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
+    @defer_to_worker_start so the import-time pass cannot create a CUDA context
+    before the handler module runs.
     """
     log.debug("Registering system resource fitness checks")
 
@@ -499,11 +503,13 @@ async def _cuda_version_check() -> None:
             await _check_cuda_versions()
 
         @register_fitness_check
+        @defer_to_worker_start
         async def _cuda_init_check() -> None:
             """CUDA device initialization check."""
             await _check_cuda_initialization()
 
         @register_fitness_check
+        @defer_to_worker_start
         async def _benchmark_check() -> None:
             """GPU compute benchmark check."""
             await _check_gpu_compute_benchmark()