From 8c2c2801bcb3a289c099e753d93a3a03a4efc38b Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 25 Aug 2026 22:31:47 -0400 Subject: [PATCH 1/4] feat(serverless): run fitness checks at import, add global skip env var Built-in GPU/system fitness checks ran in run_worker, which a handler module only reaches after loading its model. Run them when runpod.serverless is imported instead, so a broken environment fails in seconds. User-registered checks still run at start(); checks that already passed are not repeated. Adds RUNPOD_SKIP_FITNESS_CHECKS to disable all checks and RUNPOD_DEFER_FITNESS_CHECKS to restore the previous start()-only timing. --- docs/serverless/worker_fitness_checks.md | 12 +- runpod/serverless/__init__.py | 7 +- runpod/serverless/modules/rp_fitness.py | 90 ++++++++++++- .../test_modules/test_fitness/test_startup.py | 126 ++++++++++++++++++ 4 files changed, 229 insertions(+), 6 deletions(-) create mode 100644 tests/test_serverless/test_modules/test_fitness/test_startup.py diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index c50a9c93..6cbd0203 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -41,6 +41,14 @@ if __name__ == "__main__": runpod.serverless.start({"handler": handler}) ``` +## When Checks Run + +The built-in GPU and system checks run at **import time** — as soon as your handler module runs `import runpod`, before it loads a model. A worker with a broken GPU or a full disk therefore dies in seconds rather than after a multi-minute model load. + +Your own `@register_fitness_check` functions are registered after that import, so they run at `runpod.serverless.start()`. Checks that already passed at import are not repeated. + +The import-time pass is a no-op outside a real worker (no `RUNPOD_WEBHOOK_GET_JOB`), so local runs, tests, and the `runpod` CLI are unaffected. Set `RUNPOD_DEFER_FITNESS_CHECKS=true` to run everything at `start()` instead. + ## Async Fitness Checks Fitness checks support both synchronous and asynchronous functions: @@ -388,6 +396,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 @@ -399,7 +409,7 @@ os.environ["RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"] = "true" os.environ["RUNPOD_SKIP_GPU_CHECK"] = "true" ``` -User-registered checks via `@register_fitness_check` still run regardless of these flags. +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 diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index 05245207..6139b437 100644 --- a/runpod/serverless/__init__.py +++ b/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,11 @@ log = RunPodLogger() +# Validate the worker environment now, at import, rather than waiting for +# start() -- which a handler module only reaches after it has loaded its model. +# No-op outside a real worker; see run_startup_fitness_checks. +run_startup_fitness_checks() + # ---------------------------------------------------------------------------- # # Run Time Arguments # diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index 77df97e7..9fb1c95b 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/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,30 @@ def _terminate_unhealthy(code: int = 1) -> None: # Global registry for fitness check functions, preserves registration order _fitness_checks: list[Callable] = [] +# Checks that have already executed successfully. Fitness checks now run at +# import time (see run_startup_fitness_checks) as well as in run_worker, so the +# second pass must only execute checks registered after the first pass -- the +# user's own @register_fitness_check functions, which are registered between +# the two. +_completed_checks: list[Callable] = [] + +# Env var that disables every fitness check, built-in and user-registered. +SKIP_FITNESS_CHECKS_ENV = "RUNPOD_SKIP_FITNESS_CHECKS" + +# Env var that restores the old behavior: checks run only when the worker +# starts (after the handler module has loaded its model), not at import. +DEFER_FITNESS_CHECKS_ENV = "RUNPOD_DEFER_FITNESS_CHECKS" + + +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 fitness_checks_disabled() -> bool: + """True if the user has opted out of fitness checks entirely.""" + return _env_flag(SKIP_FITNESS_CHECKS_ENV) + def register_fitness_check(func: Callable) -> Callable: """ @@ -92,6 +117,7 @@ def clear_fitness_checks() -> None: Not intended for production use. """ _fitness_checks.clear() + _completed_checks.clear() _registration_state: dict[str, bool] = { @@ -237,6 +263,12 @@ 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 fitness_checks_disabled(): + 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() @@ -244,15 +276,17 @@ async def run_fitness_checks() -> None: # 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.") + pending = [check for check in _fitness_checks if check not in _completed_checks] + + 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 +300,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 +329,50 @@ 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 run_startup_fitness_checks() -> None: + """ + Run the built-in fitness checks as early as possible in the process. + + Called when ``runpod.serverless`` is imported, which on a worker is the + first line of the handler module -- before it loads its model. Running + here means a broken GPU or a full disk kills the worker in seconds instead + of after a multi-minute model load, and long before the first job. + + Only the built-in GPU/system checks can run this early; a user's + ``@register_fitness_check`` functions are registered after this import, so + they still run in ``run_worker``, which skips whatever already passed here. + + No-ops when: + - fitness checks are disabled (``RUNPOD_SKIP_FITNESS_CHECKS``) + - the early run is deferred (``RUNPOD_DEFER_FITNESS_CHECKS``), restoring + the previous run_worker-only behavior + - the process is not a Runpod worker (no ``RUNPOD_WEBHOOK_GET_JOB``), so + local development, tests and the ``runpod`` CLI are untouched + - an event loop is already running, in which case the checks are left to + ``run_worker`` + + Never raises: an unexpected failure here must not stop a worker from + booting. An actual failing check still force-exits, which is the point. + """ + if fitness_checks_disabled() or _env_flag(DEFER_FITNESS_CHECKS_ENV): + return + + if not os.environ.get("RUNPOD_WEBHOOK_GET_JOB"): + return + + try: + asyncio.get_running_loop() + except RuntimeError: + pass + else: + log.debug("Event loop already running, deferring fitness checks to startup.") + return + + try: + asyncio.run(run_fitness_checks()) + except SystemExit: + raise + except Exception as exc: # pragma: no cover - defensive + log.warn(f"Startup fitness checks could not run: {exc}") diff --git a/tests/test_serverless/test_modules/test_fitness/test_startup.py b/tests/test_serverless/test_modules/test_fitness/test_startup.py new file mode 100644 index 00000000..1aaf0c85 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -0,0 +1,126 @@ +"""Tests for fitness checks running at import/startup time (DR-1409).""" + +from unittest.mock import patch + +import pytest + +from runpod.serverless.modules import rp_fitness +from runpod.serverless.modules.rp_fitness import ( + register_fitness_check, + run_fitness_checks, + run_startup_fitness_checks, +) + + +@pytest.fixture() +def worker_env(monkeypatch): + """Make the process look like a real Runpod worker.""" + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.com/job") + monkeypatch.delenv("RUNPOD_SKIP_FITNESS_CHECKS", raising=False) + monkeypatch.delenv("RUNPOD_DEFER_FITNESS_CHECKS", raising=False) + + +class TestSkipEnvVar: + @pytest.mark.asyncio + async def test_skip_env_var_bypasses_all_checks(self, monkeypatch): + monkeypatch.setenv("RUNPOD_SKIP_FITNESS_CHECKS", "true") + called = [] + + @register_fitness_check + def check(): + called.append(True) + + await run_fitness_checks() + assert called == [] + + @pytest.mark.asyncio + async def test_checks_run_when_skip_unset(self, monkeypatch): + monkeypatch.delenv("RUNPOD_SKIP_FITNESS_CHECKS", raising=False) + called = [] + + @register_fitness_check + def check(): + called.append(True) + + await run_fitness_checks() + assert called == [True] + + +class TestRunOnce: + @pytest.mark.asyncio + async def test_passed_check_does_not_rerun(self): + calls = [] + + @register_fitness_check + def first(): + calls.append("first") + + await run_fitness_checks() + + @register_fitness_check + def second(): + calls.append("second") + + await run_fitness_checks() + + assert calls == ["first", "second"] + + +class TestStartupEntrypoint: + def test_runs_checks_on_worker(self, worker_env): + calls = [] + + @register_fitness_check + def check(): + calls.append(True) + + run_startup_fitness_checks() + assert calls == [True] + + def test_noop_outside_worker(self, monkeypatch): + monkeypatch.delenv("RUNPOD_WEBHOOK_GET_JOB", raising=False) + calls = [] + + @register_fitness_check + def check(): + calls.append(True) + + run_startup_fitness_checks() + assert calls == [] + + def test_defer_env_var_postpones_to_worker_start(self, worker_env, monkeypatch): + monkeypatch.setenv("RUNPOD_DEFER_FITNESS_CHECKS", "true") + calls = [] + + @register_fitness_check + def check(): + calls.append(True) + + run_startup_fitness_checks() + assert calls == [] + + def test_skip_env_var_respected(self, worker_env, monkeypatch): + monkeypatch.setenv("RUNPOD_SKIP_FITNESS_CHECKS", "1") + calls = [] + + @register_fitness_check + def check(): + calls.append(True) + + run_startup_fitness_checks() + assert calls == [] + + def test_unexpected_error_does_not_propagate(self, worker_env): + with patch.object(rp_fitness.asyncio, "run", side_effect=RuntimeError("boom")): + run_startup_fitness_checks() + + @pytest.mark.asyncio + async def test_noop_inside_running_loop(self, worker_env): + calls = [] + + @register_fitness_check + def check(): + calls.append(True) + + run_startup_fitness_checks() + assert calls == [] From d2b1824817c5b1544a011070116eda124cd45f8b Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 26 Aug 2026 10:29:04 -0400 Subject: [PATCH 2/4] refactor(serverless): tighten startup fitness check comments and control flow --- docs/serverless/worker_fitness_checks.md | 6 +- runpod/serverless/__init__.py | 5 +- runpod/serverless/modules/rp_fitness.py | 75 +++++++++--------------- 3 files changed, 32 insertions(+), 54 deletions(-) diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index 6cbd0203..e946a4d2 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -43,11 +43,11 @@ if __name__ == "__main__": ## When Checks Run -The built-in GPU and system checks run at **import time** — as soon as your handler module runs `import runpod`, before it loads a model. A worker with a broken GPU or a full disk therefore dies in seconds rather than after a multi-minute model load. +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. -Your own `@register_fitness_check` functions are registered after that import, so they run at `runpod.serverless.start()`. Checks that already passed at import are not repeated. +Your own `@register_fitness_check` functions are registered after that import, so they run at `runpod.serverless.start()`. Checks that already passed are not repeated. -The import-time pass is a no-op outside a real worker (no `RUNPOD_WEBHOOK_GET_JOB`), so local runs, tests, and the `runpod` CLI are unaffected. Set `RUNPOD_DEFER_FITNESS_CHECKS=true` to run everything at `start()` instead. +The import-time pass no-ops outside a real worker (no `RUNPOD_WEBHOOK_GET_JOB`), leaving local runs, tests, and the `runpod` CLI unaffected. Set `RUNPOD_DEFER_FITNESS_CHECKS=true` to run everything at `start()` instead. ## Async Fitness Checks diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index 6139b437..c6fa605e 100644 --- a/runpod/serverless/__init__.py +++ b/runpod/serverless/__init__.py @@ -29,9 +29,8 @@ log = RunPodLogger() -# Validate the worker environment now, at import, rather than waiting for -# start() -- which a handler module only reaches after it has loaded its model. -# No-op outside a real worker; see run_startup_fitness_checks. +# 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() diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index 9fb1c95b..815f1f7c 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -51,18 +51,14 @@ def _terminate_unhealthy(code: int = 1) -> None: # Global registry for fitness check functions, preserves registration order _fitness_checks: list[Callable] = [] -# Checks that have already executed successfully. Fitness checks now run at -# import time (see run_startup_fitness_checks) as well as in run_worker, so the -# second pass must only execute checks registered after the first pass -- the -# user's own @register_fitness_check functions, which are registered between -# the two. +# 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] = [] -# Env var that disables every fitness check, built-in and user-registered. +# Disables every check, built-in and user-registered. SKIP_FITNESS_CHECKS_ENV = "RUNPOD_SKIP_FITNESS_CHECKS" -# Env var that restores the old behavior: checks run only when the worker -# starts (after the handler module has loaded its model), not at import. +# Keeps the checks but runs them only in run_worker, as before. DEFER_FITNESS_CHECKS_ENV = "RUNPOD_DEFER_FITNESS_CHECKS" @@ -71,11 +67,6 @@ def _env_flag(name: str) -> bool: return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on") -def fitness_checks_disabled() -> bool: - """True if the user has opted out of fitness checks entirely.""" - return _env_flag(SKIP_FITNESS_CHECKS_ENV) - - def register_fitness_check(func: Callable) -> Callable: """ Decorator to register a fitness check function. @@ -263,10 +254,8 @@ 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 fitness_checks_disabled(): - log.info( - f"Fitness checks disabled via {SKIP_FITNESS_CHECKS_ENV}, skipping." - ) + 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 @@ -331,48 +320,38 @@ async def run_fitness_checks() -> None: 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 as early as possible in the process. - - Called when ``runpod.serverless`` is imported, which on a worker is the - first line of the handler module -- before it loads its model. Running - here means a broken GPU or a full disk kills the worker in seconds instead - of after a multi-minute model load, and long before the first job. - - Only the built-in GPU/system checks can run this early; a user's - ``@register_fitness_check`` functions are registered after this import, so - they still run in ``run_worker``, which skips whatever already passed here. - - No-ops when: - - fitness checks are disabled (``RUNPOD_SKIP_FITNESS_CHECKS``) - - the early run is deferred (``RUNPOD_DEFER_FITNESS_CHECKS``), restoring - the previous run_worker-only behavior - - the process is not a Runpod worker (no ``RUNPOD_WEBHOOK_GET_JOB``), so - local development, tests and the ``runpod`` CLI are untouched - - an event loop is already running, in which case the checks are left to - ``run_worker`` - - Never raises: an unexpected failure here must not stop a worker from - booting. An actual failing check still force-exits, which is the point. + 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. + + No-ops outside a real worker (no RUNPOD_WEBHOOK_GET_JOB), when the checks + are disabled or deferred, or inside a running event loop. Never raises: a + failure to run the checks must not stop a worker from booting. A failing + check still force-exits, which is the point. """ - if fitness_checks_disabled() or _env_flag(DEFER_FITNESS_CHECKS_ENV): + 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 - try: - asyncio.get_running_loop() - except RuntimeError: - pass - else: - log.debug("Event loop already running, deferring fitness checks to startup.") + if _event_loop_running(): + log.debug("Event loop already running, deferring fitness checks to run_worker.") return try: asyncio.run(run_fitness_checks()) - except SystemExit: - raise except Exception as exc: # pragma: no cover - defensive log.warn(f"Startup fitness checks could not run: {exc}") From 3f1cc2bda71178eda06fb708c6aa6767bc500b2c Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 26 Aug 2026 11:12:27 -0400 Subject: [PATCH 3/4] fix(serverless): keep in-process CUDA checks out of the import-time pass _cuda_init_check and _benchmark_check import torch and allocate on the device. Running them at import would leave a CUDA context in a process the handler may later fork, which CUDA does not support and vLLM/DeepSpeed trip over. Mark them @defer_to_worker_start so only subprocess-based and non-GPU checks run early. --- docs/serverless/worker_fitness_checks.md | 6 ++- runpod/serverless/modules/rp_fitness.py | 27 +++++++++-- .../serverless/modules/rp_system_fitness.py | 8 +++- .../test_modules/test_fitness/test_startup.py | 45 +++++++++++++++++++ 4 files changed, 81 insertions(+), 5 deletions(-) diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index e946a4d2..cc788849 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -45,7 +45,11 @@ if __name__ == "__main__": 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. -Your own `@register_fitness_check` functions are registered after that import, so they run at `runpod.serverless.start()`. Checks that already passed are not repeated. +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, tests, and the `runpod` CLI unaffected. Set `RUNPOD_DEFER_FITNESS_CHECKS=true` to run everything at `start()` instead. diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index 815f1f7c..c025ce05 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -67,6 +67,23 @@ def _env_flag(name: str) -> bool: 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: """ Decorator to register a fitness check function. @@ -222,7 +239,7 @@ def _ensure_system_checks_registered() -> None: log.debug("System fitness check module not found, skipping auto-registration") -async def run_fitness_checks() -> None: +async def run_fitness_checks(include_deferred: bool = True) -> None: """ Execute all registered fitness checks sequentially at startup. @@ -267,6 +284,9 @@ async def run_fitness_checks() -> None: pending = [check for check in _fitness_checks if check not 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 @@ -334,7 +354,8 @@ 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. + 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, or inside a running event loop. Never raises: a @@ -352,6 +373,6 @@ def run_startup_fitness_checks() -> None: return try: - asyncio.run(run_fitness_checks()) + asyncio.run(run_fitness_checks(include_deferred=False)) except Exception as exc: # pragma: no cover - defensive log.warn(f"Startup fitness checks could not run: {exc}") diff --git a/runpod/serverless/modules/rp_system_fitness.py b/runpod/serverless/modules/rp_system_fitness.py index 8dc8946d..1ac80f5b 100644 --- a/runpod/serverless/modules/rp_system_fitness.py +++ b/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() diff --git a/tests/test_serverless/test_modules/test_fitness/test_startup.py b/tests/test_serverless/test_modules/test_fitness/test_startup.py index 1aaf0c85..5dec39b3 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -124,3 +124,48 @@ def check(): run_startup_fitness_checks() assert calls == [] + + +class TestDeferredChecks: + """Checks that touch CUDA in-process must not run at import time.""" + + def test_deferred_check_skipped_at_import(self, worker_env): + calls = [] + + @register_fitness_check + def early(): + calls.append("early") + + @register_fitness_check + @rp_fitness.defer_to_worker_start + def late(): + calls.append("late") + + run_startup_fitness_checks() + assert calls == ["early"] + + @pytest.mark.asyncio + async def test_deferred_check_runs_at_worker_start(self, worker_env): + calls = [] + + @register_fitness_check + @rp_fitness.defer_to_worker_start + def late(): + calls.append("late") + + run_startup_fitness_checks() + assert calls == [] + + await run_fitness_checks() + assert calls == ["late"] + + def test_cuda_checks_are_marked_deferred(self): + from runpod.serverless.modules import rp_system_fitness + + with patch.object(rp_system_fitness, "gpu_available", return_value=True): + rp_system_fitness.auto_register_system_checks() + + by_name = {check.__name__: check for check in rp_fitness._fitness_checks} + assert rp_fitness._is_deferred(by_name["_cuda_init_check"]) + assert rp_fitness._is_deferred(by_name["_benchmark_check"]) + assert not rp_fitness._is_deferred(by_name["_memory_check"]) From a683d627b96ccf69010c870b07a8bda0aebe6638 Mon Sep 17 00:00:00 2001 From: justinwlin <15874969+justinwlin@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:37:52 -0400 Subject: [PATCH 4/4] fix(serverless): address review findings on import-time fitness checks - run startup pass on a dedicated event loop instead of asyncio.run, which resets the loop policy and breaks asyncio.get_event_loop() in handler code on Python 3.10+ - set RUNPOD_FITNESS_CHECKS_DONE after the startup pass so children re-importing this module under multiprocessing 'spawn' skip the checks - latch check auto-registration state only on success, so a malformed RUNPOD_MIN_*/GPU timeout value re-raises loudly in run_worker instead of silently disabling all system checks - compare completed checks by identity, not equality, so distinct registrations that compare equal (bound methods) are not skipped - bound the nvidia-smi call in rp_cuda.is_available with a 5s timeout - accept 1/true/yes/on for RUNPOD_SKIP_GPU_CHECK and RUNPOD_SKIP_AUTO_SYSTEM_CHECKS, matching the new flags - tests: pin the worker.py and import-time wiring, the full defer behavior, the done marker, the real auto-registration path (guard: no torch import), and bound-method re-registration; fix an orphaned coroutine in test_unexpected_error_does_not_propagate - docs: thresholds/skip flags must be set before import runpod, realtime API mode runs only the import-time checks, refresh stale ARCHITECTURE.md execution flow --- ARCHITECTURE.md | 2 +- docs/serverless/worker_fitness_checks.md | 19 ++- runpod/serverless/modules/rp_fitness.py | 66 ++++++--- runpod/serverless/modules/rp_gpu_fitness.py | 4 +- runpod/serverless/utils/rp_cuda.py | 6 +- .../test_modules/test_fitness/conftest.py | 3 + .../test_modules/test_fitness/test_startup.py | 140 +++++++++++++++++- tests/test_serverless/test_utils/test_cuda.py | 12 +- tests/test_serverless/test_worker.py | 5 +- 9 files changed, 227 insertions(+), 30 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 2dbca2ef..26be2681 100644 --- a/ARCHITECTURE.md +++ b/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) diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index cc788849..1717fbd3 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -51,7 +51,12 @@ Your own `@register_fitness_check` functions are registered after that import, s 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, tests, and the `runpod` CLI unaffected. Set `RUNPOD_DEFER_FITNESS_CHECKS=true` to run everything at `start()` instead. +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 @@ -383,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 @@ -411,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 ``` +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 @@ -569,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 diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index c025ce05..7110fad6 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -61,6 +61,11 @@ def _terminate_unhealthy(code: int = 1) -> None: # 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.""" @@ -199,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: @@ -216,27 +225,27 @@ 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(include_deferred: bool = True) -> None: @@ -261,6 +270,10 @@ async def run_fitness_checks(include_deferred: bool = True) -> 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 @@ -282,7 +295,13 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: # Defer system check auto-registration until fitness checks are about to run _ensure_system_checks_registered() - pending = [check for check in _fitness_checks if check not in _completed_checks] + # 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)] @@ -358,9 +377,11 @@ def run_startup_fitness_checks() -> None: 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, or inside a running event loop. Never raises: a - failure to run the checks must not stop a worker from booting. A failing - check still force-exits, which is the point. + 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 @@ -368,11 +389,22 @@ def run_startup_fitness_checks() -> None: 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: - asyncio.run(run_fitness_checks(include_deferred=False)) + # 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.warn(f"Startup fitness checks could not run: {exc}") + log.error(f"Startup fitness checks could not run: {exc}") diff --git a/runpod/serverless/modules/rp_gpu_fitness.py b/runpod/serverless/modules/rp_gpu_fitness.py index bae74cd8..fd659009 100644 --- a/runpod/serverless/modules/rp_gpu_fitness.py +++ b/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 diff --git a/runpod/serverless/utils/rp_cuda.py b/runpod/serverless/utils/rp_cuda.py index 028c7ebc..1a47108a 100644 --- a/runpod/serverless/utils/rp_cuda.py +++ b/runpod/serverless/utils/rp_cuda.py @@ -10,7 +10,11 @@ def is_available(): Returns True if CUDA is available, False otherwise. """ try: - output = subprocess.check_output(["nvidia-smi"], stderr=subprocess.DEVNULL) + # Bounded: this runs at `import runpod` on real workers, where a wedged + # nvidia-smi must not hang the boot forever. + output = subprocess.check_output( + ["nvidia-smi"], stderr=subprocess.DEVNULL, timeout=5 + ) if "NVIDIA-SMI" in output.decode(): return True except Exception: # pylint: disable=broad-except diff --git a/tests/test_serverless/test_modules/test_fitness/conftest.py b/tests/test_serverless/test_modules/test_fitness/conftest.py index f8df8614..12810382 100644 --- a/tests/test_serverless/test_modules/test_fitness/conftest.py +++ b/tests/test_serverless/test_modules/test_fitness/conftest.py @@ -23,6 +23,9 @@ def cleanup_fitness_checks(monkeypatch): """ monkeypatch.setenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "true") monkeypatch.setenv("RUNPOD_SKIP_GPU_CHECK", "true") + # run_startup_fitness_checks sets this directly on the real environment; + # clear it per-test so the marker cannot leak between tests. + monkeypatch.delenv(rp_fitness._CHECKS_DONE_ENV, raising=False) def _raise_system_exit(code=0): raise SystemExit(code) diff --git a/tests/test_serverless/test_modules/test_fitness/test_startup.py b/tests/test_serverless/test_modules/test_fitness/test_startup.py index 5dec39b3..afddfbfd 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -1,5 +1,9 @@ """Tests for fitness checks running at import/startup time (DR-1409).""" +import builtins +import os +import sys +import types from unittest.mock import patch import pytest @@ -65,6 +69,26 @@ def second(): assert calls == ["first", "second"] + @pytest.mark.asyncio + async def test_equal_but_distinct_registration_still_runs(self): + # Bound-method objects are distinct but compare equal; an == check + # against _completed_checks would wrongly skip the re-registration. + calls = [] + + class Checker: + def check(self): + calls.append("bound") + + obj = Checker() + + register_fitness_check(obj.check) + await run_fitness_checks() + + register_fitness_check(obj.check) + await run_fitness_checks() + + assert calls == ["bound", "bound"] + class TestStartupEntrypoint: def test_runs_checks_on_worker(self, worker_env): @@ -111,7 +135,11 @@ def check(): assert calls == [] def test_unexpected_error_does_not_propagate(self, worker_env): - with patch.object(rp_fitness.asyncio, "run", side_effect=RuntimeError("boom")): + # Patch loop construction, not loop execution: patching asyncio.run + # would orphan the coroutine argument and trip unraisable warnings. + with patch.object( + rp_fitness.asyncio, "new_event_loop", side_effect=RuntimeError("boom") + ): run_startup_fitness_checks() @pytest.mark.asyncio @@ -169,3 +197,113 @@ def test_cuda_checks_are_marked_deferred(self): assert rp_fitness._is_deferred(by_name["_cuda_init_check"]) assert rp_fitness._is_deferred(by_name["_benchmark_check"]) assert not rp_fitness._is_deferred(by_name["_memory_check"]) + + +class TestDoneMarker: + """Spawned children re-import this module and must not re-run the checks.""" + + def test_done_marker_skips_startup_pass(self, worker_env, monkeypatch): + monkeypatch.setenv(rp_fitness._CHECKS_DONE_ENV, "1") + calls = [] + + @register_fitness_check + def check(): + called.append(True) + + run_startup_fitness_checks() + assert calls == [] + + def test_startup_pass_sets_done_marker(self, worker_env): + run_startup_fitness_checks() + assert os.environ.get(rp_fitness._CHECKS_DONE_ENV) == "1" + + +class TestDeferFullBehavior: + """RUNPOD_DEFER_FITNESS_CHECKS restores exact pre-PR start()-only timing.""" + + @pytest.mark.asyncio + async def test_deferred_to_start_runs_everything(self, worker_env, monkeypatch): + monkeypatch.setenv("RUNPOD_DEFER_FITNESS_CHECKS", "true") + calls = [] + + @register_fitness_check + def check(): + calls.append(True) + + @register_fitness_check + @rp_fitness.defer_to_worker_start + def deferred(): + calls.append("deferred") + + run_startup_fitness_checks() + assert calls == [] + + await run_fitness_checks() + assert calls == [True, "deferred"] + + +class TestAutoRegistrationPath: + """Exercise the real _ensure_*_registered path during the startup pass.""" + + def test_startup_runs_auto_registered_checks_without_torch( + self, worker_env, monkeypatch + ): + calls = [] + + fake_gpu_module = types.SimpleNamespace( + auto_register_gpu_check=lambda: register_fitness_check( + lambda: calls.append("gpu") + ) + ) + + def register_system_checks(): + register_fitness_check(lambda: calls.append("system")) + register_fitness_check( + rp_fitness.defer_to_worker_start(lambda: calls.append("deferred")) + ) + + fake_system_module = types.SimpleNamespace( + auto_register_system_checks=register_system_checks + ) + + monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", raising=False) + monkeypatch.delenv("RUNPOD_SKIP_GPU_CHECK", raising=False) + monkeypatch.setitem( + sys.modules, "runpod.serverless.modules.rp_gpu_fitness", fake_gpu_module + ) + monkeypatch.setitem( + sys.modules, + "runpod.serverless.modules.rp_system_fitness", + fake_system_module, + ) + + real_import = builtins.__import__ + + def guard_no_torch(name, *args, **kwargs): + if name.split(".")[0] == "torch": + raise AssertionError("torch imported during startup checks") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guard_no_torch) + + run_startup_fitness_checks() + + assert calls == ["gpu", "system"] # deferred check stays for run_worker + + +class TestImportWiring: + """Deleting the wiring must fail a test, not just real workers.""" + + def test_serverless_import_calls_startup_checks(self, monkeypatch): + import importlib + + import runpod.serverless + + calls = [] + monkeypatch.setattr( + rp_fitness, "run_startup_fitness_checks", lambda: calls.append(True) + ) + + importlib.reload(runpod.serverless) + + assert calls == [True] diff --git a/tests/test_serverless/test_utils/test_cuda.py b/tests/test_serverless/test_utils/test_cuda.py index 469c2be7..69c1aab1 100644 --- a/tests/test_serverless/test_utils/test_cuda.py +++ b/tests/test_serverless/test_utils/test_cuda.py @@ -16,7 +16,9 @@ def test_is_available_true(): "subprocess.check_output", return_value=b"NVIDIA-SMI" ) as mock_check_output: assert rp_cuda.is_available() is True - mock_check_output.assert_called_once_with(["nvidia-smi"], stderr=subprocess.DEVNULL) + mock_check_output.assert_called_once_with( + ["nvidia-smi"], stderr=subprocess.DEVNULL, timeout=5 + ) def test_is_available_false(): @@ -27,7 +29,9 @@ def test_is_available_false(): "subprocess.check_output", return_value=b"Not a GPU output" ) as mock_check_output: assert rp_cuda.is_available() is False - mock_check_output.assert_called_once_with(["nvidia-smi"], stderr=subprocess.DEVNULL) + mock_check_output.assert_called_once_with( + ["nvidia-smi"], stderr=subprocess.DEVNULL, timeout=5 + ) def test_is_available_exception(): @@ -38,4 +42,6 @@ def test_is_available_exception(): "subprocess.check_output", side_effect=Exception("Bad Command") ) as mock_check: assert rp_cuda.is_available() is False - mock_check.assert_called_once_with(["nvidia-smi"], stderr=subprocess.DEVNULL) + mock_check.assert_called_once_with( + ["nvidia-smi"], stderr=subprocess.DEVNULL, timeout=5 + ) diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py index 88f969ba..547bd88f 100644 --- a/tests/test_serverless/test_worker.py +++ b/tests/test_serverless/test_worker.py @@ -185,7 +185,7 @@ def setUp(self): fitness_patcher = patch( "runpod.serverless.worker.run_fitness_checks", new=AsyncMock() ) - fitness_patcher.start() + self.mock_fitness_checks = fitness_patcher.start() self.addCleanup(fitness_patcher.stop) # Set up the config @@ -230,6 +230,9 @@ def test_run_worker( assert not mock_stream_result.called assert mock_session.called + # The wiring this class relies on: run_worker must run fitness checks. + self.mock_fitness_checks.assert_awaited_once() + @patch("runpod.serverless.modules.rp_scale.get_job") @patch("runpod.serverless.modules.rp_job.run_job") @patch("runpod.serverless.modules.rp_job.stream_result")