diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 2dbca2ef0..998eab63c 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -599,12 +599,12 @@ 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 - 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**: @@ -613,11 +613,11 @@ 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. `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, call `sys.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 @@ -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..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 -- Checks run only once at worker startup (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 c50a9c932..d32ed76f1 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -41,6 +41,37 @@ if __name__ == "__main__": runpod.serverless.start({"handler": handler}) ``` +## When Checks Run + +Early checks are automatic when the handler is launched with `runpod-worker`: + +```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. + +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. + +### Compatibility and failure handling + +- 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. + +### 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 Fitness checks support both synchronous and asynchronous functions: @@ -284,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. +Tests TCP reachability of the worker API host at worker start. -- **Default**: 5 second timeout to 8.8.8.8:53 -- **Configure**: `RUNPOD_NETWORK_CHECK_TIMEOUT=10` - -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) @@ -343,15 +372,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: ``` @@ -371,13 +400,15 @@ ENV RUNPOD_NETWORK_CHECK_TIMEOUT=10 ENV RUNPOD_GPU_BENCHMARK_TIMEOUT=2 ``` -Or in Python: +For legacy/deferred launches, settings can also be configured in Python before worker start: ```python import os os.environ["RUNPOD_MIN_MEMORY_GB"] = "8.0" os.environ["RUNPOD_MIN_DISK_PERCENT"] = "15.0" + +import runpod ``` ### Disabling Built-in Checks @@ -388,6 +419,8 @@ For testing or specialized deployments, built-in checks can be disabled via envi |---|---| | `RUNPOD_SKIP_AUTO_SYSTEM_CHECKS=true` | Skips auto-registration of memory, disk, network, CUDA version, CUDA init, and GPU benchmark checks | | `RUNPOD_SKIP_GPU_CHECK=true` | Skips auto-registration of the native GPU memory allocation test (`gpu_test` binary) | +| `RUNPOD_SKIP_FITNESS_CHECKS=true` | Skips every fitness check, built-in **and** user-registered | +| `RUNPOD_DEFER_FITNESS_CHECKS=true` | Keeps the checks but runs them only at `runpod.serverless.start()`, not at import | ```python import os @@ -397,15 +430,19 @@ os.environ["RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"] = "true" # Disable the automatic GPU memory allocation test os.environ["RUNPOD_SKIP_GPU_CHECK"] = "true" + +import runpod ``` -User-registered checks via `@register_fitness_check` still run regardless of these flags. +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. ## Behavior ### Execution Timing -- Fitness checks run **only once at worker startup** +- 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 @@ -555,7 +592,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/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..0701b8265 --- /dev/null +++ b/runpod/_health/__init__.py @@ -0,0 +1,20 @@ +"""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/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..47159ccda --- /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 . 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..b4196a284 --- /dev/null +++ b/runpod/_logger.py @@ -0,0 +1,161 @@ +""" +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, 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"{name}: [REDACTED]") + + 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..f6d312626 --- /dev/null +++ b/runpod/_startup.py @@ -0,0 +1,24 @@ +"""Process-scoped startup gate; safe to import without loading serverless.""" + +import sys + +from ._health import WORKER_PID_ENV, is_worker_process + +__all__ = ["WORKER_PID_ENV", "is_worker_process", "run_import_checks"] + + +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 052452073..f1ab29de5 100644 --- a/runpod/serverless/__init__.py +++ b/runpod/serverless/__init__.py @@ -30,6 +30,7 @@ log = RunPodLogger() + # ---------------------------------------------------------------------------- # # Run Time Arguments # # ---------------------------------------------------------------------------- # 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 77df97e79..c9e8d5343 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -1,296 +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 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] = [] - - -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() - - -_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 - - _registration_state["gpu_check"] = True - - 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") - - -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. - """ - 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": - log.debug( - "System fitness checks disabled via environment (RUNPOD_SKIP_AUTO_SYSTEM_CHECKS)" - ) - _registration_state["system_checks"] = True - return - - _registration_state["system_checks"] = True - - 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") - - -async def run_fitness_checks() -> 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 - - 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. - """ - # Defer GPU check auto-registration until fitness checks are about to run - # This avoids circular import issues during module initialization - _ensure_gpu_check_registered() - - # Defer system check auto-registration until fitness checks are about to run - _ensure_system_checks_registered() - - if not _fitness_checks: - log.debug("No fitness checks registered, skipping.") - return - - log.info(f"Running {len(_fitness_checks)} fitness check(s)...") - - total_start_time = time.perf_counter() - - for check_func in _fitness_checks: - 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 - 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)") +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 bae74cd88..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 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. - - This function is called during rp_fitness module initialization. - 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 - """ - # Allow skipping during tests - if os.environ.get("RUNPOD_SKIP_GPU_CHECK", "").lower() == "true": - 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 8dc8946d9..f531d43fc 100644 --- a/runpod/serverless/modules/rp_system_fitness.py +++ b/runpod/serverless/modules/rp_system_fitness.py @@ -1,511 +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 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. - """ - 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 - async def _cuda_init_check() -> None: - """CUDA device initialization check.""" - await _check_cuda_initialization() - - @register_fitness_check - 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 028c7ebcd..f561f993f 100644 --- a/runpod/serverless/utils/rp_cuda.py +++ b/runpod/serverless/utils/rp_cuda.py @@ -1,18 +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: - output = subprocess.check_output(["nvidia-smi"], stderr=subprocess.DEVNULL) - 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/conftest.py b/tests/test_serverless/test_modules/test_fitness/conftest.py index f8df86144..04a6f9e68 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) @@ -31,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_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 new file mode 100644 index 000000000..4beba6c82 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -0,0 +1,361 @@ +"""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 + +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.setenv("RUNPOD_FITNESS_WORKER_PID", str(os.getpid())) + 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"] + + @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): + 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): + # 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 + 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 == [] + + +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"]) + + +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.""" + + @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._health.gpu", fake_gpu_module + ) + monkeypatch.setitem( + sys.modules, + "runpod._health.system", + 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_top_level_import_calls_startup_checks(self, worker_env, monkeypatch): + import importlib + + import runpod.serverless + + calls = [] + monkeypatch.setattr( + rp_fitness, "run_startup_fitness_checks", lambda: calls.append(True) + ) + + importlib.reload(runpod) + + 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") + # 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(SystemExit) as exc: + await run_fitness_checks() + assert exc.value.code == 1 + + +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() diff --git a/tests/test_serverless/test_modules/test_logger.py b/tests/test_serverless/test_modules/test_logger.py index ea428379c..c63c360d8 100644 --- a/tests/test_serverless/test_modules/test_logger.py +++ b/tests/test_serverless/test_modules/test_logger.py @@ -105,9 +105,31 @@ 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_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. 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")