From 8c2c2801bcb3a289c099e753d93a3a03a4efc38b Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 25 Aug 2026 22:31:47 -0400 Subject: [PATCH 01/13] 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 c50a9c932..6cbd02038 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 052452073..6139b4376 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 77df97e79..9fb1c95b1 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 000000000..1aaf0c859 --- /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 02/13] 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 6cbd02038..e946a4d23 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 6139b4376..c6fa605e1 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 9fb1c95b1..815f1f7c0 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 03/13] 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 e946a4d23..cc7888496 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 815f1f7c0..c025ce05b 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 8dc8946d9..1ac80f5b3 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 1aaf0c859..5dec39b3e 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 04/13] 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 2dbca2ef0..26be26813 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 cc7888496..1717fbd30 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 c025ce05b..7110fad6a 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 bae74cd88..fd6590097 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 028c7ebcd..1a47108a4 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 f8df86144..12810382a 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 5dec39b3e..afddfbfd2 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 469c2be71..69c1aab19 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 88f969baa..547bd88f2 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") From 20d1575fb758d2f358799637b285c72b29540034 Mon Sep 17 00:00:00 2001 From: justinwlin <15874969+justinwlin@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:49:43 -0400 Subject: [PATCH 05/13] test(serverless): pin registration-latch behavior; doc and docstring touch-ups - regression test: malformed RUNPOD_MIN_* must re-raise in run_worker, never fail open (latch-on-success) - fix dormant called/calls typo in the done-marker test - README: checks run once per check, not once at startup - ARCHITECTURE.md: failure path is os._exit(1), not sys.exit(1) - docs: GPU benchmark default timeout is 2s, not 100ms - rp_gpu_fitness docstring: lazy registration + truthy flag values --- ARCHITECTURE.md | 6 +++--- README.md | 2 +- docs/serverless/worker_fitness_checks.md | 6 +++--- runpod/serverless/modules/rp_gpu_fitness.py | 4 ++-- .../test_modules/test_fitness/test_startup.py | 21 ++++++++++++++++++- 5 files changed, 29 insertions(+), 10 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 26be26813..2b5d5a9ab 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -604,7 +604,7 @@ log.error(message, job_id=None) **Responsibilities**: - Validate worker health at startup before handler initialization - Support both synchronous and asynchronous check functions -- Exit immediately with sys.exit(1) on any check failure +- Exit immediately with os._exit(1) on any check failure - Enable fail-fast deployment validation **Key Functions**: @@ -617,7 +617,7 @@ log.error(message, job_id=None) 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) -5. On failure: log detailed error, call `sys.exit(1)` +5. On failure: log detailed error, best-effort unhealthy report, force-kill via `os._exit(1)` 6. On success: log completion, proceed with worker startup **Performance**: ~0.5ms framework overhead per check, total depends on check logic @@ -765,7 +765,7 @@ sequenceDiagram CHECK->>CHECK: Log success else Check fails CHECK->>SYS: Log error + traceback - CHECK->>SYS: sys.exit(1) + CHECK->>SYS: os._exit(1) end end diff --git a/README.md b/README.md index 6ad9c6669..f5328243a 100644 --- a/README.md +++ b/README.md @@ -143,7 +143,7 @@ runpod.serverless.start({"handler": handler}) **Key Features:** - Supports both synchronous and asynchronous check functions -- Checks run only once at worker startup (production mode) +- Each check runs once per worker: built-ins at import, your checks at start (production mode) - Runs before handler initialization and job processing begins - Any check failure exits with code 1 (worker marked unhealthy) diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index 1717fbd30..77ed5a593 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -360,15 +360,15 @@ ERROR | Fitness check failed: _cuda_init_check | RuntimeError: Failed to initia Quick matrix multiplication to verify GPU compute functionality and responsiveness. Skips silently on CPU-only workers. -- **Default**: 100ms maximum execution time -- **Configure**: `RUNPOD_GPU_BENCHMARK_TIMEOUT=2` +- **Default**: 2 seconds maximum execution time +- **Configure**: `RUNPOD_GPU_BENCHMARK_TIMEOUT=2` (seconds) What it tests: - GPU compute capability (matrix multiplication) - GPU response time - Memory bandwidth to GPU -If the operation takes longer than 100ms, the worker exits as the GPU is too slow for reliable job processing. +If the operation takes longer than the timeout, the worker exits as the GPU is too slow for reliable job processing. Example log output: ``` diff --git a/runpod/serverless/modules/rp_gpu_fitness.py b/runpod/serverless/modules/rp_gpu_fitness.py index fd6590097..49db979fd 100644 --- a/runpod/serverless/modules/rp_gpu_fitness.py +++ b/runpod/serverless/modules/rp_gpu_fitness.py @@ -278,12 +278,12 @@ def auto_register_gpu_check() -> None: """ Auto-register GPU fitness check if GPUs are detected. - This function is called during rp_fitness module initialization. + Called lazily on the first fitness-check run. It detects GPU presence via nvidia-smi and registers the check if found. On CPU-only workers, the check is skipped silently. Environment variables: - - RUNPOD_SKIP_GPU_CHECK: Set to "true" to skip auto-registration + - RUNPOD_SKIP_GPU_CHECK: Set to a truthy value (1/true/yes/on) to skip auto-registration """ # Allow skipping during tests if _env_flag("RUNPOD_SKIP_GPU_CHECK"): 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 afddfbfd2..dd8dc4388 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -208,7 +208,7 @@ def test_done_marker_skips_startup_pass(self, worker_env, monkeypatch): @register_fitness_check def check(): - called.append(True) + calls.append(True) run_startup_fitness_checks() assert calls == [] @@ -307,3 +307,22 @@ def test_serverless_import_calls_startup_checks(self, monkeypatch): importlib.reload(runpod.serverless) assert calls == [True] + + +class TestRegistrationLatch: + """A malformed env value must fail loudly in run_worker, not fail open.""" + + @pytest.mark.asyncio + async def test_malformed_env_reraises_at_start(self, worker_env, monkeypatch): + monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", raising=False) + monkeypatch.setenv("RUNPOD_MIN_MEMORY_GB", "not-a-number") + # Drop the cached module so the env parse re-executes on import. + monkeypatch.delitem( + sys.modules, "runpod.serverless.modules.rp_system_fitness", raising=False + ) + + run_startup_fitness_checks() # swallowed and logged — but not latched + assert rp_fitness._registration_state["system_checks"] is False + + with pytest.raises(ValueError): + await run_fitness_checks() From 85926ae8e95bd88a8468f7bab22b600051b499f7 Mon Sep 17 00:00:00 2001 From: justinwlin <15874969+justinwlin@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:37:57 -0400 Subject: [PATCH 06/13] feat(serverless): warn when fitness-check config is set after import The import-time pass consumes RUNPOD_MIN_*/RUNPOD_SKIP_*/RUNPOD_GPU_* at import; values set from the handler afterwards were silently ignored. run_fitness_checks now diffs the current env against the values snapshot at the startup pass and warns with the exact fix (set before import, or RUNPOD_DEFER_FITNESS_CHECKS=true). --- docs/serverless/worker_fitness_checks.md | 2 +- runpod/serverless/modules/rp_fitness.py | 37 +++++++++++++++++++ .../test_modules/test_fitness/conftest.py | 2 + .../test_modules/test_fitness/test_startup.py | 34 +++++++++++++++++ 4 files changed, 74 insertions(+), 1 deletion(-) diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index 77ed5a593..e6210b595 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -55,7 +55,7 @@ The import-time pass no-ops outside a real worker (no `RUNPOD_WEBHOOK_GET_JOB`), 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. +- 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. If any of them changed since the import, `start()` logs a warning naming the ignored variables. - 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 diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index 7110fad6a..4066a4053 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -66,6 +66,23 @@ def _terminate_unhealthy(code: int = 1) -> None: # inherit the environment; the marker tells them to skip the checks. _CHECKS_DONE_ENV = "RUNPOD_FITNESS_CHECKS_DONE" +# Tuning vars consumed when the checks run. Snapshotted at the import-time +# pass so a later pass can warn about post-import changes, which would +# otherwise be silently ignored. +_CONFIG_ENV_VARS = ( + "RUNPOD_MIN_MEMORY_GB", + "RUNPOD_MIN_DISK_PERCENT", + "RUNPOD_MIN_CUDA_VERSION", + "RUNPOD_NETWORK_CHECK_TIMEOUT", + "RUNPOD_GPU_BENCHMARK_TIMEOUT", + "RUNPOD_GPU_TEST_TIMEOUT", + "RUNPOD_GPU_MAX_ERROR_MESSAGES", + "RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", + "RUNPOD_SKIP_GPU_CHECK", +) + +_config_snapshot: dict[str, str | None] = {} + def _env_flag(name: str) -> bool: """True if the env var is set to a truthy value.""" @@ -248,6 +265,19 @@ def _ensure_system_checks_registered() -> None: _registration_state["system_checks"] = True +def _warn_late_config() -> None: + """Warn if tuning vars changed since the import-time pass consumed them.""" + changed = [ + name for name, old in _config_snapshot.items() if os.environ.get(name) != old + ] + if changed: + log.warn( + f"Fitness check config changed after the startup checks ran and is " + f"ignored: {', '.join(changed)}. Set these before `import runpod`, or " + f"set {DEFER_FITNESS_CHECKS_ENV}=true to run checks at start() as before." + ) + + async def run_fitness_checks(include_deferred: bool = True) -> None: """ Execute all registered fitness checks sequentially at startup. @@ -288,6 +318,9 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: log.info(f"Fitness checks disabled via {SKIP_FITNESS_CHECKS_ENV}, skipping.") return + if _config_snapshot: + _warn_late_config() + # Defer GPU check auto-registration until fitness checks are about to run # This avoids circular import issues during module initialization _ensure_gpu_check_registered() @@ -397,6 +430,10 @@ def run_startup_fitness_checks() -> None: log.debug("Event loop already running, deferring fitness checks to run_worker.") return + # Remember the tuning values as consumed, so a later pass can warn about + # post-import changes (set in the handler, too late to apply). + _config_snapshot.update({v: os.environ.get(v) for v in _CONFIG_ENV_VARS}) + try: # Own loop rather than asyncio.run: run() resets the thread's loop # policy state, after which asyncio.get_event_loop() in handler code diff --git a/tests/test_serverless/test_modules/test_fitness/conftest.py b/tests/test_serverless/test_modules/test_fitness/conftest.py index 12810382a..04a6f9e68 100644 --- a/tests/test_serverless/test_modules/test_fitness/conftest.py +++ b/tests/test_serverless/test_modules/test_fitness/conftest.py @@ -34,6 +34,8 @@ def _raise_system_exit(code=0): _reset_registration_state() clear_fitness_checks() + rp_fitness._config_snapshot.clear() yield _reset_registration_state() clear_fitness_checks() + rp_fitness._config_snapshot.clear() 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 dd8dc4388..9ecf819d9 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -326,3 +326,37 @@ async def test_malformed_env_reraises_at_start(self, worker_env, monkeypatch): with pytest.raises(ValueError): await run_fitness_checks() + + +class TestLateConfigWarning: + """Config set in the handler after the import pass must surface loudly.""" + + @staticmethod + def _run_pass(): + # Sync context like run_worker: drive the async pass on a throwaway loop. + loop = rp_fitness.asyncio.new_event_loop() + try: + loop.run_until_complete(run_fitness_checks()) + finally: + loop.close() + + def test_warns_when_config_changes_after_startup_pass( + self, worker_env, monkeypatch + ): + run_startup_fitness_checks() # consumes + snapshots config at import + + monkeypatch.setenv("RUNPOD_MIN_MEMORY_GB", "8") # too late + + with patch.object(rp_fitness.log, "warn") as mock_warn: + self._run_pass() + + warned = " ".join(str(c.args[0]) for c in mock_warn.call_args_list) + assert "RUNPOD_MIN_MEMORY_GB" in warned + + def test_no_warning_when_config_unchanged(self, worker_env): + run_startup_fitness_checks() + + with patch.object(rp_fitness.log, "warn") as mock_warn: + self._run_pass() + + mock_warn.assert_not_called() From bc89d207ea7c7806d92fcb2e6af8d2fe4c1652e0 Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 10 Sep 2026 17:23:26 -0400 Subject: [PATCH 07/13] fix(serverless): isolate and gate early worker fitness checks --- ARCHITECTURE.md | 6 +- README.md | 4 +- docs/serverless/worker_fitness_checks.md | 52 +- pyproject.toml | 1 + runpod/__init__.py | 4 + runpod/_health/__init__.py | 1 + runpod/_health/cuda.py | 22 + runpod/_health/fitness.py | 504 ++++++++++++++++ runpod/_health/gpu.py | 327 +++++++++++ runpod/_health/system.py | 555 ++++++++++++++++++ runpod/_logger.py | 155 +++++ runpod/_startup.py | 37 ++ runpod/_worker_bootstrap.py | 47 ++ runpod/serverless/__init__.py | 5 +- runpod/serverless/modules/rp_fastapi.py | 21 + runpod/serverless/modules/rp_fitness.py | 447 +------------- runpod/serverless/modules/rp_gpu_fitness.py | 320 +--------- runpod/serverless/modules/rp_logger.py | 157 +---- .../serverless/modules/rp_system_fitness.py | 519 +--------------- runpod/serverless/utils/rp_cuda.py | 24 +- .../test_fitness/test_force_kill.py | 12 +- .../test_fitness/test_safe_startup.py | 333 +++++++++++ .../test_modules/test_fitness/test_startup.py | 17 +- 23 files changed, 2079 insertions(+), 1491 deletions(-) create mode 100644 runpod/_health/__init__.py create mode 100644 runpod/_health/cuda.py create mode 100644 runpod/_health/fitness.py create mode 100644 runpod/_health/gpu.py create mode 100644 runpod/_health/system.py create mode 100644 runpod/_logger.py create mode 100644 runpod/_startup.py create mode 100644 runpod/_worker_bootstrap.py create mode 100644 tests/test_serverless/test_modules/test_fitness/test_safe_startup.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 2b5d5a9ab..998eab63c 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -599,7 +599,7 @@ log.error(message, job_id=None) ### Fitness Checks: `modules/rp_fitness.py` -**Location**: `runpod/serverless/modules/rp_fitness.py` +**Location**: `runpod/_health/fitness.py` (legacy `serverless.modules.rp_fitness` imports remain aliases) **Responsibilities**: - Validate worker health at startup before handler initialization @@ -613,11 +613,11 @@ log.error(message, job_id=None) - `clear_fitness_checks()`: Clear registry (testing only) **Execution Flow**: -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 +1. `runpod-worker` identifies the handler process and runs early hardware checks before executing it. Existing launchers may authorize the top-level import hook using `RUNPOD_FITNESS_WORKER_PID=`. The hook and check engine do not import `serverless`; ordinary imports with only the webhook environment are exempt. Legacy launches and `RUNPOD_DEFER_FITNESS_CHECKS=true` run checks only at worker start. Network readiness, CUDA initialization, compute and custom checks run in the final pass; production realtime uses the serving process's lifespan. Successful early checks are reused unless their configuration changes. 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) -5. On failure: log detailed error, best-effort unhealthy report, force-kill via `os._exit(1)` +5. On health failure: log, best-effort unhealthy report, force-kill via `os._exit(1)`. Registration is atomic; early setup errors defer, unresolved worker-start setup errors report `fitness_check_setup` and force-exit. 6. On success: log completion, proceed with worker startup **Performance**: ~0.5ms framework overhead per check, total depends on check logic diff --git a/README.md b/README.md index f5328243a..746d09fbf 100644 --- a/README.md +++ b/README.md @@ -143,7 +143,9 @@ runpod.serverless.start({"handler": handler}) **Key Features:** - Supports both synchronous and asynchronous check functions -- Each check runs once per worker: built-ins at import, your checks at start (production mode) +- `runpod-worker handler.py` checks hardware before model loading; existing Python launches check at worker start +- Local tests/helper imports remain exempt; network readiness and custom checks run at worker start +- Successful early checks are reused unless their configuration changes - Runs before handler initialization and job processing begins - Any check failure exits with code 1 (worker marked unhealthy) diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index e6210b595..d32ed76f1 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -43,20 +43,34 @@ if __name__ == "__main__": ## 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. +Early checks are automatic when the handler is launched with `runpod-worker`: -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. +```bash +runpod-worker handler.py +# Or a module: +runpod-worker -m my_package.handler +``` + +The launcher runs memory, disk, CUDA-version and native GPU checks **before executing the handler**, then runs the same handler with its original arguments. Customers do not need to add imports or check calls. Platform-managed launchers can adopt this entrypoint without changing handler code. Existing `python handler.py` launches continue checking at worker start, preserving compatibility. + +Launchers that already manage Python directly may instead set `RUNPOD_FITNESS_WORKER_PID` to the PID of the Python handler process **before exec**. The top-level `runpod` import checks that exact PID. This hook is independent of `runpod.serverless`, including when that module is lazy-loaded. Do not set a fixed PID in a Dockerfile or template. `RUNPOD_WEBHOOK_GET_JOB` alone never authorizes import-time checks. -Your own `@register_fitness_check` functions are registered after that import, so they also run at `start()`. Checks that already passed are not repeated. +Network readiness, CUDA initialization, the GPU compute benchmark, and customer-registered checks run at worker start. Network checks use bounded retries against the worker API host; they cannot terminate a process during import. CUDA checks that initialize a context remain deferred so handler code can create child processes first. -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. +Checks that passed early are not repeated unless their settings changed. Without launcher identification, all checks run at worker start. `RUNPOD_DEFER_FITNESS_CHECKS=true` restores worker-start-only timing even with the new launcher; `RUNPOD_SKIP_FITNESS_CHECKS=true` disables all checks. -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. +### Compatibility and failure handling -Two details worth knowing: +- Imports in helper scripts and ordinary local tests are safe even when they inherit worker environment variables. `--test_input` (both argument forms) and local `--rp_serve_api` invocations skip early checks. The launcher removes its process authorization before executing the handler; children cannot inherit permission to run early checks. +- Set thresholds before launch for early validation. Settings changed afterward are applied at worker start, with a warning; affected checks are rerun, while unrelated successful checks remain completed. Earlier failures cannot be undone by changing settings later. Use deferral when the handler must configure checks before they run. +- With early checks enabled, the memory check measures available memory **before model loading**. With legacy/deferred startup it measures available memory at worker start. +- Production realtime mode (`RUNPOD_REALTIME_PORT` plus worker environment) runs the final checks in the serving process's application lifespan before accepting requests. Local API simulation remains exempt. +- A failed health check reports the failure and force-exits. An error preparing early checks is logged and retried at worker start. An unresolved setup/configuration error at worker start reports `fitness_check_setup` and force-exits, including when background threads are alive. +- Early execution inside an already-running event loop defers to worker start; it does not replace the customer's event loop. -- 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. If any of them changed since the import, `start()` logs a warning naming the ignored variables. -- 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. +### Platform rollout + +Ship the SDK first with legacy launch behavior preserved. Enable `runpod-worker` in a small set of managed worker launches, validate real GPU/fork behavior and startup failure rates, then expand. Deployments with custom entrypoints retain worker-start checks until their launcher integrates the process hook. Roll back early timing centrally with `RUNPOD_DEFER_FITNESS_CHECKS=true`; no handler edits are needed. This SDK change supplies the launcher and hook; it does not change deployed platform launch configuration. ## Async Fitness Checks @@ -301,19 +315,17 @@ Disk space check passed: 50.00GB free (50.0% available) ### Network Connectivity -Tests basic internet connectivity for API calls and job processing. - -- **Default**: 5 second timeout to 8.8.8.8:53 -- **Configure**: `RUNPOD_NETWORK_CHECK_TIMEOUT=10` +Tests TCP reachability of the worker API host at worker start. -What it checks: -- Connection to Google DNS (8.8.8.8 port 53) -- Response latency -- Overall internet accessibility +- **Default**: Up to three attempts within a 5-second total connection/cleanup budget. +- **Configure**: `RUNPOD_NETWORK_CHECK_TIMEOUT=10` (positive seconds). +- **Target**: Host and port from `RUNPOD_WEBHOOK_GET_JOB`; defaults to `api.runpod.ai:443` if absent. URL paths and credentials are not sent or logged by this probe. +- Tests connection reachability, not API authentication or full application readiness. +- Retries temporary connection failures; persistent failure exits through the worker failure path. Example log output: ``` -Network connectivity passed: Connected to 8.8.8.8 (45ms) +Network connectivity passed: Connected to api.runpod.ai:443 ``` ### CUDA Version (GPU workers only) @@ -388,7 +400,7 @@ ENV RUNPOD_NETWORK_CHECK_TIMEOUT=10 ENV RUNPOD_GPU_BENCHMARK_TIMEOUT=2 ``` -Or in Python, before `import runpod` (on the real platform the checks run at import): +For legacy/deferred launches, settings can also be configured in Python before worker start: ```python import os @@ -422,7 +434,7 @@ os.environ["RUNPOD_SKIP_GPU_CHECK"] = "true" import runpod ``` -As with the thresholds, set these before `import runpod` on the real platform. +For early checks, set these before launching the handler. For legacy/deferred launches, set them before worker start. 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. @@ -430,7 +442,7 @@ User-registered checks via `@register_fitness_check` still run regardless of `RU ### Execution Timing -- 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 +- Early checks run only in launcher-identified worker processes; the final pass runs before job processing. Successful checks are reused unless their configuration changes. - 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 diff --git a/pyproject.toml b/pyproject.toml index d4639a153..6faca1cdc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ local_scheme = "no-local-version" [project.scripts] runpod = "runpod.cli.entry:runpod_cli" +runpod-worker = "runpod._worker_bootstrap:main" [dependency-groups] diff --git a/runpod/__init__.py b/runpod/__init__.py index 6d24180ae..d1ffb33ab 100644 --- a/runpod/__init__.py +++ b/runpod/__init__.py @@ -3,6 +3,10 @@ import logging import os +from ._startup import run_import_checks + +run_import_checks() + from . import serverless from .api.ctl_commands import ( create_container_registry_auth, diff --git a/runpod/_health/__init__.py b/runpod/_health/__init__.py new file mode 100644 index 000000000..d4441d68c --- /dev/null +++ b/runpod/_health/__init__.py @@ -0,0 +1 @@ +"""Worker health checks independent of the serverless import tree.""" diff --git a/runpod/_health/cuda.py b/runpod/_health/cuda.py new file mode 100644 index 000000000..1a47108a4 --- /dev/null +++ b/runpod/_health/cuda.py @@ -0,0 +1,22 @@ +""" +Provides some of the torch.cuda functionality without requiring torch. +""" + +import subprocess + + +def is_available(): + """ + Returns True if CUDA is available, False otherwise. + """ + try: + # 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 + pass + return False diff --git a/runpod/_health/fitness.py b/runpod/_health/fitness.py new file mode 100644 index 000000000..d58e5684f --- /dev/null +++ b/runpod/_health/fitness.py @@ -0,0 +1,504 @@ +""" +Fitness check system for worker startup validation. + +Fitness checks run before handler initialization on the actual RunPod serverless +platform to validate the worker environment. Any check failure force-kills the +worker via os._exit(1), signaling unhealthy state to the container orchestrator. + +Fitness checks do NOT run in local development mode or testing mode. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import os +import sys +import time +import traceback +from collections.abc import Callable + +from runpod._logger import RunPodLogger +from runpod._startup import is_worker_process + +log = RunPodLogger() + + +def _terminate_unhealthy(code: int = 1) -> None: + """ + Force-kill the worker after a fitness check failure. + + Uses os._exit rather than sys.exit because a fitness failure means the + environment is broken and the worker must die immediately so the + orchestrator can restart it. sys.exit only raises SystemExit, which + triggers cooperative interpreter shutdown and blocks joining non-daemon + threads. Workers routinely have such threads alive by the time checks run + (e.g. vLLM's AsyncLLMEngine, constructed at import before the checks), so + sys.exit can hang forever and the worker keeps serving jobs. os._exit + bypasses thread joins, atexit handlers, and asyncgen cleanup. + + Args: + code: Process exit code (default 1, signaling unhealthy). + """ + # Best-effort flush of buffered logs before the hard exit skips normal + # cleanup. A broken worker may have a closed/None stdio stream; never let a + # flush failure stop the exit, which is the whole point of this helper. + for stream in (sys.stdout, sys.stderr): + with contextlib.suppress(Exception): + stream.flush() + os._exit(code) + + +# 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" + +# Tuning vars consumed when the checks run. Snapshotted at the import-time +# pass so a later pass can warn about post-import changes, which would +# otherwise be silently ignored. +_CONFIG_ENV_VARS = ( + "RUNPOD_MIN_MEMORY_GB", + "RUNPOD_MIN_DISK_PERCENT", + "RUNPOD_MIN_CUDA_VERSION", + "RUNPOD_NETWORK_CHECK_TIMEOUT", + "RUNPOD_GPU_BENCHMARK_TIMEOUT", + "RUNPOD_GPU_TEST_TIMEOUT", + "RUNPOD_GPU_MAX_ERROR_MESSAGES", + "RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", + "RUNPOD_SKIP_GPU_CHECK", +) + +_config_snapshot: dict[str, str | None] = {} + + +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: + """ + Decorator to register a fitness check function. + + Fitness checks validate worker health at startup before handler initialization. + If any check fails, the worker is force-killed with os._exit(1). + + Supports both sync and async functions (auto-detected via inspect.iscoroutinefunction()). + + Example: + @runpod.serverless.register_fitness_check + def check_gpu(): + import torch + if not torch.cuda.is_available(): + raise RuntimeError("GPU not available") + + @runpod.serverless.register_fitness_check + async def check_model_files(): + import aiofiles.os + if not await aiofiles.os.path.exists("/models/model.safetensors"): + raise RuntimeError("Model file not found") + + Args: + func: Function to register as fitness check. Can be sync or async. + + Returns: + Original function unchanged (allows decorator stacking). + """ + _fitness_checks.append(func) + log.debug(f"Registered fitness check: {func.__name__}") + return func + + +def clear_fitness_checks() -> None: + """ + Clear all registered fitness checks. + + Used primarily for testing to reset global state between test cases. + Not intended for production use. + """ + _fitness_checks.clear() + _completed_checks.clear() + + +_registration_state: dict[str, bool] = { + "gpu_check": False, + "system_checks": False, +} + + +def _reset_registration_state() -> None: + """ + Reset global registration state. + + Used for testing to ensure clean state between tests. + """ + _registration_state["gpu_check"] = False + _registration_state["system_checks"] = False + + +# Bound how long the best-effort unhealthy report may delay the exit. +_REPORT_TIMEOUT_SECONDS = 2 + + +def _report_unhealthy(check: str, reason: str) -> None: + """ + Best-effort report of a fitness-check failure to the host before exit. + + Sends a single GET to the ping URL (same URL/credentials the heartbeat + uses) with status=unhealthy plus the failing check name and reason, so the + host can emit a queryable worker.fitness_failed event. Any failure — no + ping URL, no API key, HTTP error, timeout — is swallowed, so this can never + prevent the os._exit that follows. It is synchronous, so it may delay that + exit by up to _REPORT_TIMEOUT_SECONDS (network phases only; it adds no + delay when there is no ping URL/API key to report to). + """ + ping_url = os.environ.get("RUNPOD_WEBHOOK_PING") + api_key = os.environ.get("RUNPOD_AI_API_KEY") + if not ping_url or ping_url == "PING_NOT_SET" or not api_key: + return + + try: + # Deferred imports: keep module import light and avoid import cycles. + from requests import Session + from runpod.version import __version__ as runpod_version + + worker_id = os.environ.get("RUNPOD_POD_ID") + if "$RUNPOD_POD_ID" in ping_url and not worker_id: + return + ping_url = ping_url.replace("$RUNPOD_POD_ID", worker_id or "") + params = { + "status": "unhealthy", + "check": check, + "reason": reason[:256], + "runpod_version": runpod_version, + } + session = Session() + try: + session.headers.update({"Authorization": api_key}) + session.get(ping_url, params=params, timeout=_REPORT_TIMEOUT_SECONDS) + finally: + session.close() + except Exception: + # Best-effort only; the exit is the guarantee, not this report. + pass + + +def _ensure_gpu_check_registered() -> None: + """ + Ensure GPU fitness check is registered. + + Deferred until first run to avoid circular import issues during module + initialization. Called from run_fitness_checks() on first invocation. + """ + if _registration_state["gpu_check"]: + return + + # 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. + from .gpu import auto_register_gpu_check + + before = len(_fitness_checks) + auto_register_gpu_check() + for check in _fitness_checks[before:]: + check._runpod_builtin = "gpu_check" + _registration_state["gpu_check"] = True + + +def _ensure_system_checks_registered() -> None: + """ + Ensure system resource fitness checks are registered. + + Deferred until first run to avoid circular import issues during module + initialization. Called from run_fitness_checks() on first invocation. + """ + if _registration_state["system_checks"]: + return + + # Allow disabling system checks for testing + 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 + + # Same latch-on-success rule as _ensure_gpu_check_registered. + from .system import auto_register_system_checks + + before = len(_fitness_checks) + auto_register_system_checks() + for check in _fitness_checks[before:]: + check._runpod_builtin = "system_checks" + _registration_state["system_checks"] = True + + +def _register_builtins() -> None: + """Register atomically: failed setup must not leave duplicate/partial checks.""" + before = list(_fitness_checks) + state = dict(_registration_state) + try: + _ensure_gpu_check_registered() + _ensure_system_checks_registered() + except Exception: + _fitness_checks[:] = before + _registration_state.update(state) + raise + + +def _refresh_late_config() -> None: + """Apply changed settings and rerun only checks whose inputs changed.""" + changed = { + name for name, old in _config_snapshot.items() if os.environ.get(name) != old + } + if not changed: + return + log.warn( + "Fitness check config changed since early checks; applying at worker start: " + + ", ".join(sorted(changed)) + ) + # Runtime tuning lives in the standalone check modules, not frozen imports. + if not _env_flag("RUNPOD_SKIP_GPU_CHECK"): + from . import gpu + + gpu.configure() + if not _env_flag("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"): + from . import system + + system.configure() + dependencies = { + "_memory_check": {"RUNPOD_MIN_MEMORY_GB"}, + "_disk_check": {"RUNPOD_MIN_DISK_PERCENT"}, + "_cuda_version_check": {"RUNPOD_MIN_CUDA_VERSION"}, + "_network_check": {"RUNPOD_NETWORK_CHECK_TIMEOUT"}, + "_benchmark_check": {"RUNPOD_GPU_BENCHMARK_TIMEOUT"}, + "_gpu_health_check": { + "RUNPOD_GPU_TEST_TIMEOUT", + "RUNPOD_GPU_MAX_ERROR_MESSAGES", + }, + } + _completed_checks[:] = [ + check + for check in _completed_checks + if not ( + getattr(check, "_runpod_builtin", False) + and dependencies.get(check.__name__, set()) & changed + ) + ] + for flag, group in ( + ("RUNPOD_SKIP_GPU_CHECK", "gpu_check"), + ("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "system_checks"), + ): + if flag in changed: + removed = [ + c + for c in _fitness_checks + if getattr(c, "_runpod_builtin", None) == group + ] + _fitness_checks[:] = [ + c for c in _fitness_checks if not any(c is old for old in removed) + ] + _completed_checks[:] = [ + c for c in _completed_checks if not any(c is old for old in removed) + ] + _registration_state[group] = False + _config_snapshot.update({v: os.environ.get(v) for v in _CONFIG_ENV_VARS}) + + +def _fail_worker(check_name: str, exc: Exception) -> None: + """Report a check/setup failure, then exit even if reporting or logging fails.""" + try: + reason = f"{type(exc).__name__}: {exc}" + with contextlib.suppress(Exception): + log.error(f"Fitness check failed: {check_name} | {reason}") + log.debug(f"Traceback: {traceback.format_exc()}") + with contextlib.suppress(Exception): + _report_unhealthy(check_name, reason) + with contextlib.suppress(Exception): + log.error("Worker is unhealthy, exiting.") + finally: + _terminate_unhealthy(1) + + +async def run_fitness_checks(include_deferred: bool = True) -> None: + """ + Execute all registered fitness checks sequentially at startup. + + Execution flow: + 1. Auto-register GPU check on first run (deferred to avoid circular imports) + 2. Check if registry is empty (early return if no checks) + 3. Log start of fitness check phase + 4. For each registered check: + - Auto-detect sync vs async using inspect.iscoroutinefunction() + - Execute check with timing instrumentation (await if async, call if sync) + - Log success or failure with check name and execution time + 5. On any exception: + - Log detailed error with check name, exception type, and message + - Log traceback at DEBUG level + - Force-kill the worker via os._exit(1) immediately (fail-fast). This is + a hard exit, not a cooperative sys.exit/SystemExit: it does not unwind + the stack or run cleanup, so callers cannot catch it and it cannot be + blocked by live non-daemon threads. + 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 + and handles checks with dependencies correctly. + Timing uses high-precision perf_counter for accurate measurements. + + Note: + 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 + + try: + if include_deferred and _config_snapshot: + _refresh_late_config() + _register_builtins() + except Exception as exc: + if not include_deferred: + log.error( + f"Fitness checks could not be prepared; retrying at worker start: {exc}" + ) + return + _fail_worker("fitness_check_setup", exc) + return + + # 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(pending)} fitness check(s)...") + + total_start_time = time.perf_counter() + + for check_func in pending: + check_name = check_func.__name__ + + try: + log.debug(f"Executing fitness check: {check_name}") + check_start_time = time.perf_counter() + + # Auto-detect async vs sync using inspect + if inspect.iscoroutinefunction(check_func): + await check_func() + else: + 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: + _fail_worker(check_name, exc) + return + + 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 without launcher process authorization, for local tests, when checks + are disabled or deferred, inside a running event loop, and in child + processes (the launcher PID must match and is consumed before the handler). 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 is_worker_process(): + return + + if os.environ.get(_CHECKS_DONE_ENV): + return + os.environ[_CHECKS_DONE_ENV] = "1" + os.environ.pop("RUNPOD_FITNESS_WORKER_PID", None) + + if _event_loop_running(): + log.debug("Event loop already running, deferring fitness checks to run_worker.") + return + + # Remember the tuning values as consumed, so a later pass can warn about + # post-import changes (set in the handler, too late to apply). + _config_snapshot.update({v: os.environ.get(v) for v in _CONFIG_ENV_VARS}) + + 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}") diff --git a/runpod/_health/gpu.py b/runpod/_health/gpu.py new file mode 100644 index 000000000..fe42b387e --- /dev/null +++ b/runpod/_health/gpu.py @@ -0,0 +1,327 @@ +""" +GPU fitness check system for worker startup validation. + +Provides comprehensive GPU health checking using: +1. Native CUDA binary (gpu_test) for memory allocation testing +2. Python fallback using nvidia-smi if binary unavailable + +Auto-registers when GPUs are detected, skips silently on CPU-only workers. +""" + +from __future__ import annotations + +import asyncio +import os +import subprocess +from pathlib import Path +from typing import Any + +from runpod._binary_helpers import get_binary_path +from .fitness import _env_flag, register_fitness_check +from runpod._logger import RunPodLogger + +log = RunPodLogger() + +# Defaults are safe to import; parse user settings when registering checks. +TIMEOUT_SECONDS = 30 +MAX_ERROR_MESSAGES = 10 + + +def configure() -> None: + """Read current fitness settings; setup errors are handled by the runner.""" + global TIMEOUT_SECONDS, MAX_ERROR_MESSAGES + TIMEOUT_SECONDS = int(os.environ.get("RUNPOD_GPU_TEST_TIMEOUT", "30")) + MAX_ERROR_MESSAGES = int(os.environ.get("RUNPOD_GPU_MAX_ERROR_MESSAGES", "10")) + + +def _get_gpu_test_binary_path() -> Path | None: + """ + Locate gpu_test binary in package. + + Returns: + Path to binary if found, None otherwise + """ + return get_binary_path("gpu_test") + + +def _parse_gpu_test_output(output: str) -> dict[str, Any]: + """ + Parse gpu_test binary output and detect success/failure. + + Looks for: + - "GPU X memory allocation test passed." for success + - Error patterns: "Failed", "error", "cannot" for failures + - GPU count from "Found X GPUs:" line + + Args: + output: Stdout from gpu_test binary + + Returns: + Dict with keys: + - success: bool - True if all GPUs passed tests + - gpu_count: int - Number of GPUs that passed tests + - found_gpus: int - Total GPUs found + - errors: List[str] - Error messages from output + - details: Dict - CUDA version, kernel version, etc + """ + lines = output.strip().split("\n") + + result = { + "success": False, + "gpu_count": 0, + "found_gpus": 0, + "errors": [], + "details": {}, + } + + passed_count = 0 + found_gpus = 0 + + for line in lines: + line = line.strip() + if not line: + continue + + # Extract metadata + if line.startswith("CUDA Driver Version:"): + result["details"]["cuda_version"] = line.split(":", 1)[1].strip() + elif line.startswith("Linux Kernel Version:"): + result["details"]["kernel"] = line.split(":", 1)[1].strip() + elif line.startswith("Found") and "GPUs" in line: + # "Found 2 GPUs:" + try: + found_gpus = int(line.split()[1]) + result["found_gpus"] = found_gpus + except (IndexError, ValueError): + # Line format doesn't match expected "Found N GPUs:" — skip + pass + + # Check for success + if "memory allocation test passed" in line.lower(): + passed_count += 1 + + # Check for errors + if any(err in line.lower() for err in ["failed", "error", "cannot", "unable"]): + result["errors"].append(line) + + result["gpu_count"] = passed_count + result["success"] = ( + passed_count > 0 and passed_count == found_gpus and len(result["errors"]) == 0 + ) + + return result + + +async def _run_gpu_test_binary() -> dict[str, Any]: + """ + Execute gpu_test binary and parse output. + + Returns: + Parsed result dict from _parse_gpu_test_output + + Raises: + RuntimeError: If binary execution fails or GPUs unhealthy + """ + binary_path = _get_gpu_test_binary_path() + + if not binary_path: + raise FileNotFoundError("gpu_test binary not found in package") + + if not os.access(binary_path, os.X_OK): + raise PermissionError(f"gpu_test binary not executable: {binary_path}") + + log.debug(f"Running gpu_test binary: {binary_path}") + + try: + # Run binary with timeout + process = await asyncio.create_subprocess_exec( + str(binary_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=TIMEOUT_SECONDS + ) + + output = stdout.decode("utf-8", errors="replace") + error_output = stderr.decode("utf-8", errors="replace") + + log.debug(f"gpu_test output:\n{output}") + + if error_output: + log.debug(f"gpu_test stderr:\n{error_output}") + + # Parse output + result = _parse_gpu_test_output(output) + + # Check for success + if not result["success"]: + error_msg = "GPU memory allocation test failed" + if result["errors"]: + error_msg += f": {'; '.join(result['errors'][:MAX_ERROR_MESSAGES])}" + raise RuntimeError(error_msg) + + log.info( + f"GPU binary test passed: {result['gpu_count']} GPU(s) healthy " + f"(CUDA {result['details'].get('cuda_version', 'unknown')})" + ) + + return result + + except asyncio.TimeoutError: + process.kill() + await process.wait() + raise RuntimeError( + f"GPU test binary timed out after {TIMEOUT_SECONDS}s" + ) from None + except FileNotFoundError: + raise + except PermissionError: + raise + except Exception as exc: + raise RuntimeError(f"GPU test binary execution failed: {exc}") from exc + + +def _run_gpu_test_fallback() -> None: + """ + Python fallback for GPU testing using nvidia-smi. + + Less comprehensive than binary (doesn't test memory allocation) but validates + basic GPU availability by checking GPU count. + + Raises: + RuntimeError: If GPUs not available or unhealthy + """ + log.debug("Running Python GPU fallback check") + + try: + # List GPUs to verify availability and count + result = subprocess.run( + ["nvidia-smi", "--list-gpus"], + capture_output=True, + text=True, + timeout=10, + check=False, + ) + + if result.returncode != 0: + raise RuntimeError(f"nvidia-smi --list-gpus failed: {result.stderr}") + + gpu_lines = [line for line in result.stdout.split("\n") if line.strip()] + gpu_count = len(gpu_lines) + + if gpu_count == 0: + raise RuntimeError("No GPUs detected by nvidia-smi") + + log.info( + f"GPU fallback check passed: {gpu_count} GPU(s) detected " + "(Note: Memory allocation NOT tested)" + ) + + except FileNotFoundError: + raise RuntimeError( + "nvidia-smi not found. Cannot validate GPU availability." + ) from None + except subprocess.TimeoutExpired: + raise RuntimeError("nvidia-smi timed out") from None + except RuntimeError: + raise + except Exception as e: + raise RuntimeError(f"nvidia-smi fallback check failed: {e}") from e + + +async def _check_gpu_health() -> None: + """ + Comprehensive GPU health check (internal implementation). + + 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 + + Raises: + RuntimeError: If GPU health check fails + """ + binary_attempted = False + binary_error = None + + # Try binary first + try: + await _run_gpu_test_binary() + return # Success! + except FileNotFoundError as exc: + log.debug(f"GPU binary not found: {exc}") + binary_error = exc + except PermissionError as exc: + log.debug(f"GPU binary not executable: {exc}") + binary_error = exc + except Exception as exc: + log.warn(f"GPU binary check failed: {exc}") + binary_attempted = True + binary_error = exc + + # Fall back to Python + log.debug("Attempting Python GPU fallback check") + try: + _run_gpu_test_fallback() + return # Success! + except Exception as fallback_exc: + # Both failed - raise composite error + if binary_attempted: + raise RuntimeError( + f"GPU health check failed. " + f"Binary test: {binary_error}. " + f"Fallback test: {fallback_exc}" + ) from fallback_exc + else: + raise RuntimeError( + f"GPU health check failed (binary disabled/missing, " + f"fallback failed): {fallback_exc}" + ) from fallback_exc + + +def auto_register_gpu_check() -> None: + """ + Auto-register GPU fitness check if GPUs are detected. + + Called lazily on the first fitness-check run. + It detects GPU presence via nvidia-smi and registers the check if found. + On CPU-only workers, the check is skipped silently. + + Environment variables: + - RUNPOD_SKIP_GPU_CHECK: Set to a truthy value (1/true/yes/on) to skip auto-registration + """ + # Allow skipping during tests + if _env_flag("RUNPOD_SKIP_GPU_CHECK"): + log.debug("GPU fitness check auto-registration disabled via environment") + return + + configure() + + # Quick GPU detection + has_gpu = False + try: + result = subprocess.run( + ["nvidia-smi"], + capture_output=True, + text=True, + timeout=5, + check=False, + ) + has_gpu = result.returncode == 0 and "NVIDIA-SMI" in result.stdout + except (FileNotFoundError, subprocess.TimeoutExpired): + has_gpu = False + except Exception: + # Catch any other exceptions and assume no GPU + has_gpu = False + + if has_gpu: + log.debug("GPU detected, registering automatic GPU fitness check") + + @register_fitness_check + async def _gpu_health_check(): + """Automatic GPU memory allocation health check.""" + await _check_gpu_health() + else: + log.debug("No GPU detected, skipping GPU fitness check registration") diff --git a/runpod/_health/system.py b/runpod/_health/system.py new file mode 100644 index 000000000..0c47af4cc --- /dev/null +++ b/runpod/_health/system.py @@ -0,0 +1,555 @@ +""" +System resource fitness checks for worker startup validation. + +Provides comprehensive checks for: +- Memory availability +- Disk space +- Network connectivity +- CUDA library versions +- GPU compute benchmark + +Auto-registers when worker starts, ensuring system readiness before accepting jobs. +""" + +from __future__ import annotations + +import asyncio +import os +import re +import shutil +import time +from urllib.parse import urlsplit + +from .fitness import defer_to_worker_start, register_fitness_check +from runpod._logger import RunPodLogger +from .cuda import is_available as gpu_available + +log = RunPodLogger() + +# Defaults are safe to import; parse user settings when registering checks. +MIN_MEMORY_GB = 4.0 +MIN_DISK_PERCENT = 10.0 +MIN_CUDA_VERSION = "11.8" +NETWORK_CHECK_TIMEOUT = 5 +GPU_BENCHMARK_TIMEOUT = 2 + + +def configure() -> None: + """Read current fitness settings; setup errors are handled by the runner.""" + global \ + MIN_MEMORY_GB, \ + MIN_DISK_PERCENT, \ + MIN_CUDA_VERSION, \ + NETWORK_CHECK_TIMEOUT, \ + GPU_BENCHMARK_TIMEOUT + MIN_MEMORY_GB = float(os.environ.get("RUNPOD_MIN_MEMORY_GB", "4.0")) + MIN_DISK_PERCENT = float(os.environ.get("RUNPOD_MIN_DISK_PERCENT", "10.0")) + MIN_CUDA_VERSION = os.environ.get("RUNPOD_MIN_CUDA_VERSION", "11.8") + NETWORK_CHECK_TIMEOUT = int(os.environ.get("RUNPOD_NETWORK_CHECK_TIMEOUT", "5")) + GPU_BENCHMARK_TIMEOUT = int(os.environ.get("RUNPOD_GPU_BENCHMARK_TIMEOUT", "2")) + if NETWORK_CHECK_TIMEOUT <= 0 or GPU_BENCHMARK_TIMEOUT <= 0: + raise ValueError( + "RUNPOD_NETWORK_CHECK_TIMEOUT and RUNPOD_GPU_BENCHMARK_TIMEOUT must be positive" + ) + + +def _parse_version(version_string: str) -> tuple[int, int]: + """ + Parse version string to tuple for comparison. + + Args: + version_string: Version string like "12.2" or "CUDA Version 12.2" + + Returns: + Tuple of ints like (12, 2) for comparison + """ + # Extract numeric version + match = re.search(r"(\d+)\.(\d+)", version_string) + if match: + return (int(match.group(1)), int(match.group(2))) + return (0, 0) + + +def _get_memory_info() -> dict[str, float]: + """ + Get system memory information. + + Returns: + Dict with total_gb, available_gb, used_percent + + Raises: + RuntimeError: If memory check fails + """ + try: + import psutil + + mem = psutil.virtual_memory() + total_gb = mem.total / (1024**3) + available_gb = mem.available / (1024**3) + used_percent = mem.percent + + return { + "total_gb": total_gb, + "available_gb": available_gb, + "used_percent": used_percent, + } + except ImportError: + # Fallback: parse /proc/meminfo + try: + with open("/proc/meminfo") as f: + meminfo_kb: dict[str, int] = {} + for line in f: + key, value = line.split(":", 1) + meminfo_kb[key.strip()] = int(value.split()[0]) + + # /proc/meminfo values are in kB; convert to GB + total_gb = meminfo_kb.get("MemTotal", 0) / (1024**2) + available_gb = meminfo_kb.get("MemAvailable", 0) / (1024**2) + used_percent = ( + 100 * (1 - available_gb / total_gb) if total_gb > 0 else 0 + ) + + return { + "total_gb": total_gb, + "available_gb": available_gb, + "used_percent": used_percent, + } + except Exception as e: + raise RuntimeError(f"Failed to read memory info: {e}") from e + + +def _check_memory_availability() -> None: + """ + Check system memory availability. + + Raises: + RuntimeError: If insufficient memory available + """ + mem_info = _get_memory_info() + available_gb = mem_info["available_gb"] + total_gb = mem_info["total_gb"] + + if available_gb < MIN_MEMORY_GB: + raise RuntimeError( + f"Insufficient memory: {available_gb:.2f}GB available, " + f"{MIN_MEMORY_GB}GB required" + ) + + log.info( + f"Memory check passed: {available_gb:.2f}GB available " + f"(of {total_gb:.2f}GB total)" + ) + + +def _check_disk_space() -> None: + """ + Check disk space availability on root filesystem. + + In containers, root (/) is typically the only filesystem. + Requires free space to be at least MIN_DISK_PERCENT% of total disk size. + + Raises: + RuntimeError: If insufficient disk space + """ + try: + usage = shutil.disk_usage("/") + total_gb = usage.total / (1024**3) + free_gb = usage.free / (1024**3) + free_percent = 100 * (free_gb / total_gb) if total_gb > 0 else 0 + + # Check if free space is below the required percentage + if free_percent < MIN_DISK_PERCENT: + raise RuntimeError( + f"Insufficient disk space: {free_gb:.2f}GB free " + f"({free_percent:.1f}%), {MIN_DISK_PERCENT}% required" + ) + + log.info( + f"Disk space check passed: {free_gb:.2f}GB free " + f"({free_percent:.1f}% available)" + ) + except FileNotFoundError: + raise RuntimeError( + "Could not check disk space: / filesystem not found" + ) from None + + +async def _check_network_connectivity() -> None: + """Probe the worker API host with three attempts within one time budget. + + This is a worker-start readiness check, never an import-time hard failure. + TCP reachability is a basic check, not a guarantee of API authentication or + application readiness. Do not send job requests or expose URL credentials. + """ + target = urlsplit( + os.environ.get("RUNPOD_WEBHOOK_GET_JOB") or "https://api.runpod.ai" + ) + if target.scheme not in ("http", "https") or not target.hostname: + raise RuntimeError("Invalid worker API URL for network connectivity check") + host = target.hostname + port = target.port or (443 if target.scheme == "https" else 80) + + async def probe() -> None: + _, writer = await asyncio.open_connection(host, port) + try: + writer.close() + await writer.wait_closed() + finally: + # Bound connection teardown too; a stuck close must not hang startup. + if writer.transport: + writer.transport.abort() + + deadline = time.monotonic() + NETWORK_CHECK_TIMEOUT + last_error = "Timeout" + for attempt in range(3): + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + await asyncio.wait_for(probe(), timeout=remaining / (3 - attempt)) + log.info(f"Network connectivity passed: Connected to {host}:{port}") + return + except asyncio.TimeoutError: + last_error = "Timeout" + except ConnectionRefusedError: + last_error = "Connection refused" + except OSError as exc: + last_error = type(exc).__name__ + if attempt < 2: + await asyncio.sleep( + min(0.1 * (attempt + 1), max(0, deadline - time.monotonic())) + ) + raise RuntimeError( + f"Network connectivity failed: {last_error} connecting to {host}:{port} " + f"after bounded retries ({NETWORK_CHECK_TIMEOUT}s budget)" + ) + + +async def _get_cuda_version() -> str | None: + """ + Get CUDA version from system. + + Returns: + Version string like "12.2" or None if not available + + Raises: + RuntimeError: If CUDA check fails critically + """ + # Try nvcc first + process = None + try: + process = await asyncio.create_subprocess_exec( + "nvcc", + "--version", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, _ = await asyncio.wait_for(process.communicate(), timeout=5) + if process.returncode == 0: + output = stdout.decode("utf-8", errors="replace") + for line in output.split("\n"): + if "release" in line.lower() or "version" in line.lower(): + return line.strip() + except Exception as e: + if process and process.returncode is None: + process.kill() + await process.wait() + log.debug(f"nvcc not available: {e}") + + # Fallback: try nvidia-smi and parse CUDA version from output + process = None + try: + process = await asyncio.create_subprocess_exec( + "nvidia-smi", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, _ = await asyncio.wait_for(process.communicate(), timeout=5) + if process.returncode == 0: + output = stdout.decode("utf-8", errors="replace") + for line in output.split("\n"): + if "CUDA Version:" in line: + parts = line.split("CUDA Version:") + if len(parts) > 1: + cuda_version = parts[1].strip().split()[0] + return f"CUDA Version: {cuda_version}" + log.debug("nvidia-smi output found but couldn't parse CUDA version") + except Exception as e: + if process and process.returncode is None: + process.kill() + await process.wait() + log.debug(f"nvidia-smi not available: {e}") + + return None + + +async def _check_cuda_versions() -> None: + """ + Check CUDA library versions meet minimum requirements. + + Raises: + RuntimeError: If CUDA version is below minimum + """ + cuda_version_str = await _get_cuda_version() + + if not cuda_version_str: + log.warn("Could not determine CUDA version, skipping check") + return + + # Parse version + cuda_version = _parse_version(cuda_version_str) + min_version = _parse_version(MIN_CUDA_VERSION) + + if cuda_version < min_version: + raise RuntimeError( + f"CUDA version too old: {cuda_version[0]}.{cuda_version[1]} found, " + f"{min_version[0]}.{min_version[1]} required" + ) + + log.info( + f"CUDA version check passed: {cuda_version[0]}.{cuda_version[1]} " + f"(minimum: {min_version[0]}.{min_version[1]})" + ) + + +async def _check_cuda_initialization() -> None: + """ + Verify CUDA can be initialized and devices are accessible. + + 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. + + Raises: + RuntimeError: If CUDA initialization or device access fails + """ + # 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 + + if not torch.cuda.is_available(): + log.debug("CUDA not available in PyTorch, skipping initialization check") + return + + # Reset CUDA state to ensure clean initialization + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + + # Verify device count + device_count = torch.cuda.device_count() + if device_count == 0: + raise RuntimeError( + "No CUDA devices available despite cuda.is_available() being True" + ) + + # Test each device + for i in range(device_count): + try: + # Get device properties + props = torch.cuda.get_device_properties(i) + if props.total_memory == 0: + raise RuntimeError(f"GPU {i} reports zero memory") + + # Try allocating a small tensor on the device + _ = torch.zeros(1024, device=f"cuda:{i}") + torch.cuda.synchronize() + + except Exception as e: + raise RuntimeError(f"Failed to initialize GPU {i}: {e}") from e + + log.info( + f"CUDA initialization passed: {device_count} device(s) initialized successfully" + ) + return + + except ImportError: + log.debug("PyTorch not available, trying CuPy...") + except Exception as e: + raise RuntimeError(f"CUDA initialization failed: {e}") from e + + # Fallback: try CuPy + try: + import cupy as cp + + # Reset CuPy state + cp.cuda.Device().synchronize() + + # Verify devices + device_count = cp.cuda.runtime.getDeviceCount() + if device_count == 0: + raise RuntimeError("No CUDA devices available via CuPy") + + # Test each device + for i in range(device_count): + try: + cp.cuda.Device(i).use() + # Try allocating memory + _ = cp.zeros(1024) + cp.cuda.Device().synchronize() + except Exception as e: + raise RuntimeError( + f"Failed to initialize GPU {i} with CuPy: {e}" + ) from e + + log.info( + f"CUDA initialization passed: {device_count} device(s) initialized successfully" + ) + return + + except ImportError: + log.debug("CuPy not available, skipping CUDA initialization check") + except Exception as e: + raise RuntimeError(f"CUDA initialization check failed: {e}") from e + + +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. + + Raises: + RuntimeError: If GPU compute fails or is too slow + """ + # 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 + + if not torch.cuda.is_available(): + log.debug("CUDA not available in PyTorch, skipping benchmark") + return + + # Create small matrix on GPU + size = 1024 + start_time = time.perf_counter() + + # Do computation + A = torch.randn(size, size, device="cuda") + B = torch.randn(size, size, device="cuda") + torch.matmul(A, B) + torch.cuda.synchronize() # Wait for GPU to finish + + elapsed_ms = (time.perf_counter() - start_time) * 1000 + max_ms = GPU_BENCHMARK_TIMEOUT * 1000 + + if elapsed_ms > max_ms: + raise RuntimeError( + f"GPU compute too slow: Matrix multiply took {elapsed_ms:.0f}ms " + f"(max: {max_ms:.0f}ms)" + ) + + log.info( + f"GPU compute benchmark passed: Matrix multiply completed in {elapsed_ms:.0f}ms" + ) + return + + except ImportError: + log.debug("PyTorch not available, trying CuPy...") + except RuntimeError: + raise # Benchmark failure is what we're testing for + except Exception as e: + log.warn(f"PyTorch GPU benchmark setup failed: {e}") + + # Fallback: try CuPy + try: + import cupy as cp + + size = 1024 + start_time = time.perf_counter() + + A = cp.random.randn(size, size) + B = cp.random.randn(size, size) + cp.matmul(A, B) + cp.cuda.Device().synchronize() + + elapsed_ms = (time.perf_counter() - start_time) * 1000 + max_ms = GPU_BENCHMARK_TIMEOUT * 1000 + + if elapsed_ms > max_ms: + raise RuntimeError( + f"GPU compute too slow: Matrix multiply took {elapsed_ms:.0f}ms " + f"(max: {max_ms:.0f}ms)" + ) + + log.info( + f"GPU compute benchmark passed: Matrix multiply completed in {elapsed_ms:.0f}ms" + ) + return + + except ImportError: + log.debug("CuPy not available, skipping GPU benchmark") + except RuntimeError: + raise # Benchmark failure is what we're testing for + except Exception as e: + log.warn(f"CuPy GPU benchmark setup failed: {e}") + + # If we get here, neither library is available + log.debug( + "PyTorch/CuPy not available for GPU benchmark, relying on gpu_test binary" + ) + + +def auto_register_system_checks() -> None: + """ + Auto-register system resource fitness checks. + + 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. + """ + configure() + log.debug("Registering system resource fitness checks") + + # Always register these checks + @register_fitness_check + def _memory_check() -> None: + """System memory availability check.""" + _check_memory_availability() + + @register_fitness_check + def _disk_check() -> None: + """System disk space check.""" + _check_disk_space() + + @register_fitness_check + @defer_to_worker_start + async def _network_check() -> None: + """Network connectivity check.""" + await _check_network_connectivity() + + # Only register GPU checks if GPU is detected + if gpu_available(): + log.debug("GPU detected, registering GPU-specific fitness checks") + + @register_fitness_check + async def _cuda_version_check() -> None: + """CUDA version check.""" + 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() + else: + log.debug("No GPU detected, skipping GPU-specific fitness checks") diff --git a/runpod/_logger.py b/runpod/_logger.py new file mode 100644 index 000000000..6ef4c5f73 --- /dev/null +++ b/runpod/_logger.py @@ -0,0 +1,155 @@ +""" +PodWorker | modules | logging.py + +Log Levels (Level - Value - Description) + +NOTSET - 0 - No logging is configured, the logging system is effectively disabled. +DEBUG - 1 - Detailed information, typically of interest only when diagnosing problems. (Default) +INFO - 2 - Confirmation that things are working as expected. +WARN - 3 - An indication that something unexpected happened. +ERROR - 4 - Serious problem, the software has not been able to perform some function. +""" + +from contextvars import ContextVar, Token +import json +import os +from typing import Optional + +MAX_MESSAGE_LENGTH = 4096 +LOG_LEVELS = ["NOTSET", "TRACE", "DEBUG", "INFO", "WARN", "ERROR"] +_batch_id: ContextVar[Optional[str]] = ContextVar("runpod_batch_id", default=None) + + +def _set_batch_id(batch_id: Optional[str]) -> Token: + """Set the batch ID associated with the current job task.""" + return _batch_id.set(batch_id) + + +def _reset_batch_id(token: Token): + """Restore the previous batch ID for the current job task.""" + _batch_id.reset(token) + + +def _validate_log_level(log_level): + """ + Checks the debug level and returns the debug level name. + """ + if isinstance(log_level, str): + log_level = log_level.upper() + + if log_level not in LOG_LEVELS: + raise ValueError(f"Invalid debug level: {log_level}") + + return log_level + + if isinstance(log_level, int): + if log_level < 0 or log_level >= len(LOG_LEVELS): + raise ValueError(f"Invalid debug level: {log_level}") + + return LOG_LEVELS[log_level] + + raise ValueError(f"Invalid debug level: {log_level}") + + +class RunPodLogger: + """Singleton class for logging.""" + + __instance = None + level = _validate_log_level( + os.environ.get( + "RUNPOD_LOG_LEVEL", os.environ.get("RUNPOD_DEBUG_LEVEL", "DEBUG") + ) + ) + + def __new__(cls): + if RunPodLogger.__instance is None: + RunPodLogger.__instance = object.__new__(cls) + return RunPodLogger.__instance + + def set_level(self, new_level): + """ + Set the debug level for logging. + Can be set to the name or value of the debug level. + """ + self.level = _validate_log_level(new_level) + self.info(f"Log level set to {self.level}") + + def log(self, message, message_level="INFO", job_id=None): + """ + Log message to stdout if RUNPOD_DEBUG is true. + """ + if self.level == "NOTSET": + return + + level_index = LOG_LEVELS.index(self.level) + if level_index > LOG_LEVELS.index(message_level) and message_level != "TIP": + return + + message = str(message) + if batch_id := _batch_id.get(): + message = f"[batchId={batch_id}] {message}" + + # Truncate message over 10MB, remove chunk from the middle + if len(message) > MAX_MESSAGE_LENGTH: + half_max_length = MAX_MESSAGE_LENGTH // 2 + truncated_amount = len(message) - MAX_MESSAGE_LENGTH + truncation_note = f"\n...TRUNCATED {truncated_amount} CHARACTERS...\n" + message = ( + message[:half_max_length] + truncation_note + message[-half_max_length:] + ) + + if os.environ.get("RUNPOD_ENDPOINT_ID"): + log_json = {"requestId": job_id, "message": message, "level": message_level} + print(json.dumps(log_json), flush=True) + return + + if job_id: + message = f"{job_id} | {message}" + + print(f"{message_level.ljust(7)}| {message}", flush=True) + return + + def secret(self, secret_name, secret): + """ + Censors secrets for logging. + Replaces everything except the first and last characters with * + """ + secret = str(secret) + redacted_secret = secret[0] + "*" * (len(secret) - 2) + secret[-1] + self.info(f"{secret_name}: {redacted_secret}") + + def debug(self, message, request_id: Optional[str] = None): + """ + debug log + """ + self.log(message, "DEBUG", request_id) + + def info(self, message, request_id: Optional[str] = None): + """ + info log + """ + self.log(message, "INFO", request_id) + + def warn(self, message, request_id: Optional[str] = None): + """ + warn log + """ + self.log(message, "WARN", request_id) + + def error(self, message, request_id: Optional[str] = None): + """ + error log + """ + self.log(message, "ERROR", request_id) + + def tip(self, message): + """ + tip log + """ + self.log(message, "TIP") + + def trace(self, message, request_id: Optional[str] = None): + """ + trace log (buffered until flushed) + """ + self.log(message, "TRACE", request_id) diff --git a/runpod/_startup.py b/runpod/_startup.py new file mode 100644 index 000000000..303f8c771 --- /dev/null +++ b/runpod/_startup.py @@ -0,0 +1,37 @@ +"""Process-scoped startup gate; safe to import without loading serverless.""" + +import os +import sys + +# Launchers may set this to the PID of the Python handler process before exec. +# A generic container-level boolean would also authorize unrelated processes. +WORKER_PID_ENV = "RUNPOD_FITNESS_WORKER_PID" + + +def is_worker_process() -> bool: + """Require explicit launcher identity and exclude local/API test invocations.""" + return ( + os.environ.get(WORKER_PID_ENV) == str(os.getpid()) + and bool(os.environ.get("RUNPOD_WEBHOOK_GET_JOB")) + and not any( + arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") + for arg in sys.argv[1:] + ) + ) + + +def run_import_checks() -> None: + """Run early checks only in the handler process selected by the launcher.""" + if not is_worker_process(): + return + try: + from ._health.fitness import run_startup_fitness_checks + + run_startup_fitness_checks() + except Exception as exc: + # Import/configuration errors are retried through the worker-start path. + # Actual failed checks force-exit and do not pass through this handler. + print( + f"Runpod startup checks could not be prepared: {type(exc).__name__}: {exc}", + file=sys.stderr, + ) diff --git a/runpod/_worker_bootstrap.py b/runpod/_worker_bootstrap.py new file mode 100644 index 000000000..a71a83c1d --- /dev/null +++ b/runpod/_worker_bootstrap.py @@ -0,0 +1,47 @@ +"""Launch a worker with health checks before executing its handler module. + +Usage: runpod-worker handler.py [handler arguments] + runpod-worker -m package.handler [handler arguments] +""" + +import os +import runpy +import sys +from pathlib import Path + +from ._startup import WORKER_PID_ENV, run_import_checks + + +def main() -> None: + """Select this process as the worker, then execute the unmodified handler.""" + args = sys.argv[1:] + module_mode = bool(args and args[0] == "-m") + if module_mode: + args = args[1:] + if not args or args[0].startswith("-"): + raise SystemExit("Usage: runpod-worker [-m] handler [arguments]") + target, *handler_args = args + sys.argv = [target, *handler_args] + if not module_mode: + # Match `python handler.py`: sibling imports resolve beside the script. + target = str(Path(target).resolve()) + if not Path(target).is_file(): + raise SystemExit(f"Worker handler not found: {target}") + sys.path.insert(0, str(Path(target).parent)) + else: + # Console entrypoints put their bin directory on sys.path, unlike python -m. + sys.path.insert(0, os.getcwd()) + os.environ[WORKER_PID_ENV] = str(os.getpid()) + try: + run_import_checks() + finally: + # Helpers, subprocesses and multiprocessing children are not workers. + os.environ.pop(WORKER_PID_ENV, None) + if module_mode: + runpy.run_module(target, run_name="__main__", alter_sys=True) + else: + runpy.run_path(target, run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index c6fa605e1..f1ab29de5 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, run_startup_fitness_checks +from .modules.rp_fitness import register_fitness_check from .utils.rp_volume_cache import VolumeCache __all__ = [ @@ -29,9 +29,6 @@ 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() # ---------------------------------------------------------------------------- # diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 5451ae40e..74646d377 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -3,6 +3,7 @@ import os import threading import uuid +from contextlib import asynccontextmanager from dataclasses import dataclass from typing import Any, Dict, Optional, Union @@ -177,6 +178,25 @@ def _send_webhook(url: str, payload: Dict[str, Any]) -> bool: class WorkerAPI: """Used to launch the FastAPI web server when the worker is running in API mode.""" + @asynccontextmanager + async def _lifespan(self, app): + """Validate production realtime workers before accepting requests. + + Run in the serving process, after any server process creation, so CUDA + initialization cannot poison a later fork. Local API simulation skips it. + """ + from ..worker import _is_local + from .rp_fitness import run_fitness_checks + + args = self.config.get("rp_args", {}) + if ( + os.environ.get("RUNPOD_REALTIME_PORT") not in (None, "", "0") + and not args.get("rp_serve_api") + and not _is_local({"rp_args": args}) + ): + await run_fitness_checks() + yield + def __init__(self, config: Dict[str, Any]): """ Initializes the WorkerAPI class. @@ -217,6 +237,7 @@ def __init__(self, config: Dict[str, Any]): version=runpod_version, docs_url="/", openapi_tags=tags_metadata, + lifespan=self._lifespan, ) # Create an APIRouter and add the route for processing jobs. diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index 4066a4053..c9e8d5343 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -1,447 +1,6 @@ -""" -Fitness check system for worker startup validation. +"""Compatibility alias for :mod:`runpod._health.fitness`.""" -Fitness checks run before handler initialization on the actual RunPod serverless -platform to validate the worker environment. Any check failure force-kills the -worker via os._exit(1), signaling unhealthy state to the container orchestrator. - -Fitness checks do NOT run in local development mode or testing mode. -""" - -from __future__ import annotations - -import asyncio -import contextlib -import inspect -import os +import importlib import sys -import time -import traceback -from collections.abc import Callable - -from .rp_logger import RunPodLogger - -log = RunPodLogger() - - -def _terminate_unhealthy(code: int = 1) -> None: - """ - Force-kill the worker after a fitness check failure. - - Uses os._exit rather than sys.exit because a fitness failure means the - environment is broken and the worker must die immediately so the - orchestrator can restart it. sys.exit only raises SystemExit, which - triggers cooperative interpreter shutdown and blocks joining non-daemon - threads. Workers routinely have such threads alive by the time checks run - (e.g. vLLM's AsyncLLMEngine, constructed at import before the checks), so - sys.exit can hang forever and the worker keeps serving jobs. os._exit - bypasses thread joins, atexit handlers, and asyncgen cleanup. - - Args: - code: Process exit code (default 1, signaling unhealthy). - """ - # Best-effort flush of buffered logs before the hard exit skips normal - # cleanup. A broken worker may have a closed/None stdio stream; never let a - # flush failure stop the exit, which is the whole point of this helper. - for stream in (sys.stdout, sys.stderr): - with contextlib.suppress(Exception): - stream.flush() - os._exit(code) - -# 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" - -# Tuning vars consumed when the checks run. Snapshotted at the import-time -# pass so a later pass can warn about post-import changes, which would -# otherwise be silently ignored. -_CONFIG_ENV_VARS = ( - "RUNPOD_MIN_MEMORY_GB", - "RUNPOD_MIN_DISK_PERCENT", - "RUNPOD_MIN_CUDA_VERSION", - "RUNPOD_NETWORK_CHECK_TIMEOUT", - "RUNPOD_GPU_BENCHMARK_TIMEOUT", - "RUNPOD_GPU_TEST_TIMEOUT", - "RUNPOD_GPU_MAX_ERROR_MESSAGES", - "RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", - "RUNPOD_SKIP_GPU_CHECK", -) - -_config_snapshot: dict[str, str | None] = {} - - -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: - """ - Decorator to register a fitness check function. - - Fitness checks validate worker health at startup before handler initialization. - If any check fails, the worker is force-killed with os._exit(1). - - Supports both sync and async functions (auto-detected via inspect.iscoroutinefunction()). - - Example: - @runpod.serverless.register_fitness_check - def check_gpu(): - import torch - if not torch.cuda.is_available(): - raise RuntimeError("GPU not available") - - @runpod.serverless.register_fitness_check - async def check_model_files(): - import aiofiles.os - if not await aiofiles.os.path.exists("/models/model.safetensors"): - raise RuntimeError("Model file not found") - - Args: - func: Function to register as fitness check. Can be sync or async. - - Returns: - Original function unchanged (allows decorator stacking). - """ - _fitness_checks.append(func) - log.debug(f"Registered fitness check: {func.__name__}") - return func - - -def clear_fitness_checks() -> None: - """ - Clear all registered fitness checks. - - Used primarily for testing to reset global state between test cases. - Not intended for production use. - """ - _fitness_checks.clear() - _completed_checks.clear() - - -_registration_state: dict[str, bool] = { - "gpu_check": False, - "system_checks": False, -} - - -def _reset_registration_state() -> None: - """ - Reset global registration state. - - Used for testing to ensure clean state between tests. - """ - _registration_state["gpu_check"] = False - _registration_state["system_checks"] = False - - -# Bound how long the best-effort unhealthy report may delay the exit. -_REPORT_TIMEOUT_SECONDS = 2 - - -def _report_unhealthy(check: str, reason: str) -> None: - """ - Best-effort report of a fitness-check failure to the host before exit. - - Sends a single GET to the ping URL (same URL/credentials the heartbeat - uses) with status=unhealthy plus the failing check name and reason, so the - host can emit a queryable worker.fitness_failed event. Any failure — no - ping URL, no API key, HTTP error, timeout — is swallowed, so this can never - prevent the os._exit that follows. It is synchronous, so it may delay that - exit by up to _REPORT_TIMEOUT_SECONDS (network phases only; it adds no - delay when there is no ping URL/API key to report to). - """ - ping_url = os.environ.get("RUNPOD_WEBHOOK_PING") - api_key = os.environ.get("RUNPOD_AI_API_KEY") - if not ping_url or ping_url == "PING_NOT_SET" or not api_key: - return - - try: - # Deferred imports: keep module import light and avoid import cycles. - from runpod.http_client import SyncClientSession - from runpod.serverless.modules.worker_state import WORKER_ID - from runpod.version import __version__ as runpod_version - - ping_url = ping_url.replace("$RUNPOD_POD_ID", WORKER_ID) - params = { - "status": "unhealthy", - "check": check, - "reason": reason[:256], - "runpod_version": runpod_version, - } - session = SyncClientSession() - try: - session.headers.update({"Authorization": api_key}) - session.get(ping_url, params=params, timeout=_REPORT_TIMEOUT_SECONDS) - finally: - session.close() - except Exception: - # Best-effort only; the exit is the guarantee, not this report. - pass - - -def _ensure_gpu_check_registered() -> None: - """ - Ensure GPU fitness check is registered. - - Deferred until first run to avoid circular import issues during module - initialization. Called from run_fitness_checks() on first invocation. - """ - if _registration_state["gpu_check"]: - return - - # 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 - 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: - """ - Ensure system resource fitness checks are registered. - - Deferred until first run to avoid circular import issues during module - initialization. Called from run_fitness_checks() on first invocation. - """ - if _registration_state["system_checks"]: - return - - # Allow disabling system checks for testing - 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 - - # Same latch-on-success rule as _ensure_gpu_check_registered. - try: - from .rp_system_fitness import 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 - - -def _warn_late_config() -> None: - """Warn if tuning vars changed since the import-time pass consumed them.""" - changed = [ - name for name, old in _config_snapshot.items() if os.environ.get(name) != old - ] - if changed: - log.warn( - f"Fitness check config changed after the startup checks ran and is " - f"ignored: {', '.join(changed)}. Set these before `import runpod`, or " - f"set {DEFER_FITNESS_CHECKS_ENV}=true to run checks at start() as before." - ) - - -async def run_fitness_checks(include_deferred: bool = True) -> None: - """ - Execute all registered fitness checks sequentially at startup. - - Execution flow: - 1. Auto-register GPU check on first run (deferred to avoid circular imports) - 2. Check if registry is empty (early return if no checks) - 3. Log start of fitness check phase - 4. For each registered check: - - Auto-detect sync vs async using inspect.iscoroutinefunction() - - Execute check with timing instrumentation (await if async, call if sync) - - Log success or failure with check name and execution time - 5. On any exception: - - Log detailed error with check name, exception type, and message - - Log traceback at DEBUG level - - Force-kill the worker via os._exit(1) immediately (fail-fast). This is - a hard exit, not a cooperative sys.exit/SystemExit: it does not unwind - the stack or run cleanup, so callers cannot catch it and it cannot be - blocked by live non-daemon threads. - 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 - and handles checks with dependencies correctly. - Timing uses high-precision perf_counter for accurate measurements. - - Note: - 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 - - if _config_snapshot: - _warn_late_config() - - # 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() - - # 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(pending)} fitness check(s)...") - - total_start_time = time.perf_counter() - - for check_func in pending: - check_name = check_func.__name__ - - try: - log.debug(f"Executing fitness check: {check_name}") - check_start_time = time.perf_counter() - - # Auto-detect async vs sync using inspect - if inspect.iscoroutinefunction(check_func): - await check_func() - else: - 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: - # Log detailed error information - error_type = type(exc).__name__ - error_message = str(exc) - full_traceback = traceback.format_exc() - - log.error( - f"Fitness check failed: {check_name} | {error_type}: {error_message}" - ) - log.debug(f"Traceback:\n{full_traceback}") - - # Best-effort report to the host so the failure is queryable. It is - # bounded (see _REPORT_TIMEOUT_SECONDS) and fully swallowed, so it - # can delay the force-exit below but can never prevent it. - try: - _report_unhealthy(check_name, f"{error_type}: {error_message}") - except Exception: # a report failure must never prevent the exit - pass - - # Force-kill immediately; see _terminate_unhealthy for why this is - # os._exit rather than sys.exit. - log.error("Worker is unhealthy, exiting.") - _terminate_unhealthy(1) - - 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 - - # Remember the tuning values as consumed, so a later pass can warn about - # post-import changes (set in the handler, too late to apply). - _config_snapshot.update({v: os.environ.get(v) for v in _CONFIG_ENV_VARS}) - 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}") +sys.modules[__name__] = importlib.import_module("runpod._health.fitness") diff --git a/runpod/serverless/modules/rp_gpu_fitness.py b/runpod/serverless/modules/rp_gpu_fitness.py index 49db979fd..f436905f0 100644 --- a/runpod/serverless/modules/rp_gpu_fitness.py +++ b/runpod/serverless/modules/rp_gpu_fitness.py @@ -1,318 +1,6 @@ -""" -GPU fitness check system for worker startup validation. +"""Compatibility alias for :mod:`runpod._health.gpu`.""" -Provides comprehensive GPU health checking using: -1. Native CUDA binary (gpu_test) for memory allocation testing -2. Python fallback using nvidia-smi if binary unavailable +import importlib +import sys -Auto-registers when GPUs are detected, skips silently on CPU-only workers. -""" - -from __future__ import annotations - -import asyncio -import os -import subprocess -from pathlib import Path -from typing import Any - -from runpod._binary_helpers import get_binary_path -from .rp_fitness import _env_flag, register_fitness_check -from .rp_logger import RunPodLogger - -log = RunPodLogger() - -# Configuration via environment variables -TIMEOUT_SECONDS = int(os.environ.get("RUNPOD_GPU_TEST_TIMEOUT", "30")) -MAX_ERROR_MESSAGES = int(os.environ.get("RUNPOD_GPU_MAX_ERROR_MESSAGES", "10")) - - -def _get_gpu_test_binary_path() -> Path | None: - """ - Locate gpu_test binary in package. - - Returns: - Path to binary if found, None otherwise - """ - return get_binary_path("gpu_test") - - -def _parse_gpu_test_output(output: str) -> dict[str, Any]: - """ - Parse gpu_test binary output and detect success/failure. - - Looks for: - - "GPU X memory allocation test passed." for success - - Error patterns: "Failed", "error", "cannot" for failures - - GPU count from "Found X GPUs:" line - - Args: - output: Stdout from gpu_test binary - - Returns: - Dict with keys: - - success: bool - True if all GPUs passed tests - - gpu_count: int - Number of GPUs that passed tests - - found_gpus: int - Total GPUs found - - errors: List[str] - Error messages from output - - details: Dict - CUDA version, kernel version, etc - """ - lines = output.strip().split("\n") - - result = { - "success": False, - "gpu_count": 0, - "found_gpus": 0, - "errors": [], - "details": {}, - } - - passed_count = 0 - found_gpus = 0 - - for line in lines: - line = line.strip() - if not line: - continue - - # Extract metadata - if line.startswith("CUDA Driver Version:"): - result["details"]["cuda_version"] = line.split(":", 1)[1].strip() - elif line.startswith("Linux Kernel Version:"): - result["details"]["kernel"] = line.split(":", 1)[1].strip() - elif line.startswith("Found") and "GPUs" in line: - # "Found 2 GPUs:" - try: - found_gpus = int(line.split()[1]) - result["found_gpus"] = found_gpus - except (IndexError, ValueError): - # Line format doesn't match expected "Found N GPUs:" — skip - pass - - # Check for success - if "memory allocation test passed" in line.lower(): - passed_count += 1 - - # Check for errors - if any(err in line.lower() for err in ["failed", "error", "cannot", "unable"]): - result["errors"].append(line) - - result["gpu_count"] = passed_count - result["success"] = ( - passed_count > 0 and passed_count == found_gpus and len(result["errors"]) == 0 - ) - - return result - - -async def _run_gpu_test_binary() -> dict[str, Any]: - """ - Execute gpu_test binary and parse output. - - Returns: - Parsed result dict from _parse_gpu_test_output - - Raises: - RuntimeError: If binary execution fails or GPUs unhealthy - """ - binary_path = _get_gpu_test_binary_path() - - if not binary_path: - raise FileNotFoundError("gpu_test binary not found in package") - - if not os.access(binary_path, os.X_OK): - raise PermissionError(f"gpu_test binary not executable: {binary_path}") - - log.debug(f"Running gpu_test binary: {binary_path}") - - try: - # Run binary with timeout - process = await asyncio.create_subprocess_exec( - str(binary_path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - - stdout, stderr = await asyncio.wait_for( - process.communicate(), timeout=TIMEOUT_SECONDS - ) - - output = stdout.decode("utf-8", errors="replace") - error_output = stderr.decode("utf-8", errors="replace") - - log.debug(f"gpu_test output:\n{output}") - - if error_output: - log.debug(f"gpu_test stderr:\n{error_output}") - - # Parse output - result = _parse_gpu_test_output(output) - - # Check for success - if not result["success"]: - error_msg = "GPU memory allocation test failed" - if result["errors"]: - error_msg += f": {'; '.join(result['errors'][:MAX_ERROR_MESSAGES])}" - raise RuntimeError(error_msg) - - log.info( - f"GPU binary test passed: {result['gpu_count']} GPU(s) healthy " - f"(CUDA {result['details'].get('cuda_version', 'unknown')})" - ) - - return result - - except asyncio.TimeoutError: - process.kill() - await process.wait() - raise RuntimeError( - f"GPU test binary timed out after {TIMEOUT_SECONDS}s" - ) from None - except FileNotFoundError: - raise - except PermissionError: - raise - except Exception as exc: - raise RuntimeError(f"GPU test binary execution failed: {exc}") from exc - - -def _run_gpu_test_fallback() -> None: - """ - Python fallback for GPU testing using nvidia-smi. - - Less comprehensive than binary (doesn't test memory allocation) but validates - basic GPU availability by checking GPU count. - - Raises: - RuntimeError: If GPUs not available or unhealthy - """ - log.debug("Running Python GPU fallback check") - - try: - # List GPUs to verify availability and count - result = subprocess.run( - ["nvidia-smi", "--list-gpus"], - capture_output=True, - text=True, - timeout=10, - check=False, - ) - - if result.returncode != 0: - raise RuntimeError(f"nvidia-smi --list-gpus failed: {result.stderr}") - - gpu_lines = [line for line in result.stdout.split("\n") if line.strip()] - gpu_count = len(gpu_lines) - - if gpu_count == 0: - raise RuntimeError("No GPUs detected by nvidia-smi") - - log.info( - f"GPU fallback check passed: {gpu_count} GPU(s) detected " - "(Note: Memory allocation NOT tested)" - ) - - except FileNotFoundError: - raise RuntimeError( - "nvidia-smi not found. Cannot validate GPU availability." - ) from None - except subprocess.TimeoutExpired: - raise RuntimeError("nvidia-smi timed out") from None - except RuntimeError: - raise - except Exception as e: - raise RuntimeError(f"nvidia-smi fallback check failed: {e}") from e - - -async def _check_gpu_health() -> None: - """ - Comprehensive GPU health check (internal implementation). - - 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 - - Raises: - RuntimeError: If GPU health check fails - """ - binary_attempted = False - binary_error = None - - # Try binary first - try: - await _run_gpu_test_binary() - return # Success! - except FileNotFoundError as exc: - log.debug(f"GPU binary not found: {exc}") - binary_error = exc - except PermissionError as exc: - log.debug(f"GPU binary not executable: {exc}") - binary_error = exc - except Exception as exc: - log.warn(f"GPU binary check failed: {exc}") - binary_attempted = True - binary_error = exc - - # Fall back to Python - log.debug("Attempting Python GPU fallback check") - try: - _run_gpu_test_fallback() - return # Success! - except Exception as fallback_exc: - # Both failed - raise composite error - if binary_attempted: - raise RuntimeError( - f"GPU health check failed. " - f"Binary test: {binary_error}. " - f"Fallback test: {fallback_exc}" - ) from fallback_exc - else: - raise RuntimeError( - f"GPU health check failed (binary disabled/missing, " - f"fallback failed): {fallback_exc}" - ) from fallback_exc - - -def auto_register_gpu_check() -> None: - """ - Auto-register GPU fitness check if GPUs are detected. - - Called lazily on the first fitness-check run. - It detects GPU presence via nvidia-smi and registers the check if found. - On CPU-only workers, the check is skipped silently. - - Environment variables: - - RUNPOD_SKIP_GPU_CHECK: Set to a truthy value (1/true/yes/on) to skip auto-registration - """ - # Allow skipping during tests - if _env_flag("RUNPOD_SKIP_GPU_CHECK"): - log.debug("GPU fitness check auto-registration disabled via environment") - return - - # Quick GPU detection - has_gpu = False - try: - result = subprocess.run( - ["nvidia-smi"], - capture_output=True, - text=True, - timeout=5, - check=False, - ) - has_gpu = result.returncode == 0 and "NVIDIA-SMI" in result.stdout - except (FileNotFoundError, subprocess.TimeoutExpired): - has_gpu = False - except Exception: - # Catch any other exceptions and assume no GPU - has_gpu = False - - if has_gpu: - log.debug("GPU detected, registering automatic GPU fitness check") - - @register_fitness_check - async def _gpu_health_check(): - """Automatic GPU memory allocation health check.""" - await _check_gpu_health() - else: - log.debug("No GPU detected, skipping GPU fitness check registration") +sys.modules[__name__] = importlib.import_module("runpod._health.gpu") diff --git a/runpod/serverless/modules/rp_logger.py b/runpod/serverless/modules/rp_logger.py index 6ef4c5f73..3be00449c 100644 --- a/runpod/serverless/modules/rp_logger.py +++ b/runpod/serverless/modules/rp_logger.py @@ -1,155 +1,6 @@ -""" -PodWorker | modules | logging.py +"""Compatibility alias for :mod:`runpod._logger`.""" -Log Levels (Level - Value - Description) +import importlib +import sys -NOTSET - 0 - No logging is configured, the logging system is effectively disabled. -DEBUG - 1 - Detailed information, typically of interest only when diagnosing problems. (Default) -INFO - 2 - Confirmation that things are working as expected. -WARN - 3 - An indication that something unexpected happened. -ERROR - 4 - Serious problem, the software has not been able to perform some function. -""" - -from contextvars import ContextVar, Token -import json -import os -from typing import Optional - -MAX_MESSAGE_LENGTH = 4096 -LOG_LEVELS = ["NOTSET", "TRACE", "DEBUG", "INFO", "WARN", "ERROR"] -_batch_id: ContextVar[Optional[str]] = ContextVar("runpod_batch_id", default=None) - - -def _set_batch_id(batch_id: Optional[str]) -> Token: - """Set the batch ID associated with the current job task.""" - return _batch_id.set(batch_id) - - -def _reset_batch_id(token: Token): - """Restore the previous batch ID for the current job task.""" - _batch_id.reset(token) - - -def _validate_log_level(log_level): - """ - Checks the debug level and returns the debug level name. - """ - if isinstance(log_level, str): - log_level = log_level.upper() - - if log_level not in LOG_LEVELS: - raise ValueError(f"Invalid debug level: {log_level}") - - return log_level - - if isinstance(log_level, int): - if log_level < 0 or log_level >= len(LOG_LEVELS): - raise ValueError(f"Invalid debug level: {log_level}") - - return LOG_LEVELS[log_level] - - raise ValueError(f"Invalid debug level: {log_level}") - - -class RunPodLogger: - """Singleton class for logging.""" - - __instance = None - level = _validate_log_level( - os.environ.get( - "RUNPOD_LOG_LEVEL", os.environ.get("RUNPOD_DEBUG_LEVEL", "DEBUG") - ) - ) - - def __new__(cls): - if RunPodLogger.__instance is None: - RunPodLogger.__instance = object.__new__(cls) - return RunPodLogger.__instance - - def set_level(self, new_level): - """ - Set the debug level for logging. - Can be set to the name or value of the debug level. - """ - self.level = _validate_log_level(new_level) - self.info(f"Log level set to {self.level}") - - def log(self, message, message_level="INFO", job_id=None): - """ - Log message to stdout if RUNPOD_DEBUG is true. - """ - if self.level == "NOTSET": - return - - level_index = LOG_LEVELS.index(self.level) - if level_index > LOG_LEVELS.index(message_level) and message_level != "TIP": - return - - message = str(message) - if batch_id := _batch_id.get(): - message = f"[batchId={batch_id}] {message}" - - # Truncate message over 10MB, remove chunk from the middle - if len(message) > MAX_MESSAGE_LENGTH: - half_max_length = MAX_MESSAGE_LENGTH // 2 - truncated_amount = len(message) - MAX_MESSAGE_LENGTH - truncation_note = f"\n...TRUNCATED {truncated_amount} CHARACTERS...\n" - message = ( - message[:half_max_length] + truncation_note + message[-half_max_length:] - ) - - if os.environ.get("RUNPOD_ENDPOINT_ID"): - log_json = {"requestId": job_id, "message": message, "level": message_level} - print(json.dumps(log_json), flush=True) - return - - if job_id: - message = f"{job_id} | {message}" - - print(f"{message_level.ljust(7)}| {message}", flush=True) - return - - def secret(self, secret_name, secret): - """ - Censors secrets for logging. - Replaces everything except the first and last characters with * - """ - secret = str(secret) - redacted_secret = secret[0] + "*" * (len(secret) - 2) + secret[-1] - self.info(f"{secret_name}: {redacted_secret}") - - def debug(self, message, request_id: Optional[str] = None): - """ - debug log - """ - self.log(message, "DEBUG", request_id) - - def info(self, message, request_id: Optional[str] = None): - """ - info log - """ - self.log(message, "INFO", request_id) - - def warn(self, message, request_id: Optional[str] = None): - """ - warn log - """ - self.log(message, "WARN", request_id) - - def error(self, message, request_id: Optional[str] = None): - """ - error log - """ - self.log(message, "ERROR", request_id) - - def tip(self, message): - """ - tip log - """ - self.log(message, "TIP") - - def trace(self, message, request_id: Optional[str] = None): - """ - trace log (buffered until flushed) - """ - self.log(message, "TRACE", request_id) +sys.modules[__name__] = importlib.import_module("runpod._logger") diff --git a/runpod/serverless/modules/rp_system_fitness.py b/runpod/serverless/modules/rp_system_fitness.py index 1ac80f5b3..f531d43fc 100644 --- a/runpod/serverless/modules/rp_system_fitness.py +++ b/runpod/serverless/modules/rp_system_fitness.py @@ -1,517 +1,6 @@ -""" -System resource fitness checks for worker startup validation. +"""Compatibility alias for :mod:`runpod._health.system`.""" -Provides comprehensive checks for: -- Memory availability -- Disk space -- Network connectivity -- CUDA library versions -- GPU compute benchmark +import importlib +import sys -Auto-registers when worker starts, ensuring system readiness before accepting jobs. -""" - -from __future__ import annotations - -import asyncio -import os -import re -import shutil -import time - -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 - -log = RunPodLogger() - -# Configuration via environment variables -MIN_MEMORY_GB = float(os.environ.get("RUNPOD_MIN_MEMORY_GB", "4.0")) -MIN_DISK_PERCENT = float(os.environ.get("RUNPOD_MIN_DISK_PERCENT", "10.0")) -MIN_CUDA_VERSION = os.environ.get("RUNPOD_MIN_CUDA_VERSION", "11.8") -NETWORK_CHECK_TIMEOUT = int(os.environ.get("RUNPOD_NETWORK_CHECK_TIMEOUT", "5")) -GPU_BENCHMARK_TIMEOUT = int(os.environ.get("RUNPOD_GPU_BENCHMARK_TIMEOUT", "2")) - - -def _parse_version(version_string: str) -> tuple[int, int]: - """ - Parse version string to tuple for comparison. - - Args: - version_string: Version string like "12.2" or "CUDA Version 12.2" - - Returns: - Tuple of ints like (12, 2) for comparison - """ - # Extract numeric version - match = re.search(r"(\d+)\.(\d+)", version_string) - if match: - return (int(match.group(1)), int(match.group(2))) - return (0, 0) - - -def _get_memory_info() -> dict[str, float]: - """ - Get system memory information. - - Returns: - Dict with total_gb, available_gb, used_percent - - Raises: - RuntimeError: If memory check fails - """ - try: - import psutil - - mem = psutil.virtual_memory() - total_gb = mem.total / (1024**3) - available_gb = mem.available / (1024**3) - used_percent = mem.percent - - return { - "total_gb": total_gb, - "available_gb": available_gb, - "used_percent": used_percent, - } - except ImportError: - # Fallback: parse /proc/meminfo - try: - with open("/proc/meminfo") as f: - meminfo_kb: dict[str, int] = {} - for line in f: - key, value = line.split(":", 1) - meminfo_kb[key.strip()] = int(value.split()[0]) - - # /proc/meminfo values are in kB; convert to GB - total_gb = meminfo_kb.get("MemTotal", 0) / (1024**2) - available_gb = meminfo_kb.get("MemAvailable", 0) / (1024**2) - used_percent = ( - 100 * (1 - available_gb / total_gb) if total_gb > 0 else 0 - ) - - return { - "total_gb": total_gb, - "available_gb": available_gb, - "used_percent": used_percent, - } - except Exception as e: - raise RuntimeError(f"Failed to read memory info: {e}") from e - - -def _check_memory_availability() -> None: - """ - Check system memory availability. - - Raises: - RuntimeError: If insufficient memory available - """ - mem_info = _get_memory_info() - available_gb = mem_info["available_gb"] - total_gb = mem_info["total_gb"] - - if available_gb < MIN_MEMORY_GB: - raise RuntimeError( - f"Insufficient memory: {available_gb:.2f}GB available, " - f"{MIN_MEMORY_GB}GB required" - ) - - log.info( - f"Memory check passed: {available_gb:.2f}GB available " - f"(of {total_gb:.2f}GB total)" - ) - - -def _check_disk_space() -> None: - """ - Check disk space availability on root filesystem. - - In containers, root (/) is typically the only filesystem. - Requires free space to be at least MIN_DISK_PERCENT% of total disk size. - - Raises: - RuntimeError: If insufficient disk space - """ - try: - usage = shutil.disk_usage("/") - total_gb = usage.total / (1024**3) - free_gb = usage.free / (1024**3) - free_percent = 100 * (free_gb / total_gb) if total_gb > 0 else 0 - - # Check if free space is below the required percentage - if free_percent < MIN_DISK_PERCENT: - raise RuntimeError( - f"Insufficient disk space: {free_gb:.2f}GB free " - f"({free_percent:.1f}%), {MIN_DISK_PERCENT}% required" - ) - - log.info( - f"Disk space check passed: {free_gb:.2f}GB free " - f"({free_percent:.1f}% available)" - ) - except FileNotFoundError: - raise RuntimeError( - "Could not check disk space: / filesystem not found" - ) from None - - -async def _check_network_connectivity() -> None: - """ - Check basic network connectivity to 8.8.8.8:53. - - Raises: - RuntimeError: If network connectivity fails - """ - host = "8.8.8.8" - port = 53 - - try: - start_time = time.perf_counter() - _, writer = await asyncio.wait_for( - asyncio.open_connection(host, port), timeout=NETWORK_CHECK_TIMEOUT - ) - elapsed_ms = (time.perf_counter() - start_time) * 1000 - writer.close() - await writer.wait_closed() - - log.info( - f"Network connectivity passed: Connected to {host} ({elapsed_ms:.0f}ms)" - ) - except asyncio.TimeoutError: - raise RuntimeError( - f"Network connectivity failed: Timeout connecting to {host}:{port} " - f"({NETWORK_CHECK_TIMEOUT}s)" - ) from None - except ConnectionRefusedError: - raise RuntimeError( - f"Network connectivity failed: Connection refused to {host}:{port}" - ) from None - except Exception as e: - raise RuntimeError(f"Network connectivity check failed: {e}") from e - - -async def _get_cuda_version() -> str | None: - """ - Get CUDA version from system. - - Returns: - Version string like "12.2" or None if not available - - Raises: - RuntimeError: If CUDA check fails critically - """ - # Try nvcc first - process = None - try: - process = await asyncio.create_subprocess_exec( - "nvcc", - "--version", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, _ = await asyncio.wait_for(process.communicate(), timeout=5) - if process.returncode == 0: - output = stdout.decode("utf-8", errors="replace") - for line in output.split("\n"): - if "release" in line.lower() or "version" in line.lower(): - return line.strip() - except Exception as e: - if process and process.returncode is None: - process.kill() - await process.wait() - log.debug(f"nvcc not available: {e}") - - # Fallback: try nvidia-smi and parse CUDA version from output - process = None - try: - process = await asyncio.create_subprocess_exec( - "nvidia-smi", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, _ = await asyncio.wait_for(process.communicate(), timeout=5) - if process.returncode == 0: - output = stdout.decode("utf-8", errors="replace") - for line in output.split("\n"): - if "CUDA Version:" in line: - parts = line.split("CUDA Version:") - if len(parts) > 1: - cuda_version = parts[1].strip().split()[0] - return f"CUDA Version: {cuda_version}" - log.debug("nvidia-smi output found but couldn't parse CUDA version") - except Exception as e: - if process and process.returncode is None: - process.kill() - await process.wait() - log.debug(f"nvidia-smi not available: {e}") - - return None - - -async def _check_cuda_versions() -> None: - """ - Check CUDA library versions meet minimum requirements. - - Raises: - RuntimeError: If CUDA version is below minimum - """ - cuda_version_str = await _get_cuda_version() - - if not cuda_version_str: - log.warn("Could not determine CUDA version, skipping check") - return - - # Parse version - cuda_version = _parse_version(cuda_version_str) - min_version = _parse_version(MIN_CUDA_VERSION) - - if cuda_version < min_version: - raise RuntimeError( - f"CUDA version too old: {cuda_version[0]}.{cuda_version[1]} found, " - f"{min_version[0]}.{min_version[1]} required" - ) - - log.info( - f"CUDA version check passed: {cuda_version[0]}.{cuda_version[1]} " - f"(minimum: {min_version[0]}.{min_version[1]})" - ) - - -async def _check_cuda_initialization() -> None: - """ - Verify CUDA can be initialized and devices are accessible. - - 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. - - Raises: - RuntimeError: If CUDA initialization or device access fails - """ - # 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 - - if not torch.cuda.is_available(): - log.debug("CUDA not available in PyTorch, skipping initialization check") - return - - # Reset CUDA state to ensure clean initialization - torch.cuda.reset_peak_memory_stats() - torch.cuda.synchronize() - - # Verify device count - device_count = torch.cuda.device_count() - if device_count == 0: - raise RuntimeError( - "No CUDA devices available despite cuda.is_available() being True" - ) - - # Test each device - for i in range(device_count): - try: - # Get device properties - props = torch.cuda.get_device_properties(i) - if props.total_memory == 0: - raise RuntimeError(f"GPU {i} reports zero memory") - - # Try allocating a small tensor on the device - _ = torch.zeros(1024, device=f"cuda:{i}") - torch.cuda.synchronize() - - except Exception as e: - raise RuntimeError(f"Failed to initialize GPU {i}: {e}") from e - - log.info( - f"CUDA initialization passed: {device_count} device(s) initialized successfully" - ) - return - - except ImportError: - log.debug("PyTorch not available, trying CuPy...") - except Exception as e: - raise RuntimeError(f"CUDA initialization failed: {e}") from e - - # Fallback: try CuPy - try: - import cupy as cp - - # Reset CuPy state - cp.cuda.Device().synchronize() - - # Verify devices - device_count = cp.cuda.runtime.getDeviceCount() - if device_count == 0: - raise RuntimeError("No CUDA devices available via CuPy") - - # Test each device - for i in range(device_count): - try: - cp.cuda.Device(i).use() - # Try allocating memory - _ = cp.zeros(1024) - cp.cuda.Device().synchronize() - except Exception as e: - raise RuntimeError( - f"Failed to initialize GPU {i} with CuPy: {e}" - ) from e - - log.info( - f"CUDA initialization passed: {device_count} device(s) initialized successfully" - ) - return - - except ImportError: - log.debug("CuPy not available, skipping CUDA initialization check") - except Exception as e: - raise RuntimeError(f"CUDA initialization check failed: {e}") from e - - -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. - - Raises: - RuntimeError: If GPU compute fails or is too slow - """ - # 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 - - if not torch.cuda.is_available(): - log.debug("CUDA not available in PyTorch, skipping benchmark") - return - - # Create small matrix on GPU - size = 1024 - start_time = time.perf_counter() - - # Do computation - A = torch.randn(size, size, device="cuda") - B = torch.randn(size, size, device="cuda") - torch.matmul(A, B) - torch.cuda.synchronize() # Wait for GPU to finish - - elapsed_ms = (time.perf_counter() - start_time) * 1000 - max_ms = GPU_BENCHMARK_TIMEOUT * 1000 - - if elapsed_ms > max_ms: - raise RuntimeError( - f"GPU compute too slow: Matrix multiply took {elapsed_ms:.0f}ms " - f"(max: {max_ms:.0f}ms)" - ) - - log.info( - f"GPU compute benchmark passed: Matrix multiply completed in {elapsed_ms:.0f}ms" - ) - return - - except ImportError: - log.debug("PyTorch not available, trying CuPy...") - except RuntimeError: - raise # Benchmark failure is what we're testing for - except Exception as e: - log.warn(f"PyTorch GPU benchmark setup failed: {e}") - - # Fallback: try CuPy - try: - import cupy as cp - - size = 1024 - start_time = time.perf_counter() - - A = cp.random.randn(size, size) - B = cp.random.randn(size, size) - cp.matmul(A, B) - cp.cuda.Device().synchronize() - - elapsed_ms = (time.perf_counter() - start_time) * 1000 - max_ms = GPU_BENCHMARK_TIMEOUT * 1000 - - if elapsed_ms > max_ms: - raise RuntimeError( - f"GPU compute too slow: Matrix multiply took {elapsed_ms:.0f}ms " - f"(max: {max_ms:.0f}ms)" - ) - - log.info( - f"GPU compute benchmark passed: Matrix multiply completed in {elapsed_ms:.0f}ms" - ) - return - - except ImportError: - log.debug("CuPy not available, skipping GPU benchmark") - except RuntimeError: - raise # Benchmark failure is what we're testing for - except Exception as e: - log.warn(f"CuPy GPU benchmark setup failed: {e}") - - # If we get here, neither library is available - log.debug( - "PyTorch/CuPy not available for GPU benchmark, relying on gpu_test binary" - ) - - -def auto_register_system_checks() -> None: - """ - Auto-register system resource fitness checks. - - 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") - - # Always register these checks - @register_fitness_check - def _memory_check() -> None: - """System memory availability check.""" - _check_memory_availability() - - @register_fitness_check - def _disk_check() -> None: - """System disk space check.""" - _check_disk_space() - - @register_fitness_check - async def _network_check() -> None: - """Network connectivity check.""" - await _check_network_connectivity() - - # Only register GPU checks if GPU is detected - if gpu_available(): - log.debug("GPU detected, registering GPU-specific fitness checks") - - @register_fitness_check - async def _cuda_version_check() -> None: - """CUDA version check.""" - 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() - else: - log.debug("No GPU detected, skipping GPU-specific fitness checks") +sys.modules[__name__] = importlib.import_module("runpod._health.system") diff --git a/runpod/serverless/utils/rp_cuda.py b/runpod/serverless/utils/rp_cuda.py index 1a47108a4..f561f993f 100644 --- a/runpod/serverless/utils/rp_cuda.py +++ b/runpod/serverless/utils/rp_cuda.py @@ -1,22 +1,6 @@ -""" -Provides some of the torch.cuda functionality without requiring torch. -""" +"""Compatibility alias for :mod:`runpod._health.cuda`.""" -import subprocess +import importlib +import sys - -def is_available(): - """ - Returns True if CUDA is available, False otherwise. - """ - try: - # 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 - pass - return False +sys.modules[__name__] = importlib.import_module("runpod._health.cuda") diff --git a/tests/test_serverless/test_modules/test_fitness/test_force_kill.py b/tests/test_serverless/test_modules/test_fitness/test_force_kill.py index 9bb4f2130..f49c1a705 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_force_kill.py +++ b/tests/test_serverless/test_modules/test_fitness/test_force_kill.py @@ -103,9 +103,9 @@ def test_report_unhealthy_posts_check_and_reason(monkeypatch): monkeypatch.setenv("RUNPOD_WEBHOOK_PING", "https://api.test/ping/$RUNPOD_POD_ID") monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") + monkeypatch.setenv("RUNPOD_POD_ID", "podABC") fake_session = MagicMock() - with patch("runpod.http_client.SyncClientSession", return_value=fake_session), \ - patch("runpod.serverless.modules.worker_state.WORKER_ID", "podABC"): + with patch("requests.Session", return_value=fake_session): rp_fitness._report_unhealthy("_cuda_init_check", "RuntimeError: boom") assert fake_session.get.call_count == 1 @@ -122,7 +122,7 @@ def test_report_unhealthy_posts_check_and_reason(monkeypatch): def test_report_unhealthy_skipped_without_ping_url(monkeypatch): monkeypatch.delenv("RUNPOD_WEBHOOK_PING", raising=False) monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") - with patch("runpod.http_client.SyncClientSession") as session_cls: + with patch("requests.Session") as session_cls: rp_fitness._report_unhealthy("_memory_check", "RuntimeError: low") session_cls.assert_not_called() @@ -132,7 +132,7 @@ def test_report_unhealthy_truncates_long_reason(monkeypatch): monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") fake_session = MagicMock() - with patch("runpod.http_client.SyncClientSession", return_value=fake_session): + with patch("requests.Session", return_value=fake_session): rp_fitness._report_unhealthy("_disk_check", "x" * 300) params = fake_session.get.call_args.kwargs["params"] @@ -142,7 +142,7 @@ def test_report_unhealthy_truncates_long_reason(monkeypatch): def test_report_unhealthy_skipped_without_api_key(monkeypatch): monkeypatch.setenv("RUNPOD_WEBHOOK_PING", "https://api.test/ping") monkeypatch.delenv("RUNPOD_AI_API_KEY", raising=False) - with patch("runpod.http_client.SyncClientSession") as session_cls: + with patch("requests.Session") as session_cls: rp_fitness._report_unhealthy("_memory_check", "RuntimeError: low") session_cls.assert_not_called() @@ -152,7 +152,7 @@ def test_report_unhealthy_swallows_errors(monkeypatch): monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") fake_session = MagicMock() fake_session.get.side_effect = RuntimeError("network down") - with patch("runpod.http_client.SyncClientSession", return_value=fake_session): + with patch("requests.Session", return_value=fake_session): rp_fitness._report_unhealthy("_disk_check", "RuntimeError: full") # must not raise diff --git a/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py b/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py new file mode 100644 index 000000000..c2497ec3f --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py @@ -0,0 +1,333 @@ +"""Customer-safety regressions for automatic early worker checks.""" + +import asyncio +import os +import subprocess +import sys +import time +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from runpod._health import fitness, system +from runpod._startup import WORKER_PID_ENV, run_import_checks + + +@pytest.mark.parametrize( + "args", + [ + ["handler.py"], + ["handler.py", "--test_input", "{}"], + ["handler.py", "--test_input={}"], + ["handler.py", "--rp_serve_api"], + ], +) +def test_import_does_not_run_for_unmarked_or_local_process(monkeypatch, args): + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + monkeypatch.setattr(sys, "argv", args) + if len(args) > 1: + monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + else: + monkeypatch.delenv(WORKER_PID_ENV, raising=False) + with patch.object(fitness, "run_startup_fitness_checks") as run: + run_import_checks() + run.assert_not_called() + + +def test_inherited_worker_pid_does_not_authorize_child(monkeypatch): + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid() + 1)) + with patch.object(fitness, "run_startup_fitness_checks") as run: + run_import_checks() + run.assert_not_called() + + +def test_initial_pass_does_not_compare_config(monkeypatch): + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + with patch.object(fitness, "_refresh_late_config") as refresh: + run_import_checks() + refresh.assert_not_called() + assert WORKER_PID_ENV not in os.environ + + +@pytest.mark.asyncio +async def test_registration_failure_rolls_back_and_reports_at_worker_start(monkeypatch): + def partial_registration(): + fitness.register_fitness_check(lambda: None) + raise ValueError("bad threshold") + + monkeypatch.setattr( + fitness, "_ensure_system_checks_registered", partial_registration + ) + await fitness.run_fitness_checks(include_deferred=False) + assert fitness._fitness_checks == [] + assert fitness._registration_state == {"gpu_check": False, "system_checks": False} + with patch.object(fitness, "_report_unhealthy") as report: + with pytest.raises(SystemExit): + await fitness.run_fitness_checks() + assert report.call_args.args == ("fitness_check_setup", "ValueError: bad threshold") + + +@pytest.mark.asyncio +async def test_setup_failure_exits_even_if_reporting_breaks(monkeypatch): + monkeypatch.setattr( + fitness, "_register_builtins", MagicMock(side_effect=ValueError("bad")) + ) + monkeypatch.setattr( + fitness, "_report_unhealthy", MagicMock(side_effect=RuntimeError("offline")) + ) + with pytest.raises(SystemExit) as exc: + await fitness.run_fitness_checks() + assert exc.value.code == 1 + + +@pytest.mark.asyncio +async def test_network_retries_then_succeeds_on_worker_api_host(monkeypatch): + monkeypatch.setenv( + "RUNPOD_WEBHOOK_GET_JOB", "https://worker.example:8443/job?token=secret" + ) + writer = MagicMock() + writer.wait_closed = AsyncMock() + with patch("asyncio.open_connection", new_callable=AsyncMock) as connect: + connect.side_effect = [ConnectionRefusedError(), (MagicMock(), writer)] + await system._check_network_connectivity() + assert connect.await_count == 2 + connect.assert_awaited_with("worker.example", 8443) + writer.close.assert_called_once() + + +@pytest.mark.asyncio +async def test_network_stuck_close_is_bounded(monkeypatch): + monkeypatch.setattr(system, "NETWORK_CHECK_TIMEOUT", 0.1) + writer = MagicMock() + writer.wait_closed.side_effect = lambda: asyncio.sleep(60) + with patch( + "asyncio.open_connection", + new_callable=AsyncMock, + return_value=(MagicMock(), writer), + ): + started = time.monotonic() + with pytest.raises(RuntimeError, match="Timeout"): + await system._check_network_connectivity() + assert time.monotonic() - started < 1 + writer.transport.abort.assert_called() + + +def test_network_is_deferred_even_in_authorized_worker(monkeypatch): + monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://worker.example/job") + monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + with ( + patch.object(system, "gpu_available", return_value=False), + patch.object(system, "_check_memory_availability"), + patch.object(system, "_check_disk_space"), + patch.object( + system, "_check_network_connectivity", new_callable=AsyncMock + ) as network, + ): + run_import_checks() + network.assert_not_awaited() + assert any(c.__name__ == "_network_check" for c in fitness._fitness_checks) + + +def test_changed_threshold_is_applied_without_rerunning_unrelated_checks(monkeypatch): + monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://worker.example/job") + monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + with ( + patch.object(system, "gpu_available", return_value=False), + patch.object(system, "_check_memory_availability") as memory, + patch.object(system, "_check_disk_space") as disk, + patch.object(system, "_check_network_connectivity", new_callable=AsyncMock), + ): + run_import_checks() + monkeypatch.setenv("RUNPOD_MIN_DISK_PERCENT", "2") + asyncio.run(fitness.run_fitness_checks()) + assert system.MIN_DISK_PERCENT == 2 + assert memory.call_count == 1 + assert disk.call_count == 2 + + +@pytest.mark.parametrize("local", [False, True]) +@pytest.mark.asyncio +async def test_realtime_checks_before_serving_but_local_api_exempt(monkeypatch, local): + from runpod.serverless.modules.rp_fastapi import WorkerAPI + + monkeypatch.setenv("RUNPOD_REALTIME_PORT", "8000") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://worker.example/job") + api = object.__new__(WorkerAPI) + api.config = {"rp_args": {"rp_serve_api": local}} + with patch.object(fitness, "run_fitness_checks", new_callable=AsyncMock) as run: + async with api._lifespan(None): + assert run.await_count == (0 if local else 1) + + +def run_child(code, **kwargs): + env = {k: v for k, v in os.environ.items() if not k.startswith("RUNPOD_")} + env.update(RUNPOD_SKIP_GPU_CHECK="true", RUNPOD_SKIP_AUTO_SYSTEM_CHECKS="true") + return subprocess.run( + [sys.executable, "-c", code], + env=env, + text=True, + capture_output=True, + timeout=15, + **kwargs, + ) + + +def test_actual_import_is_safe_with_inherited_worker_environment(): + result = run_child(""" +import os +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +os.environ['RUNPOD_SKIP_AUTO_SYSTEM_CHECKS'] = 'false' +os.environ['RUNPOD_MIN_MEMORY_GB'] = 'invalid' +import runpod +print('IMPORT_SURVIVED') +""") + assert result.returncode == 0, result.stderr + assert "IMPORT_SURVIVED" in result.stdout + + +@pytest.mark.parametrize("module_mode", [False, True]) +def test_launcher_checks_before_handler_and_preserves_arguments(tmp_path, module_mode): + handler = tmp_path / "handler.py" + handler.write_text( + "import os, sys\nassert 'RUNPOD_FITNESS_WORKER_PID' not in os.environ\n" + "assert sys.argv[1:] == ['--customer-arg', 'value']\nprint('MODEL_LOAD')\n" + ) + result = run_child(f""" +import os, sys +import runpod._worker_bootstrap as bootstrap +from runpod._health import fitness +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +fitness.register_fitness_check(lambda: print('EARLY_CHECK')) +os.chdir({str(tmp_path)!r}) +# Model console entrypoint sys.path: the current directory is not pre-added. +sys.path = [p for p in sys.path if p] +sys.argv = ['runpod-worker'] + {(["-m", "handler"] if module_mode else [str(handler)])!r} + ['--customer-arg', 'value'] +bootstrap.main() +""") + assert result.returncode == 0, result.stderr + assert result.stdout.index("EARLY_CHECK") < result.stdout.index("MODEL_LOAD") + + +def test_setup_failure_exits_with_live_thread(): + result = run_child(""" +import asyncio, os, threading, time +from runpod._health import fitness +os.environ['RUNPOD_SKIP_AUTO_SYSTEM_CHECKS'] = 'false' +os.environ['RUNPOD_MIN_MEMORY_GB'] = 'invalid' +threading.Thread(target=lambda: time.sleep(60), daemon=False).start() +asyncio.run(fitness.run_fitness_checks()) +""") + assert result.returncode == 1, result.stderr + assert "fitness_check_setup" in result.stdout + + +def test_lazy_parent_early_checks_never_import_serverless_or_cuda_libraries(): + root = str(Path(__file__).resolve().parents[5] / "runpod") + result = run_child(f""" +import asyncio, importlib.abc, os, sys, types +# Model the apps-sdk lazy package: no eager serverless import. +package = types.ModuleType('runpod') +package.__path__ = [{root!r}] +sys.modules['runpod'] = package +class Guard(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith('runpod.serverless') or fullname.split('.')[0] in ('torch', 'cupy'): + raise AssertionError('early check loaded ' + fullname) +sys.meta_path.insert(0, Guard()) +from unittest.mock import patch, MagicMock +from runpod._health import fitness, system +from runpod._startup import run_import_checks +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +os.environ['RUNPOD_FITNESS_WORKER_PID'] = str(os.getpid()) +os.environ['RUNPOD_SKIP_AUTO_SYSTEM_CHECKS'] = 'false' +loop = asyncio.new_event_loop() +asyncio.set_event_loop(loop) +with patch.object(system, 'gpu_available', return_value=False), patch.object(system, '_check_memory_availability'), patch.object(system, '_check_disk_space'): + run_import_checks() +assert asyncio.get_event_loop() is loop +loop.close() +assert sorted(c.__name__ for c in fitness._completed_checks) == ['_disk_check', '_memory_check'] +os.environ['RUNPOD_WEBHOOK_PING'] = 'https://example.test/ping' +os.environ['RUNPOD_AI_API_KEY'] = 'fake-test-key' +with patch('requests.Session') as session: + fitness._report_unhealthy('test', 'failure') + session.return_value.get.assert_called_once() +print('LAZY_PASS') +""") + assert result.returncode == 0, result.stderr + assert "LAZY_PASS" in result.stdout + + +@pytest.mark.parametrize("method", ["spawn", "fork"]) +def test_child_processes_do_not_repeat_early_checks(tmp_path, method): + import multiprocessing + + if method not in multiprocessing.get_all_start_methods(): + pytest.skip(f"{method} is not supported") + handler = tmp_path / "child_handler.py" + handler.write_text(""" +import multiprocessing, os +from runpod._startup import is_worker_process + +def child(): + assert not is_worker_process() + print('CHILD_SAFE', flush=True) + +if __name__ == '__main__': + child_process = multiprocessing.get_context(os.environ['TEST_START_METHOD']).Process(target=child) + child_process.start() + child_process.join(5) + assert child_process.exitcode == 0 +""") + result = run_child(f""" +import os, sys +from runpod._worker_bootstrap import main +os.environ['TEST_START_METHOD'] = {method!r} +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +sys.argv = ['runpod-worker', {str(handler)!r}] +main() +""") + assert result.returncode == 0, result.stderr + assert "CHILD_SAFE" in result.stdout + + +def test_launcher_failed_early_check_prevents_model_load(tmp_path): + handler = tmp_path / "handler.py" + handler.write_text("print('MODEL_LOAD')\n") + result = run_child(f""" +import os, sys +from runpod._worker_bootstrap import main +from runpod._health import fitness +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +def fail(): + raise RuntimeError('broken hardware') +fitness.register_fitness_check(fail) +sys.argv = ['runpod-worker', {str(handler)!r}] +main() +""") + assert result.returncode == 1 + assert "broken hardware" in result.stdout + assert "MODEL_LOAD" not in result.stdout + + +def test_launcher_local_test_does_not_run_early_checks(tmp_path): + handler = tmp_path / "handler.py" + handler.write_text("print('LOCAL_TEST')\n") + result = run_child(f""" +import os, sys +from runpod._worker_bootstrap import main +from runpod._health import fitness +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +def fail(): + raise RuntimeError('must not run') +fitness.register_fitness_check(fail) +sys.argv = ['runpod-worker', {str(handler)!r}, '--test_input={{}}'] +main() +""") + assert result.returncode == 0, result.stderr + assert "LOCAL_TEST" in result.stdout 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 9ecf819d9..4beba6c82 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -20,6 +20,7 @@ def worker_env(monkeypatch): """Make the process look like a real Runpod worker.""" monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.com/job") + monkeypatch.setenv("RUNPOD_FITNESS_WORKER_PID", str(os.getpid())) monkeypatch.delenv("RUNPOD_SKIP_FITNESS_CHECKS", raising=False) monkeypatch.delenv("RUNPOD_DEFER_FITNESS_CHECKS", raising=False) @@ -269,11 +270,11 @@ def 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 + sys.modules, "runpod._health.gpu", fake_gpu_module ) monkeypatch.setitem( sys.modules, - "runpod.serverless.modules.rp_system_fitness", + "runpod._health.system", fake_system_module, ) @@ -294,7 +295,7 @@ def guard_no_torch(name, *args, **kwargs): class TestImportWiring: """Deleting the wiring must fail a test, not just real workers.""" - def test_serverless_import_calls_startup_checks(self, monkeypatch): + def test_top_level_import_calls_startup_checks(self, worker_env, monkeypatch): import importlib import runpod.serverless @@ -304,7 +305,7 @@ def test_serverless_import_calls_startup_checks(self, monkeypatch): rp_fitness, "run_startup_fitness_checks", lambda: calls.append(True) ) - importlib.reload(runpod.serverless) + importlib.reload(runpod) assert calls == [True] @@ -316,16 +317,14 @@ class TestRegistrationLatch: async def test_malformed_env_reraises_at_start(self, worker_env, monkeypatch): monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", raising=False) monkeypatch.setenv("RUNPOD_MIN_MEMORY_GB", "not-a-number") - # Drop the cached module so the env parse re-executes on import. - monkeypatch.delitem( - sys.modules, "runpod.serverless.modules.rp_system_fitness", raising=False - ) + # Configuration is parsed when preparing checks, even if already imported. run_startup_fitness_checks() # swallowed and logged — but not latched assert rp_fitness._registration_state["system_checks"] is False - with pytest.raises(ValueError): + with pytest.raises(SystemExit) as exc: await run_fitness_checks() + assert exc.value.code == 1 class TestLateConfigWarning: From 813a29c15fef74c6385469917d9eaa21fdf0eecc Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 10 Sep 2026 17:27:41 -0400 Subject: [PATCH 08/13] fix(serverless): fully redact secrets and remove startup import cycle --- runpod/_health/__init__.py | 21 ++++++++++++++++++- runpod/_health/fitness.py | 2 +- runpod/_logger.py | 12 +++++------ runpod/_startup.py | 17 ++------------- .../test_modules/test_logger.py | 13 +++++++++++- 5 files changed, 40 insertions(+), 25 deletions(-) diff --git a/runpod/_health/__init__.py b/runpod/_health/__init__.py index d4441d68c..0701b8265 100644 --- a/runpod/_health/__init__.py +++ b/runpod/_health/__init__.py @@ -1 +1,20 @@ -"""Worker health checks independent of the serverless import tree.""" +"""Worker-process identity shared by startup and health checks.""" + +import os +import sys + +# Launchers may set this to the PID of the Python handler process before exec. +# A generic container-level boolean would also authorize unrelated processes. +WORKER_PID_ENV = "RUNPOD_FITNESS_WORKER_PID" + + +def is_worker_process() -> bool: + """Require explicit launcher identity and exclude local/API test invocations.""" + return ( + os.environ.get(WORKER_PID_ENV) == str(os.getpid()) + and bool(os.environ.get("RUNPOD_WEBHOOK_GET_JOB")) + and not any( + arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") + for arg in sys.argv[1:] + ) + ) diff --git a/runpod/_health/fitness.py b/runpod/_health/fitness.py index d58e5684f..47159ccda 100644 --- a/runpod/_health/fitness.py +++ b/runpod/_health/fitness.py @@ -20,7 +20,7 @@ from collections.abc import Callable from runpod._logger import RunPodLogger -from runpod._startup import is_worker_process +from . import is_worker_process log = RunPodLogger() diff --git a/runpod/_logger.py b/runpod/_logger.py index 6ef4c5f73..9a4010b0a 100644 --- a/runpod/_logger.py +++ b/runpod/_logger.py @@ -110,13 +110,11 @@ def log(self, message, message_level="INFO", job_id=None): return def secret(self, secret_name, secret): - """ - Censors secrets for logging. - Replaces everything except the first and last characters with * - """ - secret = str(secret) - redacted_secret = secret[0] + "*" * (len(secret) - 2) + secret[-1] - self.info(f"{secret_name}: {redacted_secret}") + """Log the secret's name without exposing its value or length.""" + # Even a one-character value must be completely redacted. Do not call + # str(secret): custom objects may reveal sensitive data or raise. + self.info(f"{secret_name}: [REDACTED]") + def debug(self, message, request_id: Optional[str] = None): """ diff --git a/runpod/_startup.py b/runpod/_startup.py index 303f8c771..f6d312626 100644 --- a/runpod/_startup.py +++ b/runpod/_startup.py @@ -1,23 +1,10 @@ """Process-scoped startup gate; safe to import without loading serverless.""" -import os import sys -# Launchers may set this to the PID of the Python handler process before exec. -# A generic container-level boolean would also authorize unrelated processes. -WORKER_PID_ENV = "RUNPOD_FITNESS_WORKER_PID" +from ._health import WORKER_PID_ENV, is_worker_process - -def is_worker_process() -> bool: - """Require explicit launcher identity and exclude local/API test invocations.""" - return ( - os.environ.get(WORKER_PID_ENV) == str(os.getpid()) - and bool(os.environ.get("RUNPOD_WEBHOOK_GET_JOB")) - and not any( - arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") - for arg in sys.argv[1:] - ) - ) +__all__ = ["WORKER_PID_ENV", "is_worker_process", "run_import_checks"] def run_import_checks() -> None: diff --git a/tests/test_serverless/test_modules/test_logger.py b/tests/test_serverless/test_modules/test_logger.py index ea428379c..83d7813b4 100644 --- a/tests/test_serverless/test_modules/test_logger.py +++ b/tests/test_serverless/test_modules/test_logger.py @@ -105,9 +105,20 @@ def test_log_secret(self): with patch("runpod.serverless.modules.rp_logger.RunPodLogger.log") as mock_log: self.logger.secret("test_secret", "test_secret_value") mock_log.assert_called_once_with( - "test_secret: t***************e", "INFO", None + "test_secret: [REDACTED]", "INFO", None ) + def test_secret_redacts_short_empty_and_object_values(self): + class Sensitive: + def __str__(self): + raise AssertionError("A secret must not be converted to text") + + for value in ("", "a", "ab", "long-secret", None, Sensitive()): + with self.subTest(value_type=type(value).__name__): + with patch.object(self.logger, "log") as mock_log: + self.logger.secret("credential", value) + mock_log.assert_called_once_with("credential: [REDACTED]", "INFO", None) + def test_log_tip(self): """ Tests that the tip method logs a tip. From f05a8a36fc12aa64de3a6ec0b17f28953b9452e9 Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 10 Sep 2026 17:29:55 -0400 Subject: [PATCH 09/13] fix(logging): distinguish credential labels from secret values --- runpod/_logger.py | 16 ++++++++++++---- .../test_serverless/test_modules/test_logger.py | 11 +++++++++++ 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/runpod/_logger.py b/runpod/_logger.py index 9a4010b0a..b4196a284 100644 --- a/runpod/_logger.py +++ b/runpod/_logger.py @@ -109,12 +109,20 @@ def log(self, message, message_level="INFO", job_id=None): print(f"{message_level.ljust(7)}| {message}", flush=True) return - def secret(self, secret_name, secret): - """Log the secret's name without exposing its value or length.""" + def secret(self, name=None, secret=None, **kwargs): + """Log a credential label without exposing its value or length. + + `secret_name=` remains accepted for compatibility with older callers. + """ + if "secret_name" in kwargs: + if name is not None: + raise TypeError("Pass either name or secret_name, not both") + name = kwargs.pop("secret_name") + if kwargs: + raise TypeError("Unexpected keyword argument to secret()") # Even a one-character value must be completely redacted. Do not call # str(secret): custom objects may reveal sensitive data or raise. - self.info(f"{secret_name}: [REDACTED]") - + self.info(f"{name}: [REDACTED]") def debug(self, message, request_id: Optional[str] = None): """ diff --git a/tests/test_serverless/test_modules/test_logger.py b/tests/test_serverless/test_modules/test_logger.py index 83d7813b4..c63c360d8 100644 --- a/tests/test_serverless/test_modules/test_logger.py +++ b/tests/test_serverless/test_modules/test_logger.py @@ -119,6 +119,17 @@ def __str__(self): self.logger.secret("credential", value) mock_log.assert_called_once_with("credential: [REDACTED]", "INFO", None) + def test_secret_legacy_keyword_label(self): + with patch.object(self.logger, "log") as mock_log: + self.logger.secret(secret_name="credential", secret="sensitive") + mock_log.assert_called_once_with("credential: [REDACTED]", "INFO", None) + + def test_secret_rejects_conflicting_or_unknown_labels(self): + with self.assertRaises(TypeError): + self.logger.secret("first", "sensitive", secret_name="second") + with self.assertRaises(TypeError): + self.logger.secret("credential", "sensitive", unexpected="value") + def test_log_tip(self): """ Tests that the tip method logs a tip. From cf1650f50785d46168c5eeb7e1cd820e66c5bae2 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 11 Sep 2026 10:36:56 -0400 Subject: [PATCH 10/13] Run early health checks once per serverless container startup --- ARCHITECTURE.md | 2 +- README.md | 2 +- docs/serverless/worker_fitness_checks.md | 35 +-- pyproject.toml | 1 - runpod/_health/__init__.py | 26 ++- runpod/_health/coordination.py | 100 +++++++++ runpod/_health/fitness.py | 101 +++++++-- runpod/_startup.py | 8 +- runpod/_worker_bootstrap.py | 47 ---- .../test_modules/test_fitness/conftest.py | 14 +- .../test_fitness/test_coordination.py | 210 ++++++++++++++++++ .../test_fitness/test_safe_startup.py | 123 +--------- .../test_modules/test_fitness/test_startup.py | 59 ++--- 13 files changed, 469 insertions(+), 259 deletions(-) create mode 100644 runpod/_health/coordination.py delete mode 100644 runpod/_worker_bootstrap.py create mode 100644 tests/test_serverless/test_modules/test_fitness/test_coordination.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 998eab63c..99a81e8dc 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. `runpod-worker` identifies the handler process and runs early hardware checks before executing it. Existing launchers may authorize the top-level import hook using `RUNPOD_FITNESS_WORKER_PID=`. The hook and check engine do not import `serverless`; ordinary imports with only the webhook environment are exempt. Legacy launches and `RUNPOD_DEFER_FITNESS_CHECKS=true` run checks only at worker start. Network readiness, CUDA initialization, compute and custom checks run in the final pass; production realtime uses the serving process's lifespan. Successful early checks are reused unless their configuration changes. +1. The first top-level import with both `RUNPOD_ENDPOINT_ID` and `RUNPOD_WEBHOOK_GET_JOB` runs shared hardware checks, excluding test invocations and `RUNPOD_TEST`. A Linux file lock and container-start-scoped result prevent concurrent/repeated execution across processes; saved failures propagate to later workers. Identity includes host boot, PID namespace, and PID 1 start time. Network, Python CUDA initialization, compute, and custom checks remain at worker start (realtime uses serving lifespan). Unsupported/unwritable coordination defers to worker-start checks. Lock waiting is bounded at 35 seconds; a worker-start timeout fails closed. `RUNPOD_DEFER_FITNESS_CHECKS=true` postpones early checks. No custom launcher or PID environment variable is required. 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/README.md b/README.md index 746d09fbf..a273a7049 100644 --- a/README.md +++ b/README.md @@ -143,7 +143,7 @@ runpod.serverless.start({"handler": handler}) **Key Features:** - Supports both synchronous and asynchronous check functions -- `runpod-worker handler.py` checks hardware before model loading; existing Python launches check at worker start +- Shared hardware checks run once per container at the first Serverless import; network and process-specific checks run at worker start - Local tests/helper imports remain exempt; network readiness and custom checks run at worker start - Successful early checks are reused unless their configuration changes - Runs before handler initialization and job processing begins diff --git a/docs/serverless/worker_fitness_checks.md b/docs/serverless/worker_fitness_checks.md index d32ed76f1..4af57b8d9 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -43,34 +43,19 @@ if __name__ == "__main__": ## When Checks Run -Early checks are automatic when the handler is launched with `runpod-worker`: +On Serverless, the first `import runpod` runs RAM, disk, CUDA-version, and native GPU health checks. Eligibility requires both `RUNPOD_ENDPOINT_ID` and `RUNPOD_WEBHOOK_GET_JOB`. Platform tests (`RUNPOD_TEST`), `--test_input`, and local `--rp_serve_api` invocations skip early checks. -```bash -runpod-worker handler.py -# Or a module: -runpod-worker -m my_package.handler -``` - -The launcher runs memory, disk, CUDA-version and native GPU checks **before executing the handler**, then runs the same handler with its original arguments. Customers do not need to add imports or check calls. Platform-managed launchers can adopt this entrypoint without changing handler code. Existing `python handler.py` launches continue checking at worker start, preserving compatibility. - -Launchers that already manage Python directly may instead set `RUNPOD_FITNESS_WORKER_PID` to the PID of the Python handler process **before exec**. The top-level `runpod` import checks that exact PID. This hook is independent of `runpod.serverless`, including when that module is lazy-loaded. Do not set a fixed PID in a Dockerfile or template. `RUNPOD_WEBHOOK_GET_JOB` alone never authorizes import-time checks. - -Network readiness, CUDA initialization, the GPU compute benchmark, and customer-registered checks run at worker start. Network checks use bounded retries against the worker API host; they cannot terminate a process during import. CUDA checks that initialize a context remain deferred so handler code can create child processes first. +A shared Linux file lock serializes these checks across Python processes. Successful results are reused by helpers and the actual worker; failures are saved before reporting unhealthy and exiting, so another process cannot silently ignore the failure. Results are scoped to the host boot, PID namespace, and container init process start time, rather than just the Pod ID. Checks with changed settings or SDK version are rerun. The first import may occur after model loading; no earlier timing is guaranteed in that case. -Checks that passed early are not repeated unless their settings changed. Without launcher identification, all checks run at worker start. `RUNPOD_DEFER_FITNESS_CHECKS=true` restores worker-start-only timing even with the new launcher; `RUNPOD_SKIP_FITNESS_CHECKS=true` disables all checks. +Network connectivity, Python CUDA initialization, GPU compute, and customer-registered checks run in the worker process at `.start()`, before accepting jobs. Network checks retry against the worker API with a bounded budget. Keeping Python CUDA initialization out of imports protects subsequent customer forks. Production realtime mode runs its final checks in serving lifespan. -### Compatibility and failure handling +Coordination uses a fixed `/tmp` path shared by container processes. If procfs or shared state is unavailable, early checks defer to worker start. Imports wait at most 35 seconds for another checking process; a timeout defers to worker start. If the lock is still busy at worker start, the worker reports failure and exits rather than accepting jobs without validation. An owner crash releases its OS lock, allowing a later process to retry unfinished checks. Processes with separate filesystems or incompatible file permissions cannot share results and use the fallback. -- Imports in helper scripts and ordinary local tests are safe even when they inherit worker environment variables. `--test_input` (both argument forms) and local `--rp_serve_api` invocations skip early checks. The launcher removes its process authorization before executing the handler; children cannot inherit permission to run early checks. -- Set thresholds before launch for early validation. Settings changed afterward are applied at worker start, with a warning; affected checks are rerun, while unrelated successful checks remain completed. Earlier failures cannot be undone by changing settings later. Use deferral when the handler must configure checks before they run. -- With early checks enabled, the memory check measures available memory **before model loading**. With legacy/deferred startup it measures available memory at worker start. -- Production realtime mode (`RUNPOD_REALTIME_PORT` plus worker environment) runs the final checks in the serving process's application lifespan before accepting requests. Local API simulation remains exempt. -- A failed health check reports the failure and force-exits. An error preparing early checks is logged and retried at worker start. An unresolved setup/configuration error at worker start reports `fitness_check_setup` and force-exits, including when background threads are alive. -- Early execution inside an already-running event loop defers to worker start; it does not replace the customer's event loop. +`RUNPOD_DEFER_FITNESS_CHECKS=true` restores worker-start timing. `RUNPOD_SKIP_FITNESS_CHECKS=true` disables all checks. Set early thresholds before importing the SDK; late changes are applied at worker start, but cannot undo an earlier failure. No launcher or Docker entrypoint changes are needed. -### Platform rollout +### Rollout -Ship the SDK first with legacy launch behavior preserved. Enable `runpod-worker` in a small set of managed worker launches, validate real GPU/fork behavior and startup failure rates, then expand. Deployments with custom entrypoints retain worker-start checks until their launcher integrates the process hook. Roll back early timing centrally with `RUNPOD_DEFER_FITNESS_CHECKS=true`; no handler edits are needed. This SDK change supplies the launcher and hook; it does not change deployed platform launch configuration. +Validate in a small set of workers before broader rollout. The deferral variable provides a rollback of early timing without handler edits. This SDK change does not itself alter deployed platform configuration. ## Async Fitness Checks @@ -400,7 +385,7 @@ ENV RUNPOD_NETWORK_CHECK_TIMEOUT=10 ENV RUNPOD_GPU_BENCHMARK_TIMEOUT=2 ``` -For legacy/deferred launches, settings can also be configured in Python before worker start: +For deferred launches, settings can also be configured in Python before worker start: ```python import os @@ -434,7 +419,7 @@ os.environ["RUNPOD_SKIP_GPU_CHECK"] = "true" import runpod ``` -For early checks, set these before launching the handler. For legacy/deferred launches, set them before worker start. +For early checks, set these before launching the handler. For deferred launches, set them before worker start. 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. @@ -442,7 +427,7 @@ User-registered checks via `@register_fitness_check` still run regardless of `RU ### Execution Timing -- Early checks run only in launcher-identified worker processes; the final pass runs before job processing. Successful checks are reused unless their configuration changes. +- Early checks run in eligible Serverless containers; the final pass runs before job processing. Successful checks are reused unless their configuration changes. - 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 diff --git a/pyproject.toml b/pyproject.toml index 6faca1cdc..d4639a153 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,6 @@ local_scheme = "no-local-version" [project.scripts] runpod = "runpod.cli.entry:runpod_cli" -runpod-worker = "runpod._worker_bootstrap:main" [dependency-groups] diff --git a/runpod/_health/__init__.py b/runpod/_health/__init__.py index 0701b8265..89f67ee89 100644 --- a/runpod/_health/__init__.py +++ b/runpod/_health/__init__.py @@ -1,20 +1,22 @@ -"""Worker-process identity shared by startup and health checks.""" +"""Lightweight Serverless environment detection; no SDK imports.""" import os import sys -# Launchers may set this to the PID of the Python handler process before exec. -# A generic container-level boolean would also authorize unrelated processes. -WORKER_PID_ENV = "RUNPOD_FITNESS_WORKER_PID" + +def is_serverless_environment() -> bool: + """Recognize production worker configuration, excluding platform tests.""" + return ( + bool(os.environ.get("RUNPOD_ENDPOINT_ID", "").strip()) + and bool(os.environ.get("RUNPOD_WEBHOOK_GET_JOB", "").strip()) + and os.environ.get("RUNPOD_TEST", "").strip().lower() + not in ("1", "true", "yes", "on") + ) def is_worker_process() -> bool: - """Require explicit launcher identity and exclude local/API test invocations.""" - return ( - os.environ.get(WORKER_PID_ENV) == str(os.getpid()) - and bool(os.environ.get("RUNPOD_WEBHOOK_GET_JOB")) - and not any( - arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") - for arg in sys.argv[1:] - ) + """Eligibility for shared early checks, not an assertion of process identity.""" + return is_serverless_environment() and not any( + arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") + for arg in sys.argv[1:] ) diff --git a/runpod/_health/coordination.py b/runpod/_health/coordination.py new file mode 100644 index 000000000..fd0fc27d4 --- /dev/null +++ b/runpod/_health/coordination.py @@ -0,0 +1,100 @@ +"""Linux container-start-scoped coordination for shared health checks.""" + +import asyncio +import hashlib +import json +import os +import time +from pathlib import Path + + +class CoordinationUnavailable(Exception): + """Shared state cannot be used; worker-start checks remain available.""" + + +class CoordinationBusy(Exception): + """Another process did not finish within the bounded wait.""" + + +def container_start_id() -> str: + """PID namespace + init start ticks + host boot distinguish container restarts. + + Use fixed /tmp rather than TMPDIR (which can differ between processes). + No pod-id-only marker: container files can survive a restart. + """ + boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip() + stat = Path("/proc/1/stat").read_text() + start_ticks = stat.rsplit(")", 1)[1].split()[19] + namespace = os.readlink("/proc/1/ns/pid") + return hashlib.sha256(f"{boot}:{namespace}:{start_ticks}".encode()).hexdigest() + + +class ContainerChecks: + """Hold one flock while reading, executing, and recording early checks.""" + + def __init__(self, timeout=35): + self.timeout = timeout + self.fd = None + self.state = {"passed": [], "failure": None} + + async def __aenter__(self): + try: + import fcntl + + identity = container_start_id() + path = f"/tmp/runpod-fitness-{identity}.json" + self.fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600) + os.set_inheritable(self.fd, False) + except (OSError, ValueError, IndexError, ImportError) as exc: + self.close() + raise CoordinationUnavailable(str(exc)) from exc + deadline = time.monotonic() + self.timeout + try: + while True: + try: + fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() >= deadline: + raise CoordinationBusy( + "Timed out waiting for early health checks" + ) + await asyncio.sleep(0.05) + raw = os.read(self.fd, 65536) + if raw: + self.state = json.loads(raw) + if ( + not isinstance(self.state, dict) + or not isinstance(self.state.get("passed"), list) + or not all(isinstance(key, str) for key in self.state["passed"]) + or not isinstance(self.state.get("failure"), (str, type(None))) + ): + raise ValueError("Invalid health-check state") + return self + except (OSError, ValueError) as exc: + self.close() + raise CoordinationUnavailable(str(exc)) from exc + except BaseException: + self.close() + raise + + def save(self): + """Persist before releasing the lock or terminating on failure.""" + data = json.dumps(self.state).encode() + os.lseek(self.fd, 0, os.SEEK_SET) + remaining = memoryview(data) + while remaining: + written = os.write(self.fd, remaining) + if written <= 0: + raise OSError("Unable to persist health-check state") + remaining = remaining[written:] + os.ftruncate(self.fd, len(data)) + os.fsync(self.fd) + + def close(self): + if self.fd is not None: + os.close(self.fd) + self.fd = None + + async def __aexit__(self, *args): + self.close() diff --git a/runpod/_health/fitness.py b/runpod/_health/fitness.py index 47159ccda..55e9fb533 100644 --- a/runpod/_health/fitness.py +++ b/runpod/_health/fitness.py @@ -13,6 +13,8 @@ import asyncio import contextlib import inspect +import hashlib +import json import os import sys import time @@ -21,6 +23,7 @@ from runpod._logger import RunPodLogger from . import is_worker_process +from .coordination import ContainerChecks, CoordinationUnavailable, CoordinationBusy log = RunPodLogger() @@ -63,11 +66,6 @@ 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" - # Tuning vars consumed when the checks run. Snapshotted at the import-time # pass so a later pass can warn about post-import changes, which would # otherwise be silently ignored. @@ -352,6 +350,74 @@ def _fail_worker(check_name: str, exc: Exception) -> None: _terminate_unhealthy(1) +async def _run_shared_checks(include_deferred: bool) -> None: + """Reuse container checks across imports, including independent helpers.""" + try: + async with ContainerChecks() as shared: + if shared.state.get("failure"): + _fail_worker( + "early_container_check", RuntimeError(shared.state["failure"]) + ) + return + for check in _fitness_checks: + if not getattr(check, "_runpod_builtin", False) or _is_deferred(check): + continue + # Version and relevant settings prevent reusing incompatible results. + from runpod.version import __version__ + + dependencies = { + "_memory_check": ("RUNPOD_MIN_MEMORY_GB",), + "_disk_check": ("RUNPOD_MIN_DISK_PERCENT",), + "_cuda_version_check": ("RUNPOD_MIN_CUDA_VERSION",), + "_gpu_health_check": ( + "RUNPOD_GPU_TEST_TIMEOUT", + "RUNPOD_GPU_MAX_ERROR_MESSAGES", + "RUNPOD_BINARY_GPU_TEST_PATH", + ), + } + settings = { + k: os.environ.get(k) for k in dependencies.get(check.__name__, ()) + } + key = hashlib.sha256( + json.dumps( + [__version__, check._runpod_builtin, check.__name__, settings], + sort_keys=True, + ).encode() + ).hexdigest() + if key not in shared.state["passed"]: + try: + if inspect.iscoroutinefunction(check): + await check() + else: + check() + except Exception as exc: + shared.state["failure"] = ( + f"{check.__name__}: {type(exc).__name__}" + ) + try: + shared.save() + finally: + _fail_worker(check.__name__, exc) + return + shared.state["passed"].append(key) + shared.save() + if not any(check is done for done in _completed_checks): + _completed_checks.append(check) + except CoordinationUnavailable as exc: + log.warn( + f"Early check coordination unavailable; using worker-start checks: {exc}" + ) + except CoordinationBusy as exc: + if include_deferred: + _fail_worker("fitness_check_coordination", exc) + else: + log.warn( + "Early checks still running in another process; deferring to worker start." + ) + except OSError as exc: + log.warn(f"Cannot save shared checks; using worker-start checks: {exc}") + + async def run_fitness_checks(include_deferred: bool = True) -> None: """ Execute all registered fitness checks sequentially at startup. @@ -405,6 +471,13 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: _fail_worker("fitness_check_setup", exc) return + if is_worker_process() and ( + include_deferred or not _env_flag(DEFER_FITNESS_CHECKS_ENV) + ): + await _run_shared_checks(include_deferred) + if not include_deferred: + return + # Identity, not equality: two distinct registrations may compare equal # (e.g. fresh bound-method objects of one method), and `==` would skip one. pending = [ @@ -414,7 +487,11 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: ] if not include_deferred: - pending = [check for check in pending if not _is_deferred(check)] + pending = [ + check + for check in pending + if getattr(check, "_runpod_builtin", False) and not _is_deferred(check) + ] if not pending: log.debug("No pending fitness checks, skipping.") @@ -466,11 +543,8 @@ def run_startup_fitness_checks() -> None: 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 without launcher process authorization, for local tests, when checks - are disabled or deferred, inside a running event loop, and in child - processes (the launcher PID must match and is consumed before the handler). 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. + Shared built-ins run once per container startup. Child processes reuse the + result; process-specific and customer checks wait for worker start. """ if _env_flag(SKIP_FITNESS_CHECKS_ENV) or _env_flag(DEFER_FITNESS_CHECKS_ENV): return @@ -478,11 +552,6 @@ def run_startup_fitness_checks() -> None: if not is_worker_process(): return - if os.environ.get(_CHECKS_DONE_ENV): - return - os.environ[_CHECKS_DONE_ENV] = "1" - os.environ.pop("RUNPOD_FITNESS_WORKER_PID", None) - if _event_loop_running(): log.debug("Event loop already running, deferring fitness checks to run_worker.") return diff --git a/runpod/_startup.py b/runpod/_startup.py index f6d312626..be80cb939 100644 --- a/runpod/_startup.py +++ b/runpod/_startup.py @@ -1,14 +1,14 @@ -"""Process-scoped startup gate; safe to import without loading serverless.""" +"""Container-scoped startup gate; safe to import without loading serverless.""" import sys -from ._health import WORKER_PID_ENV, is_worker_process +from ._health import is_worker_process -__all__ = ["WORKER_PID_ENV", "is_worker_process", "run_import_checks"] +__all__ = ["is_worker_process", "run_import_checks"] def run_import_checks() -> None: - """Run early checks only in the handler process selected by the launcher.""" + """Run shared early checks in eligible Serverless containers.""" if not is_worker_process(): return try: diff --git a/runpod/_worker_bootstrap.py b/runpod/_worker_bootstrap.py deleted file mode 100644 index a71a83c1d..000000000 --- a/runpod/_worker_bootstrap.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Launch a worker with health checks before executing its handler module. - -Usage: runpod-worker handler.py [handler arguments] - runpod-worker -m package.handler [handler arguments] -""" - -import os -import runpy -import sys -from pathlib import Path - -from ._startup import WORKER_PID_ENV, run_import_checks - - -def main() -> None: - """Select this process as the worker, then execute the unmodified handler.""" - args = sys.argv[1:] - module_mode = bool(args and args[0] == "-m") - if module_mode: - args = args[1:] - if not args or args[0].startswith("-"): - raise SystemExit("Usage: runpod-worker [-m] handler [arguments]") - target, *handler_args = args - sys.argv = [target, *handler_args] - if not module_mode: - # Match `python handler.py`: sibling imports resolve beside the script. - target = str(Path(target).resolve()) - if not Path(target).is_file(): - raise SystemExit(f"Worker handler not found: {target}") - sys.path.insert(0, str(Path(target).parent)) - else: - # Console entrypoints put their bin directory on sys.path, unlike python -m. - sys.path.insert(0, os.getcwd()) - os.environ[WORKER_PID_ENV] = str(os.getpid()) - try: - run_import_checks() - finally: - # Helpers, subprocesses and multiprocessing children are not workers. - os.environ.pop(WORKER_PID_ENV, None) - if module_mode: - runpy.run_module(target, run_name="__main__", alter_sys=True) - else: - runpy.run_path(target, run_name="__main__") - - -if __name__ == "__main__": - main() diff --git a/tests/test_serverless/test_modules/test_fitness/conftest.py b/tests/test_serverless/test_modules/test_fitness/conftest.py index 04a6f9e68..42e801896 100644 --- a/tests/test_serverless/test_modules/test_fitness/conftest.py +++ b/tests/test_serverless/test_modules/test_fitness/conftest.py @@ -10,7 +10,7 @@ @pytest.fixture(autouse=True) -def cleanup_fitness_checks(monkeypatch): +def cleanup_fitness_checks(monkeypatch, tmp_path): """Automatically clean up fitness checks before and after each test. Disables auto-registration of system checks to avoid interference @@ -21,11 +21,17 @@ def cleanup_fitness_checks(monkeypatch): to raise SystemExit(1) so tests can assert exit behavior in-process. Tests that need the real os._exit patch it themselves. """ + from runpod._health import coordination + + monkeypatch.setattr( + coordination, + "container_start_id", + lambda: tmp_path.parent.name + "-" + tmp_path.name, + ) + monkeypatch.delenv("RUNPOD_ENDPOINT_ID", raising=False) + monkeypatch.delenv("RUNPOD_TEST", raising=False) 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_coordination.py b/tests/test_serverless/test_modules/test_fitness/test_coordination.py new file mode 100644 index 000000000..11adaa298 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_coordination.py @@ -0,0 +1,210 @@ +"""Real interprocess locks/results, and container-start identity regressions.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from runpod._health import coordination, fitness, is_worker_process +from runpod._health.coordination import container_start_id as real_container_start_id + + +@pytest.mark.parametrize( + "endpoint,webhook,test,expected", + [ + ("", "", "", False), + ("ep", "", "", False), + ("", "https://example.test/job", "", False), + ("ep", "https://example.test/job", "", True), + ("ep", "https://example.test/job", "TRUE", False), + ("ep", "https://example.test/job", "1", False), + ], +) +def test_environment_gate(monkeypatch, endpoint, webhook, test, expected): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", endpoint) + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", webhook) + monkeypatch.setenv("RUNPOD_TEST", test) + monkeypatch.setattr(sys, "argv", ["handler.py"]) + assert is_worker_process() is expected + + +def process_code(tmp_path, fail=False): + return f""" +import asyncio, os, time +from pathlib import Path +from runpod._health import fitness, coordination +coordination.container_start_id = lambda: {(tmp_path.parent.name + "-" + tmp_path.name)!r} +os.environ['RUNPOD_ENDPOINT_ID'] = 'ep' +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +fitness._report_unhealthy = lambda *args: None +@fitness.register_fitness_check +def shared_check(): + with open({str(tmp_path / "calls")!r}, 'a') as out: + out.write('check\\n') + time.sleep(0.2) + if {fail!r}: + raise RuntimeError('failed') +shared_check._runpod_builtin = 'system_checks' +asyncio.run(fitness.run_fitness_checks(include_deferred=False)) +asyncio.run(fitness.run_fitness_checks()) +""" + + +def launch(code): + env = {k: v for k, v in os.environ.items() if not k.startswith("RUNPOD_")} + env.update(RUNPOD_SKIP_GPU_CHECK="true", RUNPOD_SKIP_AUTO_SYSTEM_CHECKS="true") + return subprocess.Popen( + [sys.executable, "-c", code], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def test_independent_processes_share_success(tmp_path): + children = [launch(process_code(tmp_path)) for _ in range(3)] + for child in children: + out, err = child.communicate(timeout=10) + assert child.returncode == 0, out + err + assert (tmp_path / "calls").read_text().splitlines() == ["check"] + + +def test_failure_propagates_without_rerunning(tmp_path): + for _ in range(2): + child = launch(process_code(tmp_path, fail=True)) + out, err = child.communicate(timeout=10) + assert child.returncode == 1, out + err + assert (tmp_path / "calls").read_text().splitlines() == ["check"] + + +@pytest.mark.asyncio +async def test_busy_is_bounded_and_owner_crash_releases_lock(tmp_path): + code = f""" +import asyncio, time +from runpod._health import coordination +coordination.container_start_id = lambda: {(tmp_path.parent.name + "-" + tmp_path.name)!r} +async def main(): + async with coordination.ContainerChecks(): + print('LOCKED', flush=True) + time.sleep(30) +asyncio.run(main()) +""" + child = launch(code) + try: + assert child.stdout.readline().strip() == "LOCKED" + with pytest.raises(coordination.CoordinationBusy): + async with coordination.ContainerChecks(timeout=0.1): + pass + finally: + child.kill() + child.communicate(timeout=5) + async with coordination.ContainerChecks(timeout=0.1) as shared: + assert shared.state["passed"] == [] + + +@pytest.mark.asyncio +async def test_restart_does_not_reuse_previous_success(monkeypatch, tmp_path): + async with coordination.ContainerChecks() as shared: + shared.state["passed"] = ["old-success"] + shared.save() + monkeypatch.setattr( + coordination, + "container_start_id", + lambda: (tmp_path.parent.name + "-" + tmp_path.name) + "-restart", + ) + async with coordination.ContainerChecks() as shared: + assert shared.state["passed"] == [] + + +def test_identity_changes_when_init_restarts(monkeypatch): + # Fields following comm start at field 3; starttime is field 22. + ticks = ["100"] + + def read(path): + if str(path).endswith("boot_id"): + return "host-boot" + return "1 (name with spaces) " + " ".join(["0"] * 19 + ticks) + + monkeypatch.setattr(Path, "read_text", read) + monkeypatch.setattr(os, "readlink", lambda path: "pid:[123]") + first = real_container_start_id() + ticks[0] = "200" + assert real_container_start_id() != first + + +@pytest.mark.asyncio +async def test_unavailable_coordination_defers_then_checks_at_start(monkeypatch): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + + def unavailable(): + raise OSError("read-only") + + monkeypatch.setattr(coordination, "container_start_id", unavailable) + calls = [] + + @fitness.register_fitness_check + def early(): + calls.append("early") + + early._runpod_builtin = "system_checks" + + @fitness.register_fitness_check + def customer(): + calls.append("customer") + + await fitness.run_fitness_checks(include_deferred=False) + assert calls == [] + await fitness.run_fitness_checks() + assert calls == ["early", "customer"] + + +@pytest.mark.asyncio +async def test_customer_and_process_checks_only_at_start(monkeypatch): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + calls = [] + + @fitness.register_fitness_check + def customer(): + calls.append("customer") + + @fitness.register_fitness_check + @fitness.defer_to_worker_start + def cuda(): + calls.append("cuda") + + cuda._runpod_builtin = "system_checks" + await fitness.run_fitness_checks(include_deferred=False) + assert calls == [] + await fitness.run_fitness_checks() + assert calls == ["customer", "cuda"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("contents", ["[]", '{"passed": null}', "{broken"]) +async def test_corrupt_state_uses_fallback(tmp_path, contents): + identity = coordination.container_start_id() + Path(f"/tmp/runpod-fitness-{identity}.json").write_text(contents) + with pytest.raises(coordination.CoordinationUnavailable): + async with coordination.ContainerChecks(): + pass + + +@pytest.mark.asyncio +async def test_lock_timeout_defers_import_but_blocks_worker(monkeypatch): + class Busy: + async def __aenter__(self): + raise coordination.CoordinationBusy("still checking") + + async def __aexit__(self, *args): + pass + + monkeypatch.setattr(fitness, "ContainerChecks", Busy) + monkeypatch.setattr(fitness, "_report_unhealthy", lambda *args: None) + await fitness._run_shared_checks(include_deferred=False) + with pytest.raises(SystemExit): + await fitness._run_shared_checks(include_deferred=True) diff --git a/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py b/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py index c2497ec3f..905b0df74 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py @@ -10,8 +10,10 @@ import pytest -from runpod._health import fitness, system -from runpod._startup import WORKER_PID_ENV, run_import_checks +from runpod._health import fitness, system, coordination + +coordination.container_start_id = lambda: "lazy-test-" + str(os.getpid()) +from runpod._startup import run_import_checks @pytest.mark.parametrize( @@ -27,17 +29,9 @@ def test_import_does_not_run_for_unmarked_or_local_process(monkeypatch, args): monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") monkeypatch.setattr(sys, "argv", args) if len(args) > 1: - monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint") else: - monkeypatch.delenv(WORKER_PID_ENV, raising=False) - with patch.object(fitness, "run_startup_fitness_checks") as run: - run_import_checks() - run.assert_not_called() - - -def test_inherited_worker_pid_does_not_authorize_child(monkeypatch): - monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") - monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid() + 1)) + monkeypatch.delenv("RUNPOD_ENDPOINT_ID", raising=False) with patch.object(fitness, "run_startup_fitness_checks") as run: run_import_checks() run.assert_not_called() @@ -45,11 +39,10 @@ def test_inherited_worker_pid_does_not_authorize_child(monkeypatch): def test_initial_pass_does_not_compare_config(monkeypatch): monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") - monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint") with patch.object(fitness, "_refresh_late_config") as refresh: run_import_checks() refresh.assert_not_called() - assert WORKER_PID_ENV not in os.environ @pytest.mark.asyncio @@ -118,7 +111,7 @@ async def test_network_stuck_close_is_bounded(monkeypatch): def test_network_is_deferred_even_in_authorized_worker(monkeypatch): monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS") monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://worker.example/job") - monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint") with ( patch.object(system, "gpu_available", return_value=False), patch.object(system, "_check_memory_availability"), @@ -135,7 +128,7 @@ def test_network_is_deferred_even_in_authorized_worker(monkeypatch): def test_changed_threshold_is_applied_without_rerunning_unrelated_checks(monkeypatch): monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS") monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://worker.example/job") - monkeypatch.setenv(WORKER_PID_ENV, str(os.getpid())) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint") with ( patch.object(system, "gpu_available", return_value=False), patch.object(system, "_check_memory_availability") as memory, @@ -190,29 +183,6 @@ def test_actual_import_is_safe_with_inherited_worker_environment(): assert "IMPORT_SURVIVED" in result.stdout -@pytest.mark.parametrize("module_mode", [False, True]) -def test_launcher_checks_before_handler_and_preserves_arguments(tmp_path, module_mode): - handler = tmp_path / "handler.py" - handler.write_text( - "import os, sys\nassert 'RUNPOD_FITNESS_WORKER_PID' not in os.environ\n" - "assert sys.argv[1:] == ['--customer-arg', 'value']\nprint('MODEL_LOAD')\n" - ) - result = run_child(f""" -import os, sys -import runpod._worker_bootstrap as bootstrap -from runpod._health import fitness -os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' -fitness.register_fitness_check(lambda: print('EARLY_CHECK')) -os.chdir({str(tmp_path)!r}) -# Model console entrypoint sys.path: the current directory is not pre-added. -sys.path = [p for p in sys.path if p] -sys.argv = ['runpod-worker'] + {(["-m", "handler"] if module_mode else [str(handler)])!r} + ['--customer-arg', 'value'] -bootstrap.main() -""") - assert result.returncode == 0, result.stderr - assert result.stdout.index("EARLY_CHECK") < result.stdout.index("MODEL_LOAD") - - def test_setup_failure_exits_with_live_thread(): result = run_child(""" import asyncio, os, threading, time @@ -240,10 +210,11 @@ def find_spec(self, fullname, path=None, target=None): raise AssertionError('early check loaded ' + fullname) sys.meta_path.insert(0, Guard()) from unittest.mock import patch, MagicMock -from runpod._health import fitness, system +from runpod._health import fitness, system, coordination +coordination.container_start_id = lambda: 'lazy-test-' + str(os.getpid()) from runpod._startup import run_import_checks os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' -os.environ['RUNPOD_FITNESS_WORKER_PID'] = str(os.getpid()) +os.environ['RUNPOD_ENDPOINT_ID'] = 'endpoint' os.environ['RUNPOD_SKIP_AUTO_SYSTEM_CHECKS'] = 'false' loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) @@ -261,73 +232,3 @@ def find_spec(self, fullname, path=None, target=None): """) assert result.returncode == 0, result.stderr assert "LAZY_PASS" in result.stdout - - -@pytest.mark.parametrize("method", ["spawn", "fork"]) -def test_child_processes_do_not_repeat_early_checks(tmp_path, method): - import multiprocessing - - if method not in multiprocessing.get_all_start_methods(): - pytest.skip(f"{method} is not supported") - handler = tmp_path / "child_handler.py" - handler.write_text(""" -import multiprocessing, os -from runpod._startup import is_worker_process - -def child(): - assert not is_worker_process() - print('CHILD_SAFE', flush=True) - -if __name__ == '__main__': - child_process = multiprocessing.get_context(os.environ['TEST_START_METHOD']).Process(target=child) - child_process.start() - child_process.join(5) - assert child_process.exitcode == 0 -""") - result = run_child(f""" -import os, sys -from runpod._worker_bootstrap import main -os.environ['TEST_START_METHOD'] = {method!r} -os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' -sys.argv = ['runpod-worker', {str(handler)!r}] -main() -""") - assert result.returncode == 0, result.stderr - assert "CHILD_SAFE" in result.stdout - - -def test_launcher_failed_early_check_prevents_model_load(tmp_path): - handler = tmp_path / "handler.py" - handler.write_text("print('MODEL_LOAD')\n") - result = run_child(f""" -import os, sys -from runpod._worker_bootstrap import main -from runpod._health import fitness -os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' -def fail(): - raise RuntimeError('broken hardware') -fitness.register_fitness_check(fail) -sys.argv = ['runpod-worker', {str(handler)!r}] -main() -""") - assert result.returncode == 1 - assert "broken hardware" in result.stdout - assert "MODEL_LOAD" not in result.stdout - - -def test_launcher_local_test_does_not_run_early_checks(tmp_path): - handler = tmp_path / "handler.py" - handler.write_text("print('LOCAL_TEST')\n") - result = run_child(f""" -import os, sys -from runpod._worker_bootstrap import main -from runpod._health import fitness -os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' -def fail(): - raise RuntimeError('must not run') -fitness.register_fitness_check(fail) -sys.argv = ['runpod-worker', {str(handler)!r}, '--test_input={{}}'] -main() -""") - assert result.returncode == 0, result.stderr - assert "LOCAL_TEST" in result.stdout 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 4beba6c82..f3204b6ab 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -16,11 +16,17 @@ ) +def register_early_check(func): + """Synthetic built-in for exercising the early runner.""" + func._runpod_builtin = "system_checks" + return register_fitness_check(func) + + @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.setenv("RUNPOD_FITNESS_WORKER_PID", str(os.getpid())) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint") monkeypatch.delenv("RUNPOD_SKIP_FITNESS_CHECKS", raising=False) monkeypatch.delenv("RUNPOD_DEFER_FITNESS_CHECKS", raising=False) @@ -31,7 +37,7 @@ async def test_skip_env_var_bypasses_all_checks(self, monkeypatch): monkeypatch.setenv("RUNPOD_SKIP_FITNESS_CHECKS", "true") called = [] - @register_fitness_check + @register_early_check def check(): called.append(True) @@ -43,7 +49,7 @@ async def test_checks_run_when_skip_unset(self, monkeypatch): monkeypatch.delenv("RUNPOD_SKIP_FITNESS_CHECKS", raising=False) called = [] - @register_fitness_check + @register_early_check def check(): called.append(True) @@ -56,13 +62,13 @@ class TestRunOnce: async def test_passed_check_does_not_rerun(self): calls = [] - @register_fitness_check + @register_early_check def first(): calls.append("first") await run_fitness_checks() - @register_fitness_check + @register_early_check def second(): calls.append("second") @@ -95,7 +101,7 @@ class TestStartupEntrypoint: def test_runs_checks_on_worker(self, worker_env): calls = [] - @register_fitness_check + @register_early_check def check(): calls.append(True) @@ -106,7 +112,7 @@ def test_noop_outside_worker(self, monkeypatch): monkeypatch.delenv("RUNPOD_WEBHOOK_GET_JOB", raising=False) calls = [] - @register_fitness_check + @register_early_check def check(): calls.append(True) @@ -117,7 +123,7 @@ def test_defer_env_var_postpones_to_worker_start(self, worker_env, monkeypatch): monkeypatch.setenv("RUNPOD_DEFER_FITNESS_CHECKS", "true") calls = [] - @register_fitness_check + @register_early_check def check(): calls.append(True) @@ -128,7 +134,7 @@ def test_skip_env_var_respected(self, worker_env, monkeypatch): monkeypatch.setenv("RUNPOD_SKIP_FITNESS_CHECKS", "1") calls = [] - @register_fitness_check + @register_early_check def check(): calls.append(True) @@ -147,7 +153,7 @@ def test_unexpected_error_does_not_propagate(self, worker_env): async def test_noop_inside_running_loop(self, worker_env): calls = [] - @register_fitness_check + @register_early_check def check(): calls.append(True) @@ -161,11 +167,11 @@ class TestDeferredChecks: def test_deferred_check_skipped_at_import(self, worker_env): calls = [] - @register_fitness_check + @register_early_check def early(): calls.append("early") - @register_fitness_check + @register_early_check @rp_fitness.defer_to_worker_start def late(): calls.append("late") @@ -177,7 +183,7 @@ def late(): async def test_deferred_check_runs_at_worker_start(self, worker_env): calls = [] - @register_fitness_check + @register_early_check @rp_fitness.defer_to_worker_start def late(): calls.append("late") @@ -200,25 +206,6 @@ def test_cuda_checks_are_marked_deferred(self): 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(): - calls.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.""" @@ -227,11 +214,11 @@ async def test_deferred_to_start_runs_everything(self, worker_env, monkeypatch): monkeypatch.setenv("RUNPOD_DEFER_FITNESS_CHECKS", "true") calls = [] - @register_fitness_check + @register_early_check def check(): calls.append(True) - @register_fitness_check + @register_early_check @rp_fitness.defer_to_worker_start def deferred(): calls.append("deferred") @@ -269,9 +256,7 @@ def 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._health.gpu", fake_gpu_module - ) + monkeypatch.setitem(sys.modules, "runpod._health.gpu", fake_gpu_module) monkeypatch.setitem( sys.modules, "runpod._health.system", From f2884dbcc9c1c88b566d8bbb79d986b287fadd98 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 11 Sep 2026 10:49:20 -0400 Subject: [PATCH 11/13] Simplify health-check coordination and configuration flow --- runpod/_health/__init__.py | 2 +- runpod/_health/coordination.py | 56 +++-- runpod/_health/fitness.py | 206 ++++++++---------- runpod/_startup.py | 6 +- .../test_fitness/test_coordination.py | 4 +- 5 files changed, 126 insertions(+), 148 deletions(-) diff --git a/runpod/_health/__init__.py b/runpod/_health/__init__.py index 89f67ee89..e8c1d713b 100644 --- a/runpod/_health/__init__.py +++ b/runpod/_health/__init__.py @@ -14,7 +14,7 @@ def is_serverless_environment() -> bool: ) -def is_worker_process() -> bool: +def is_early_check_eligible() -> bool: """Eligibility for shared early checks, not an assertion of process identity.""" return is_serverless_environment() and not any( arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") diff --git a/runpod/_health/coordination.py b/runpod/_health/coordination.py index fd0fc27d4..2f4725214 100644 --- a/runpod/_health/coordination.py +++ b/runpod/_health/coordination.py @@ -32,7 +32,7 @@ def container_start_id() -> str: class ContainerChecks: """Hold one flock while reading, executing, and recording early checks.""" - def __init__(self, timeout=35): + def __init__(self, timeout: float = 35): self.timeout = timeout self.fd = None self.state = {"passed": [], "failure": None} @@ -48,28 +48,9 @@ async def __aenter__(self): except (OSError, ValueError, IndexError, ImportError) as exc: self.close() raise CoordinationUnavailable(str(exc)) from exc - deadline = time.monotonic() + self.timeout try: - while True: - try: - fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - break - except BlockingIOError: - if time.monotonic() >= deadline: - raise CoordinationBusy( - "Timed out waiting for early health checks" - ) - await asyncio.sleep(0.05) - raw = os.read(self.fd, 65536) - if raw: - self.state = json.loads(raw) - if ( - not isinstance(self.state, dict) - or not isinstance(self.state.get("passed"), list) - or not all(isinstance(key, str) for key in self.state["passed"]) - or not isinstance(self.state.get("failure"), (str, type(None))) - ): - raise ValueError("Invalid health-check state") + await self._acquire_lock(fcntl) + self._load_state() return self except (OSError, ValueError) as exc: self.close() @@ -78,7 +59,34 @@ async def __aenter__(self): self.close() raise - def save(self): + async def _acquire_lock(self, fcntl) -> None: + deadline = time.monotonic() + self.timeout + while True: + try: + fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return + except BlockingIOError: + if time.monotonic() >= deadline: + raise CoordinationBusy("Timed out waiting for early health checks") + await asyncio.sleep(0.05) + + def _load_state(self) -> None: + raw = os.read(self.fd, 65536) + if not raw: + return + state = json.loads(raw) + if not isinstance(state, dict): + raise ValueError("Health-check state must be an object") + passed = state.get("passed") + if not isinstance(passed, list) or not all( + isinstance(key, str) for key in passed + ): + raise ValueError("Passed health checks must be a list of cache keys") + if not isinstance(state.get("failure"), (str, type(None))): + raise ValueError("Health-check failure must be a string or null") + self.state = state + + def save(self) -> None: """Persist before releasing the lock or terminating on failure.""" data = json.dumps(self.state).encode() os.lseek(self.fd, 0, os.SEEK_SET) @@ -91,7 +99,7 @@ def save(self): os.ftruncate(self.fd, len(data)) os.fsync(self.fd) - def close(self): + def close(self) -> None: if self.fd is not None: os.close(self.fd) self.fd = None diff --git a/runpod/_health/fitness.py b/runpod/_health/fitness.py index 55e9fb533..fde368471 100644 --- a/runpod/_health/fitness.py +++ b/runpod/_health/fitness.py @@ -22,7 +22,7 @@ from collections.abc import Callable from runpod._logger import RunPodLogger -from . import is_worker_process +from . import is_early_check_eligible from .coordination import ContainerChecks, CoordinationUnavailable, CoordinationBusy log = RunPodLogger() @@ -77,10 +77,24 @@ def _terminate_unhealthy(code: int = 1) -> None: "RUNPOD_GPU_BENCHMARK_TIMEOUT", "RUNPOD_GPU_TEST_TIMEOUT", "RUNPOD_GPU_MAX_ERROR_MESSAGES", + "RUNPOD_BINARY_GPU_TEST_PATH", "RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "RUNPOD_SKIP_GPU_CHECK", ) +_CHECK_CONFIG_DEPENDENCIES = { + "_memory_check": {"RUNPOD_MIN_MEMORY_GB"}, + "_disk_check": {"RUNPOD_MIN_DISK_PERCENT"}, + "_cuda_version_check": {"RUNPOD_MIN_CUDA_VERSION"}, + "_network_check": {"RUNPOD_NETWORK_CHECK_TIMEOUT"}, + "_benchmark_check": {"RUNPOD_GPU_BENCHMARK_TIMEOUT"}, + "_gpu_health_check": { + "RUNPOD_GPU_TEST_TIMEOUT", + "RUNPOD_GPU_MAX_ERROR_MESSAGES", + "RUNPOD_BINARY_GPU_TEST_PATH", + }, +} + _config_snapshot: dict[str, str | None] = {} @@ -296,42 +310,31 @@ def _refresh_late_config() -> None: from . import system system.configure() - dependencies = { - "_memory_check": {"RUNPOD_MIN_MEMORY_GB"}, - "_disk_check": {"RUNPOD_MIN_DISK_PERCENT"}, - "_cuda_version_check": {"RUNPOD_MIN_CUDA_VERSION"}, - "_network_check": {"RUNPOD_NETWORK_CHECK_TIMEOUT"}, - "_benchmark_check": {"RUNPOD_GPU_BENCHMARK_TIMEOUT"}, - "_gpu_health_check": { - "RUNPOD_GPU_TEST_TIMEOUT", - "RUNPOD_GPU_MAX_ERROR_MESSAGES", - }, - } _completed_checks[:] = [ check for check in _completed_checks if not ( getattr(check, "_runpod_builtin", False) - and dependencies.get(check.__name__, set()) & changed + and _CHECK_CONFIG_DEPENDENCIES.get(check.__name__, set()) & changed ) ] for flag, group in ( ("RUNPOD_SKIP_GPU_CHECK", "gpu_check"), ("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "system_checks"), ): - if flag in changed: - removed = [ - c - for c in _fitness_checks - if getattr(c, "_runpod_builtin", None) == group - ] - _fitness_checks[:] = [ - c for c in _fitness_checks if not any(c is old for old in removed) - ] - _completed_checks[:] = [ - c for c in _completed_checks if not any(c is old for old in removed) - ] - _registration_state[group] = False + if flag not in changed: + continue + _fitness_checks[:] = [ + check + for check in _fitness_checks + if getattr(check, "_runpod_builtin", None) != group + ] + _completed_checks[:] = [ + check + for check in _completed_checks + if getattr(check, "_runpod_builtin", None) != group + ] + _registration_state[group] = False _config_snapshot.update({v: os.environ.get(v) for v in _CONFIG_ENV_VARS}) @@ -350,6 +353,47 @@ def _fail_worker(check_name: str, exc: Exception) -> None: _terminate_unhealthy(1) +def _is_shared_check(check: Callable) -> bool: + return bool(getattr(check, "_runpod_builtin", False)) and not _is_deferred(check) + + +def _shared_check_key(check: Callable) -> str: + """Identify a built-in result by SDK version and relevant configuration.""" + from runpod.version import __version__ + + settings = { + name: os.environ.get(name) + for name in _CHECK_CONFIG_DEPENDENCIES.get(check.__name__, ()) + } + identity = [__version__, check._runpod_builtin, check.__name__, settings] + return hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() + + +async def _invoke_check(check: Callable) -> None: + if inspect.iscoroutinefunction(check): + await check() + else: + check() + + +async def _run_and_save_shared_check(check: Callable, shared: ContainerChecks) -> None: + """Save failures before exiting; save successes only after execution.""" + key = _shared_check_key(check) + if key in shared.state["passed"]: + return + try: + await _invoke_check(check) + except Exception as exc: + shared.state["failure"] = f"{check.__name__}: {type(exc).__name__}" + try: + shared.save() + finally: + _fail_worker(check.__name__, exc) + return + shared.state["passed"].append(key) + shared.save() + + async def _run_shared_checks(include_deferred: bool) -> None: """Reuse container checks across imports, including independent helpers.""" try: @@ -359,48 +403,8 @@ async def _run_shared_checks(include_deferred: bool) -> None: "early_container_check", RuntimeError(shared.state["failure"]) ) return - for check in _fitness_checks: - if not getattr(check, "_runpod_builtin", False) or _is_deferred(check): - continue - # Version and relevant settings prevent reusing incompatible results. - from runpod.version import __version__ - - dependencies = { - "_memory_check": ("RUNPOD_MIN_MEMORY_GB",), - "_disk_check": ("RUNPOD_MIN_DISK_PERCENT",), - "_cuda_version_check": ("RUNPOD_MIN_CUDA_VERSION",), - "_gpu_health_check": ( - "RUNPOD_GPU_TEST_TIMEOUT", - "RUNPOD_GPU_MAX_ERROR_MESSAGES", - "RUNPOD_BINARY_GPU_TEST_PATH", - ), - } - settings = { - k: os.environ.get(k) for k in dependencies.get(check.__name__, ()) - } - key = hashlib.sha256( - json.dumps( - [__version__, check._runpod_builtin, check.__name__, settings], - sort_keys=True, - ).encode() - ).hexdigest() - if key not in shared.state["passed"]: - try: - if inspect.iscoroutinefunction(check): - await check() - else: - check() - except Exception as exc: - shared.state["failure"] = ( - f"{check.__name__}: {type(exc).__name__}" - ) - try: - shared.save() - finally: - _fail_worker(check.__name__, exc) - return - shared.state["passed"].append(key) - shared.save() + for check in filter(_is_shared_check, _fitness_checks): + await _run_and_save_shared_check(check, shared) if not any(check is done for done in _completed_checks): _completed_checks.append(check) except CoordinationUnavailable as exc: @@ -410,49 +414,23 @@ async def _run_shared_checks(include_deferred: bool) -> None: except CoordinationBusy as exc: if include_deferred: _fail_worker("fitness_check_coordination", exc) - else: - log.warn( - "Early checks still running in another process; deferring to worker start." - ) + return + log.warn( + "Early checks still running in another process; deferring to worker start." + ) except OSError as exc: log.warn(f"Cannot save shared checks; using worker-start checks: {exc}") async def run_fitness_checks(include_deferred: bool = True) -> None: - """ - Execute all registered fitness checks sequentially at startup. - - Execution flow: - 1. Auto-register GPU check on first run (deferred to avoid circular imports) - 2. Check if registry is empty (early return if no checks) - 3. Log start of fitness check phase - 4. For each registered check: - - Auto-detect sync vs async using inspect.iscoroutinefunction() - - Execute check with timing instrumentation (await if async, call if sync) - - Log success or failure with check name and execution time - 5. On any exception: - - Log detailed error with check name, exception type, and message - - Log traceback at DEBUG level - - Force-kill the worker via os._exit(1) immediately (fail-fast). This is - a hard exit, not a cooperative sys.exit/SystemExit: it does not unwind - the stack or run cleanup, so callers cannot catch it and it cannot be - blocked by live non-daemon threads. - 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 - and handles checks with dependencies correctly. - Timing uses high-precision perf_counter for accurate measurements. - - Note: - A failing check terminates the process via os._exit(1); this function - does not return in that case and does not raise SystemExit. + """Validate startup health before accepting jobs. + + Shared built-ins reuse container results; process-specific and customer + checks run only in the final pass. Successful registrations are tracked by + identity so repeated calls skip them unless their configuration changes. + + Failed checks report unhealthy and force-exit, even with live threads. + Setup/coordination unavailability during import defers to worker start. """ if _env_flag(SKIP_FITNESS_CHECKS_ENV): log.info(f"Fitness checks disabled via {SKIP_FITNESS_CHECKS_ENV}, skipping.") @@ -471,7 +449,7 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: _fail_worker("fitness_check_setup", exc) return - if is_worker_process() and ( + if is_early_check_eligible() and ( include_deferred or not _env_flag(DEFER_FITNESS_CHECKS_ENV) ): await _run_shared_checks(include_deferred) @@ -487,11 +465,7 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: ] if not include_deferred: - pending = [ - check - for check in pending - if getattr(check, "_runpod_builtin", False) and not _is_deferred(check) - ] + pending = [check for check in pending if _is_shared_check(check)] if not pending: log.debug("No pending fitness checks, skipping.") @@ -508,11 +482,7 @@ async def run_fitness_checks(include_deferred: bool = True) -> None: log.debug(f"Executing fitness check: {check_name}") check_start_time = time.perf_counter() - # Auto-detect async vs sync using inspect - if inspect.iscoroutinefunction(check_func): - await check_func() - else: - check_func() + await _invoke_check(check_func) check_elapsed_ms = (time.perf_counter() - check_start_time) * 1000 _completed_checks.append(check_func) @@ -549,7 +519,7 @@ def run_startup_fitness_checks() -> None: if _env_flag(SKIP_FITNESS_CHECKS_ENV) or _env_flag(DEFER_FITNESS_CHECKS_ENV): return - if not is_worker_process(): + if not is_early_check_eligible(): return if _event_loop_running(): diff --git a/runpod/_startup.py b/runpod/_startup.py index be80cb939..fd69c57c9 100644 --- a/runpod/_startup.py +++ b/runpod/_startup.py @@ -2,14 +2,14 @@ import sys -from ._health import is_worker_process +from ._health import is_early_check_eligible -__all__ = ["is_worker_process", "run_import_checks"] +__all__ = ["is_early_check_eligible", "run_import_checks"] def run_import_checks() -> None: """Run shared early checks in eligible Serverless containers.""" - if not is_worker_process(): + if not is_early_check_eligible(): return try: from ._health.fitness import run_startup_fitness_checks diff --git a/tests/test_serverless/test_modules/test_fitness/test_coordination.py b/tests/test_serverless/test_modules/test_fitness/test_coordination.py index 11adaa298..9a5d652b9 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_coordination.py +++ b/tests/test_serverless/test_modules/test_fitness/test_coordination.py @@ -7,7 +7,7 @@ import pytest -from runpod._health import coordination, fitness, is_worker_process +from runpod._health import coordination, fitness, is_early_check_eligible from runpod._health.coordination import container_start_id as real_container_start_id @@ -27,7 +27,7 @@ def test_environment_gate(monkeypatch, endpoint, webhook, test, expected): monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", webhook) monkeypatch.setenv("RUNPOD_TEST", test) monkeypatch.setattr(sys, "argv", ["handler.py"]) - assert is_worker_process() is expected + assert is_early_check_eligible() is expected def process_code(tmp_path, fail=False): From 9fb1c97d0879ad57eff980f2405cf7199c4fad82 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 11 Sep 2026 10:51:07 -0400 Subject: [PATCH 12/13] Remove unused import flagged by CodeQL --- tests/test_serverless/test_modules/test_fitness/test_startup.py | 1 - 1 file changed, 1 deletion(-) 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 f3204b6ab..bdc075b7d 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_startup.py +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -1,7 +1,6 @@ """Tests for fitness checks running at import/startup time (DR-1409).""" import builtins -import os import sys import types from unittest.mock import patch From 3779487bba883b5f3bdc108120e0126b38993a44 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 11 Sep 2026 10:55:02 -0400 Subject: [PATCH 13/13] Size worker coordination wait for configured health-check timeouts --- ARCHITECTURE.md | 2 +- docs/serverless/worker_fitness_checks.md | 2 +- runpod/_health/fitness.py | 18 ++++++++- runpod/_health/gpu.py | 3 +- runpod/_health/system.py | 9 ++++- .../test_fitness/test_coordination.py | 40 +++++++++++++++++++ 6 files changed, 68 insertions(+), 6 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 99a81e8dc..8156224e1 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. The first top-level import with both `RUNPOD_ENDPOINT_ID` and `RUNPOD_WEBHOOK_GET_JOB` runs shared hardware checks, excluding test invocations and `RUNPOD_TEST`. A Linux file lock and container-start-scoped result prevent concurrent/repeated execution across processes; saved failures propagate to later workers. Identity includes host boot, PID namespace, and PID 1 start time. Network, Python CUDA initialization, compute, and custom checks remain at worker start (realtime uses serving lifespan). Unsupported/unwritable coordination defers to worker-start checks. Lock waiting is bounded at 35 seconds; a worker-start timeout fails closed. `RUNPOD_DEFER_FITNESS_CHECKS=true` postpones early checks. No custom launcher or PID environment variable is required. +1. The first top-level import with both `RUNPOD_ENDPOINT_ID` and `RUNPOD_WEBHOOK_GET_JOB` runs shared hardware checks, excluding test invocations and `RUNPOD_TEST`. A Linux file lock and container-start-scoped result prevent concurrent/repeated execution across processes; saved failures propagate to later workers. Identity includes host boot, PID namespace, and PID 1 start time. Network, Python CUDA initialization, compute, and custom checks remain at worker start (realtime uses serving lifespan). Unsupported/unwritable coordination defers to worker-start checks. Import lock waiting is bounded at 35 seconds; worker-start waiting covers the configured GPU timeout, GPU fallback, both CUDA-version probes, and a five-second overhead allowance (minimum 35 seconds). A worker-start timeout fails closed. `RUNPOD_DEFER_FITNESS_CHECKS=true` postpones early checks. No custom launcher or PID environment variable is required. 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 4af57b8d9..8edc4897b 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -49,7 +49,7 @@ A shared Linux file lock serializes these checks across Python processes. Succes Network connectivity, Python CUDA initialization, GPU compute, and customer-registered checks run in the worker process at `.start()`, before accepting jobs. Network checks retry against the worker API with a bounded budget. Keeping Python CUDA initialization out of imports protects subsequent customer forks. Production realtime mode runs its final checks in serving lifespan. -Coordination uses a fixed `/tmp` path shared by container processes. If procfs or shared state is unavailable, early checks defer to worker start. Imports wait at most 35 seconds for another checking process; a timeout defers to worker start. If the lock is still busy at worker start, the worker reports failure and exits rather than accepting jobs without validation. An owner crash releases its OS lock, allowing a later process to retry unfinished checks. Processes with separate filesystems or incompatible file permissions cannot share results and use the fallback. +Coordination uses a fixed `/tmp` path shared by container processes. If procfs or shared state is unavailable, early checks defer to worker start. Imports wait at most 35 seconds for another checking process; a timeout defers to worker start. At worker start the wait budget covers the configured GPU timeout, its fallback, both CUDA-version probes, and five seconds of overhead (minimum 35 seconds). If the lock remains busy after that budget, the worker reports failure and exits rather than accepting jobs without validation. An owner crash releases its OS lock, allowing a later process to retry unfinished checks. Processes with separate filesystems or incompatible file permissions cannot share results and use the fallback. `RUNPOD_DEFER_FITNESS_CHECKS=true` restores worker-start timing. `RUNPOD_SKIP_FITNESS_CHECKS=true` disables all checks. Set early thresholds before importing the SDK; late changes are applied at worker start, but cannot undo an earlier failure. No launcher or Docker entrypoint changes are needed. diff --git a/runpod/_health/fitness.py b/runpod/_health/fitness.py index fde368471..49367b718 100644 --- a/runpod/_health/fitness.py +++ b/runpod/_health/fitness.py @@ -394,10 +394,26 @@ async def _run_and_save_shared_check(check: Callable, shared: ContainerChecks) - shared.save() +def _coordination_wait_seconds() -> float: + """Cover sequential shared probes plus scheduling and result-write overhead.""" + budget = 5.0 + if not _env_flag("RUNPOD_SKIP_GPU_CHECK"): + from . import gpu + + budget += gpu.TIMEOUT_SECONDS + gpu.FALLBACK_TIMEOUT_SECONDS + if not _env_flag("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"): + from . import system + + # CUDA version probes nvcc, then nvidia-smi if nvcc fails. + budget += 2 * system.CUDA_VERSION_PROBE_TIMEOUT + return max(35.0, budget) + + async def _run_shared_checks(include_deferred: bool) -> None: """Reuse container checks across imports, including independent helpers.""" try: - async with ContainerChecks() as shared: + timeout = _coordination_wait_seconds() if include_deferred else 35.0 + async with ContainerChecks(timeout=timeout) as shared: if shared.state.get("failure"): _fail_worker( "early_container_check", RuntimeError(shared.state["failure"]) diff --git a/runpod/_health/gpu.py b/runpod/_health/gpu.py index fe42b387e..c40ee2922 100644 --- a/runpod/_health/gpu.py +++ b/runpod/_health/gpu.py @@ -24,6 +24,7 @@ # Defaults are safe to import; parse user settings when registering checks. TIMEOUT_SECONDS = 30 +FALLBACK_TIMEOUT_SECONDS = 10 MAX_ERROR_MESSAGES = 10 @@ -201,7 +202,7 @@ def _run_gpu_test_fallback() -> None: ["nvidia-smi", "--list-gpus"], capture_output=True, text=True, - timeout=10, + timeout=FALLBACK_TIMEOUT_SECONDS, check=False, ) diff --git a/runpod/_health/system.py b/runpod/_health/system.py index 0c47af4cc..d6e3a676c 100644 --- a/runpod/_health/system.py +++ b/runpod/_health/system.py @@ -32,6 +32,7 @@ MIN_CUDA_VERSION = "11.8" NETWORK_CHECK_TIMEOUT = 5 GPU_BENCHMARK_TIMEOUT = 2 +CUDA_VERSION_PROBE_TIMEOUT = 5 def configure() -> None: @@ -244,7 +245,9 @@ async def _get_cuda_version() -> str | None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, _ = await asyncio.wait_for(process.communicate(), timeout=5) + stdout, _ = await asyncio.wait_for( + process.communicate(), timeout=CUDA_VERSION_PROBE_TIMEOUT + ) if process.returncode == 0: output = stdout.decode("utf-8", errors="replace") for line in output.split("\n"): @@ -264,7 +267,9 @@ async def _get_cuda_version() -> str | None: stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) - stdout, _ = await asyncio.wait_for(process.communicate(), timeout=5) + stdout, _ = await asyncio.wait_for( + process.communicate(), timeout=CUDA_VERSION_PROBE_TIMEOUT + ) if process.returncode == 0: output = stdout.decode("utf-8", errors="replace") for line in output.split("\n"): diff --git a/tests/test_serverless/test_modules/test_fitness/test_coordination.py b/tests/test_serverless/test_modules/test_fitness/test_coordination.py index 9a5d652b9..f714a5d30 100644 --- a/tests/test_serverless/test_modules/test_fitness/test_coordination.py +++ b/tests/test_serverless/test_modules/test_fitness/test_coordination.py @@ -197,6 +197,9 @@ async def test_corrupt_state_uses_fallback(tmp_path, contents): @pytest.mark.asyncio async def test_lock_timeout_defers_import_but_blocks_worker(monkeypatch): class Busy: + def __init__(self, **kwargs): + pass + async def __aenter__(self): raise coordination.CoordinationBusy("still checking") @@ -208,3 +211,40 @@ async def __aexit__(self, *args): await fitness._run_shared_checks(include_deferred=False) with pytest.raises(SystemExit): await fitness._run_shared_checks(include_deferred=True) + + +@pytest.mark.asyncio +async def test_deferred_worker_wait_covers_long_gpu_check(monkeypatch): + from runpod._health import gpu, system + + monkeypatch.delenv("RUNPOD_SKIP_GPU_CHECK") + monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS") + monkeypatch.setenv("RUNPOD_DEFER_FITNESS_CHECKS", "true") + monkeypatch.setenv("RUNPOD_GPU_TEST_TIMEOUT", "60") + monkeypatch.setattr(gpu, "TIMEOUT_SECONDS", gpu.TIMEOUT_SECONDS) + monkeypatch.setattr(gpu, "MAX_ERROR_MESSAGES", gpu.MAX_ERROR_MESSAGES) + gpu.configure() + waits = [] + + class ConcurrentCheck: + def __init__(self, timeout): + waits.append(timeout) + self.state = {"passed": [], "failure": None} + + async def __aenter__(self): + # Model an owner finishing after 45 seconds without a slow test. + if waits[-1] < 45: + raise coordination.CoordinationBusy("healthy check still running") + return self + + async def __aexit__(self, *args): + pass + + monkeypatch.setattr(fitness, "ContainerChecks", ConcurrentCheck) + await fitness._run_shared_checks(include_deferred=True) + assert waits == [ + 60 + gpu.FALLBACK_TIMEOUT_SECONDS + 2 * system.CUDA_VERSION_PROBE_TIMEOUT + 5 + ] + # Imports retain a short bounded wait and defer rather than terminating. + await fitness._run_shared_checks(include_deferred=False) + assert waits[-1] == 35