diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 2dbca2ef0..8156224e1 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. The first top-level import with both `RUNPOD_ENDPOINT_ID` and `RUNPOD_WEBHOOK_GET_JOB` runs shared hardware checks, excluding test invocations and `RUNPOD_TEST`. A Linux file lock and container-start-scoped result prevent concurrent/repeated execution across processes; saved failures propagate to later workers. Identity includes host boot, PID namespace, and PID 1 start time. Network, Python CUDA initialization, compute, and custom checks remain at worker start (realtime uses serving lifespan). Unsupported/unwritable coordination defers to worker-start checks. Import lock waiting is bounded at 35 seconds; worker-start waiting covers the configured GPU timeout, GPU fallback, both CUDA-version probes, and a five-second overhead allowance (minimum 35 seconds). A worker-start timeout fails closed. `RUNPOD_DEFER_FITNESS_CHECKS=true` postpones early checks. No custom launcher or PID environment variable is required. 2. Runs only in production mode (skipped for local testing) 3. Auto-detects sync vs async using `inspect.iscoroutinefunction()` 4. Executes checks in registration order (list preserves order) -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..a273a7049 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) +- Shared hardware checks run once per container at the first Serverless import; network and process-specific checks run at worker start +- Local tests/helper imports remain exempt; network readiness and custom checks run at worker start +- Successful early checks are reused unless their configuration changes - Runs before handler initialization and job processing begins - 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..8edc4897b 100644 --- a/docs/serverless/worker_fitness_checks.md +++ b/docs/serverless/worker_fitness_checks.md @@ -41,6 +41,22 @@ if __name__ == "__main__": runpod.serverless.start({"handler": handler}) ``` +## When Checks Run + +On Serverless, the first `import runpod` runs RAM, disk, CUDA-version, and native GPU health checks. Eligibility requires both `RUNPOD_ENDPOINT_ID` and `RUNPOD_WEBHOOK_GET_JOB`. Platform tests (`RUNPOD_TEST`), `--test_input`, and local `--rp_serve_api` invocations skip early checks. + +A shared Linux file lock serializes these checks across Python processes. Successful results are reused by helpers and the actual worker; failures are saved before reporting unhealthy and exiting, so another process cannot silently ignore the failure. Results are scoped to the host boot, PID namespace, and container init process start time, rather than just the Pod ID. Checks with changed settings or SDK version are rerun. The first import may occur after model loading; no earlier timing is guaranteed in that case. + +Network connectivity, Python CUDA initialization, GPU compute, and customer-registered checks run in the worker process at `.start()`, before accepting jobs. Network checks retry against the worker API with a bounded budget. Keeping Python CUDA initialization out of imports protects subsequent customer forks. Production realtime mode runs its final checks in serving lifespan. + +Coordination uses a fixed `/tmp` path shared by container processes. If procfs or shared state is unavailable, early checks defer to worker start. Imports wait at most 35 seconds for another checking process; a timeout defers to worker start. At worker start the wait budget covers the configured GPU timeout, its fallback, both CUDA-version probes, and five seconds of overhead (minimum 35 seconds). If the lock remains busy after that budget, the worker reports failure and exits rather than accepting jobs without validation. An owner crash releases its OS lock, allowing a later process to retry unfinished checks. Processes with separate filesystems or incompatible file permissions cannot share results and use the fallback. + +`RUNPOD_DEFER_FITNESS_CHECKS=true` restores worker-start timing. `RUNPOD_SKIP_FITNESS_CHECKS=true` disables all checks. Set early thresholds before importing the SDK; late changes are applied at worker start, but cannot undo an earlier failure. No launcher or Docker entrypoint changes are needed. + +### Rollout + +Validate in a small set of workers before broader rollout. The deferral variable provides a rollback of early timing without handler edits. This SDK change does not itself alter deployed platform configuration. + ## Async Fitness Checks Fitness checks support both synchronous and asynchronous functions: @@ -284,19 +300,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 +357,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 +385,15 @@ ENV RUNPOD_NETWORK_CHECK_TIMEOUT=10 ENV RUNPOD_GPU_BENCHMARK_TIMEOUT=2 ``` -Or in Python: +For 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 +404,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 +415,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 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 in eligible Serverless containers; the final pass runs before job processing. Successful checks are reused unless their configuration changes. - They run **before the first job is processed** - They run **only on the actual Runpod serverless platform** - Local development and testing modes skip fitness checks @@ -555,7 +577,7 @@ async def check_api_with_retry(): ## Testing -When developing locally, fitness checks don't run. To test them, you can manually invoke the runner: +When developing locally, fitness checks don't run. To test them, you can manually invoke the runner. Note that each check runs once per process: a second `run_fitness_checks()` call skips checks that already passed, so call `clear_fitness_checks()` (as below) between runs: ```python import asyncio diff --git a/runpod/__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..e8c1d713b --- /dev/null +++ b/runpod/_health/__init__.py @@ -0,0 +1,22 @@ +"""Lightweight Serverless environment detection; no SDK imports.""" + +import os +import sys + + +def is_serverless_environment() -> bool: + """Recognize production worker configuration, excluding platform tests.""" + return ( + bool(os.environ.get("RUNPOD_ENDPOINT_ID", "").strip()) + and bool(os.environ.get("RUNPOD_WEBHOOK_GET_JOB", "").strip()) + and os.environ.get("RUNPOD_TEST", "").strip().lower() + not in ("1", "true", "yes", "on") + ) + + +def is_early_check_eligible() -> bool: + """Eligibility for shared early checks, not an assertion of process identity.""" + return is_serverless_environment() and not any( + arg.split("=", 1)[0] in ("--test_input", "--rp_serve_api") + for arg in sys.argv[1:] + ) diff --git a/runpod/_health/coordination.py b/runpod/_health/coordination.py new file mode 100644 index 000000000..2f4725214 --- /dev/null +++ b/runpod/_health/coordination.py @@ -0,0 +1,108 @@ +"""Linux container-start-scoped coordination for shared health checks.""" + +import asyncio +import hashlib +import json +import os +import time +from pathlib import Path + + +class CoordinationUnavailable(Exception): + """Shared state cannot be used; worker-start checks remain available.""" + + +class CoordinationBusy(Exception): + """Another process did not finish within the bounded wait.""" + + +def container_start_id() -> str: + """PID namespace + init start ticks + host boot distinguish container restarts. + + Use fixed /tmp rather than TMPDIR (which can differ between processes). + No pod-id-only marker: container files can survive a restart. + """ + boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip() + stat = Path("/proc/1/stat").read_text() + start_ticks = stat.rsplit(")", 1)[1].split()[19] + namespace = os.readlink("/proc/1/ns/pid") + return hashlib.sha256(f"{boot}:{namespace}:{start_ticks}".encode()).hexdigest() + + +class ContainerChecks: + """Hold one flock while reading, executing, and recording early checks.""" + + def __init__(self, timeout: float = 35): + self.timeout = timeout + self.fd = None + self.state = {"passed": [], "failure": None} + + async def __aenter__(self): + try: + import fcntl + + identity = container_start_id() + path = f"/tmp/runpod-fitness-{identity}.json" + self.fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600) + os.set_inheritable(self.fd, False) + except (OSError, ValueError, IndexError, ImportError) as exc: + self.close() + raise CoordinationUnavailable(str(exc)) from exc + try: + await self._acquire_lock(fcntl) + self._load_state() + return self + except (OSError, ValueError) as exc: + self.close() + raise CoordinationUnavailable(str(exc)) from exc + except BaseException: + self.close() + raise + + async def _acquire_lock(self, fcntl) -> None: + deadline = time.monotonic() + self.timeout + while True: + try: + fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return + except BlockingIOError: + if time.monotonic() >= deadline: + raise CoordinationBusy("Timed out waiting for early health checks") + await asyncio.sleep(0.05) + + def _load_state(self) -> None: + raw = os.read(self.fd, 65536) + if not raw: + return + state = json.loads(raw) + if not isinstance(state, dict): + raise ValueError("Health-check state must be an object") + passed = state.get("passed") + if not isinstance(passed, list) or not all( + isinstance(key, str) for key in passed + ): + raise ValueError("Passed health checks must be a list of cache keys") + if not isinstance(state.get("failure"), (str, type(None))): + raise ValueError("Health-check failure must be a string or null") + self.state = state + + def save(self) -> None: + """Persist before releasing the lock or terminating on failure.""" + data = json.dumps(self.state).encode() + os.lseek(self.fd, 0, os.SEEK_SET) + remaining = memoryview(data) + while remaining: + written = os.write(self.fd, remaining) + if written <= 0: + raise OSError("Unable to persist health-check state") + remaining = remaining[written:] + os.ftruncate(self.fd, len(data)) + os.fsync(self.fd) + + def close(self) -> None: + if self.fd is not None: + os.close(self.fd) + self.fd = None + + async def __aexit__(self, *args): + self.close() diff --git a/runpod/_health/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..49367b718 --- /dev/null +++ b/runpod/_health/fitness.py @@ -0,0 +1,559 @@ +""" +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 hashlib +import json +import os +import sys +import time +import traceback +from collections.abc import Callable + +from runpod._logger import RunPodLogger +from . import is_early_check_eligible +from .coordination import ContainerChecks, CoordinationUnavailable, CoordinationBusy + +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" + +# 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_BINARY_GPU_TEST_PATH", + "RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", + "RUNPOD_SKIP_GPU_CHECK", +) + +_CHECK_CONFIG_DEPENDENCIES = { + "_memory_check": {"RUNPOD_MIN_MEMORY_GB"}, + "_disk_check": {"RUNPOD_MIN_DISK_PERCENT"}, + "_cuda_version_check": {"RUNPOD_MIN_CUDA_VERSION"}, + "_network_check": {"RUNPOD_NETWORK_CHECK_TIMEOUT"}, + "_benchmark_check": {"RUNPOD_GPU_BENCHMARK_TIMEOUT"}, + "_gpu_health_check": { + "RUNPOD_GPU_TEST_TIMEOUT", + "RUNPOD_GPU_MAX_ERROR_MESSAGES", + "RUNPOD_BINARY_GPU_TEST_PATH", + }, +} + +_config_snapshot: dict[str, str | None] = {} + + +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() + _completed_checks[:] = [ + check + for check in _completed_checks + if not ( + getattr(check, "_runpod_builtin", False) + and _CHECK_CONFIG_DEPENDENCIES.get(check.__name__, set()) & changed + ) + ] + for flag, group in ( + ("RUNPOD_SKIP_GPU_CHECK", "gpu_check"), + ("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "system_checks"), + ): + if flag not in changed: + continue + _fitness_checks[:] = [ + check + for check in _fitness_checks + if getattr(check, "_runpod_builtin", None) != group + ] + _completed_checks[:] = [ + check + for check in _completed_checks + if getattr(check, "_runpod_builtin", None) != group + ] + _registration_state[group] = False + _config_snapshot.update({v: os.environ.get(v) for v in _CONFIG_ENV_VARS}) + + +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) + + +def _is_shared_check(check: Callable) -> bool: + return bool(getattr(check, "_runpod_builtin", False)) and not _is_deferred(check) + + +def _shared_check_key(check: Callable) -> str: + """Identify a built-in result by SDK version and relevant configuration.""" + from runpod.version import __version__ + + settings = { + name: os.environ.get(name) + for name in _CHECK_CONFIG_DEPENDENCIES.get(check.__name__, ()) + } + identity = [__version__, check._runpod_builtin, check.__name__, settings] + return hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() + + +async def _invoke_check(check: Callable) -> None: + if inspect.iscoroutinefunction(check): + await check() + else: + check() + + +async def _run_and_save_shared_check(check: Callable, shared: ContainerChecks) -> None: + """Save failures before exiting; save successes only after execution.""" + key = _shared_check_key(check) + if key in shared.state["passed"]: + return + try: + await _invoke_check(check) + except Exception as exc: + shared.state["failure"] = f"{check.__name__}: {type(exc).__name__}" + try: + shared.save() + finally: + _fail_worker(check.__name__, exc) + return + shared.state["passed"].append(key) + shared.save() + + +def _coordination_wait_seconds() -> float: + """Cover sequential shared probes plus scheduling and result-write overhead.""" + budget = 5.0 + if not _env_flag("RUNPOD_SKIP_GPU_CHECK"): + from . import gpu + + budget += gpu.TIMEOUT_SECONDS + gpu.FALLBACK_TIMEOUT_SECONDS + if not _env_flag("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"): + from . import system + + # CUDA version probes nvcc, then nvidia-smi if nvcc fails. + budget += 2 * system.CUDA_VERSION_PROBE_TIMEOUT + return max(35.0, budget) + + +async def _run_shared_checks(include_deferred: bool) -> None: + """Reuse container checks across imports, including independent helpers.""" + try: + timeout = _coordination_wait_seconds() if include_deferred else 35.0 + async with ContainerChecks(timeout=timeout) as shared: + if shared.state.get("failure"): + _fail_worker( + "early_container_check", RuntimeError(shared.state["failure"]) + ) + return + for check in filter(_is_shared_check, _fitness_checks): + await _run_and_save_shared_check(check, shared) + if not any(check is done for done in _completed_checks): + _completed_checks.append(check) + except CoordinationUnavailable as exc: + log.warn( + f"Early check coordination unavailable; using worker-start checks: {exc}" + ) + except CoordinationBusy as exc: + if include_deferred: + _fail_worker("fitness_check_coordination", exc) + return + log.warn( + "Early checks still running in another process; deferring to worker start." + ) + except OSError as exc: + log.warn(f"Cannot save shared checks; using worker-start checks: {exc}") + + +async def run_fitness_checks(include_deferred: bool = True) -> None: + """Validate startup health before accepting jobs. + + Shared built-ins reuse container results; process-specific and customer + checks run only in the final pass. Successful registrations are tracked by + identity so repeated calls skip them unless their configuration changes. + + Failed checks report unhealthy and force-exit, even with live threads. + Setup/coordination unavailability during import defers to worker start. + """ + if _env_flag(SKIP_FITNESS_CHECKS_ENV): + log.info(f"Fitness checks disabled via {SKIP_FITNESS_CHECKS_ENV}, skipping.") + 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 + + if is_early_check_eligible() and ( + include_deferred or not _env_flag(DEFER_FITNESS_CHECKS_ENV) + ): + await _run_shared_checks(include_deferred) + if not include_deferred: + return + + # Identity, not equality: two distinct registrations may compare equal + # (e.g. fresh bound-method objects of one method), and `==` would skip one. + pending = [ + 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 _is_shared_check(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() + + await _invoke_check(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. + + Shared built-ins run once per container startup. Child processes reuse the + result; process-specific and customer checks wait for worker start. + """ + if _env_flag(SKIP_FITNESS_CHECKS_ENV) or _env_flag(DEFER_FITNESS_CHECKS_ENV): + return + + if not is_early_check_eligible(): + return + + 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..c40ee2922 --- /dev/null +++ b/runpod/_health/gpu.py @@ -0,0 +1,328 @@ +""" +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 +FALLBACK_TIMEOUT_SECONDS = 10 +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=FALLBACK_TIMEOUT_SECONDS, + 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..d6e3a676c --- /dev/null +++ b/runpod/_health/system.py @@ -0,0 +1,560 @@ +""" +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 +CUDA_VERSION_PROBE_TIMEOUT = 5 + + +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=CUDA_VERSION_PROBE_TIMEOUT + ) + 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=CUDA_VERSION_PROBE_TIMEOUT + ) + 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..fd69c57c9 --- /dev/null +++ b/runpod/_startup.py @@ -0,0 +1,24 @@ +"""Container-scoped startup gate; safe to import without loading serverless.""" + +import sys + +from ._health import is_early_check_eligible + +__all__ = ["is_early_check_eligible", "run_import_checks"] + + +def run_import_checks() -> None: + """Run shared early checks in eligible Serverless containers.""" + if not is_early_check_eligible(): + 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/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..42e801896 100644 --- a/tests/test_serverless/test_modules/test_fitness/conftest.py +++ b/tests/test_serverless/test_modules/test_fitness/conftest.py @@ -10,7 +10,7 @@ @pytest.fixture(autouse=True) -def cleanup_fitness_checks(monkeypatch): +def cleanup_fitness_checks(monkeypatch, tmp_path): """Automatically clean up fitness checks before and after each test. Disables auto-registration of system checks to avoid interference @@ -21,6 +21,15 @@ def cleanup_fitness_checks(monkeypatch): to raise SystemExit(1) so tests can assert exit behavior in-process. Tests that need the real os._exit patch it themselves. """ + from runpod._health import coordination + + monkeypatch.setattr( + coordination, + "container_start_id", + lambda: tmp_path.parent.name + "-" + tmp_path.name, + ) + monkeypatch.delenv("RUNPOD_ENDPOINT_ID", raising=False) + monkeypatch.delenv("RUNPOD_TEST", raising=False) monkeypatch.setenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "true") monkeypatch.setenv("RUNPOD_SKIP_GPU_CHECK", "true") @@ -31,6 +40,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_coordination.py b/tests/test_serverless/test_modules/test_fitness/test_coordination.py new file mode 100644 index 000000000..f714a5d30 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_coordination.py @@ -0,0 +1,250 @@ +"""Real interprocess locks/results, and container-start identity regressions.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from runpod._health import coordination, fitness, is_early_check_eligible +from runpod._health.coordination import container_start_id as real_container_start_id + + +@pytest.mark.parametrize( + "endpoint,webhook,test,expected", + [ + ("", "", "", False), + ("ep", "", "", False), + ("", "https://example.test/job", "", False), + ("ep", "https://example.test/job", "", True), + ("ep", "https://example.test/job", "TRUE", False), + ("ep", "https://example.test/job", "1", False), + ], +) +def test_environment_gate(monkeypatch, endpoint, webhook, test, expected): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", endpoint) + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", webhook) + monkeypatch.setenv("RUNPOD_TEST", test) + monkeypatch.setattr(sys, "argv", ["handler.py"]) + assert is_early_check_eligible() is expected + + +def process_code(tmp_path, fail=False): + return f""" +import asyncio, os, time +from pathlib import Path +from runpod._health import fitness, coordination +coordination.container_start_id = lambda: {(tmp_path.parent.name + "-" + tmp_path.name)!r} +os.environ['RUNPOD_ENDPOINT_ID'] = 'ep' +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +fitness._report_unhealthy = lambda *args: None +@fitness.register_fitness_check +def shared_check(): + with open({str(tmp_path / "calls")!r}, 'a') as out: + out.write('check\\n') + time.sleep(0.2) + if {fail!r}: + raise RuntimeError('failed') +shared_check._runpod_builtin = 'system_checks' +asyncio.run(fitness.run_fitness_checks(include_deferred=False)) +asyncio.run(fitness.run_fitness_checks()) +""" + + +def launch(code): + env = {k: v for k, v in os.environ.items() if not k.startswith("RUNPOD_")} + env.update(RUNPOD_SKIP_GPU_CHECK="true", RUNPOD_SKIP_AUTO_SYSTEM_CHECKS="true") + return subprocess.Popen( + [sys.executable, "-c", code], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def test_independent_processes_share_success(tmp_path): + children = [launch(process_code(tmp_path)) for _ in range(3)] + for child in children: + out, err = child.communicate(timeout=10) + assert child.returncode == 0, out + err + assert (tmp_path / "calls").read_text().splitlines() == ["check"] + + +def test_failure_propagates_without_rerunning(tmp_path): + for _ in range(2): + child = launch(process_code(tmp_path, fail=True)) + out, err = child.communicate(timeout=10) + assert child.returncode == 1, out + err + assert (tmp_path / "calls").read_text().splitlines() == ["check"] + + +@pytest.mark.asyncio +async def test_busy_is_bounded_and_owner_crash_releases_lock(tmp_path): + code = f""" +import asyncio, time +from runpod._health import coordination +coordination.container_start_id = lambda: {(tmp_path.parent.name + "-" + tmp_path.name)!r} +async def main(): + async with coordination.ContainerChecks(): + print('LOCKED', flush=True) + time.sleep(30) +asyncio.run(main()) +""" + child = launch(code) + try: + assert child.stdout.readline().strip() == "LOCKED" + with pytest.raises(coordination.CoordinationBusy): + async with coordination.ContainerChecks(timeout=0.1): + pass + finally: + child.kill() + child.communicate(timeout=5) + async with coordination.ContainerChecks(timeout=0.1) as shared: + assert shared.state["passed"] == [] + + +@pytest.mark.asyncio +async def test_restart_does_not_reuse_previous_success(monkeypatch, tmp_path): + async with coordination.ContainerChecks() as shared: + shared.state["passed"] = ["old-success"] + shared.save() + monkeypatch.setattr( + coordination, + "container_start_id", + lambda: (tmp_path.parent.name + "-" + tmp_path.name) + "-restart", + ) + async with coordination.ContainerChecks() as shared: + assert shared.state["passed"] == [] + + +def test_identity_changes_when_init_restarts(monkeypatch): + # Fields following comm start at field 3; starttime is field 22. + ticks = ["100"] + + def read(path): + if str(path).endswith("boot_id"): + return "host-boot" + return "1 (name with spaces) " + " ".join(["0"] * 19 + ticks) + + monkeypatch.setattr(Path, "read_text", read) + monkeypatch.setattr(os, "readlink", lambda path: "pid:[123]") + first = real_container_start_id() + ticks[0] = "200" + assert real_container_start_id() != first + + +@pytest.mark.asyncio +async def test_unavailable_coordination_defers_then_checks_at_start(monkeypatch): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + + def unavailable(): + raise OSError("read-only") + + monkeypatch.setattr(coordination, "container_start_id", unavailable) + calls = [] + + @fitness.register_fitness_check + def early(): + calls.append("early") + + early._runpod_builtin = "system_checks" + + @fitness.register_fitness_check + def customer(): + calls.append("customer") + + await fitness.run_fitness_checks(include_deferred=False) + assert calls == [] + await fitness.run_fitness_checks() + assert calls == ["early", "customer"] + + +@pytest.mark.asyncio +async def test_customer_and_process_checks_only_at_start(monkeypatch): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.test/job") + calls = [] + + @fitness.register_fitness_check + def customer(): + calls.append("customer") + + @fitness.register_fitness_check + @fitness.defer_to_worker_start + def cuda(): + calls.append("cuda") + + cuda._runpod_builtin = "system_checks" + await fitness.run_fitness_checks(include_deferred=False) + assert calls == [] + await fitness.run_fitness_checks() + assert calls == ["customer", "cuda"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("contents", ["[]", '{"passed": null}', "{broken"]) +async def test_corrupt_state_uses_fallback(tmp_path, contents): + identity = coordination.container_start_id() + Path(f"/tmp/runpod-fitness-{identity}.json").write_text(contents) + with pytest.raises(coordination.CoordinationUnavailable): + async with coordination.ContainerChecks(): + pass + + +@pytest.mark.asyncio +async def test_lock_timeout_defers_import_but_blocks_worker(monkeypatch): + class Busy: + def __init__(self, **kwargs): + pass + + async def __aenter__(self): + raise coordination.CoordinationBusy("still checking") + + async def __aexit__(self, *args): + pass + + monkeypatch.setattr(fitness, "ContainerChecks", Busy) + monkeypatch.setattr(fitness, "_report_unhealthy", lambda *args: None) + await fitness._run_shared_checks(include_deferred=False) + with pytest.raises(SystemExit): + await fitness._run_shared_checks(include_deferred=True) + + +@pytest.mark.asyncio +async def test_deferred_worker_wait_covers_long_gpu_check(monkeypatch): + from runpod._health import gpu, system + + monkeypatch.delenv("RUNPOD_SKIP_GPU_CHECK") + monkeypatch.delenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS") + monkeypatch.setenv("RUNPOD_DEFER_FITNESS_CHECKS", "true") + monkeypatch.setenv("RUNPOD_GPU_TEST_TIMEOUT", "60") + monkeypatch.setattr(gpu, "TIMEOUT_SECONDS", gpu.TIMEOUT_SECONDS) + monkeypatch.setattr(gpu, "MAX_ERROR_MESSAGES", gpu.MAX_ERROR_MESSAGES) + gpu.configure() + waits = [] + + class ConcurrentCheck: + def __init__(self, timeout): + waits.append(timeout) + self.state = {"passed": [], "failure": None} + + async def __aenter__(self): + # Model an owner finishing after 45 seconds without a slow test. + if waits[-1] < 45: + raise coordination.CoordinationBusy("healthy check still running") + return self + + async def __aexit__(self, *args): + pass + + monkeypatch.setattr(fitness, "ContainerChecks", ConcurrentCheck) + await fitness._run_shared_checks(include_deferred=True) + assert waits == [ + 60 + gpu.FALLBACK_TIMEOUT_SECONDS + 2 * system.CUDA_VERSION_PROBE_TIMEOUT + 5 + ] + # Imports retain a short bounded wait and defer rather than terminating. + await fitness._run_shared_checks(include_deferred=False) + assert waits[-1] == 35 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..905b0df74 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_safe_startup.py @@ -0,0 +1,234 @@ +"""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, coordination + +coordination.container_start_id = lambda: "lazy-test-" + str(os.getpid()) +from runpod._startup import 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("RUNPOD_ENDPOINT_ID", "endpoint") + else: + monkeypatch.delenv("RUNPOD_ENDPOINT_ID", raising=False) + 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("RUNPOD_ENDPOINT_ID", "endpoint") + with patch.object(fitness, "_refresh_late_config") as refresh: + run_import_checks() + refresh.assert_not_called() + + +@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("RUNPOD_ENDPOINT_ID", "endpoint") + 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("RUNPOD_ENDPOINT_ID", "endpoint") + 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 + + +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, coordination +coordination.container_start_id = lambda: 'lazy-test-' + str(os.getpid()) +from runpod._startup import run_import_checks +os.environ['RUNPOD_WEBHOOK_GET_JOB'] = 'https://example.test/job' +os.environ['RUNPOD_ENDPOINT_ID'] = 'endpoint' +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 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..bdc075b7d --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_startup.py @@ -0,0 +1,345 @@ +"""Tests for fitness checks running at import/startup time (DR-1409).""" + +import builtins +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, +) + + +def register_early_check(func): + """Synthetic built-in for exercising the early runner.""" + func._runpod_builtin = "system_checks" + return register_fitness_check(func) + + +@pytest.fixture() +def worker_env(monkeypatch): + """Make the process look like a real Runpod worker.""" + monkeypatch.setenv("RUNPOD_WEBHOOK_GET_JOB", "https://example.com/job") + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint") + 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_early_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_early_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_early_check + def first(): + calls.append("first") + + await run_fitness_checks() + + @register_early_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_early_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_early_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_early_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_early_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_early_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_early_check + def early(): + calls.append("early") + + @register_early_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_early_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 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_early_check + def check(): + calls.append(True) + + @register_early_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")