feat(serverless): add safe early fitness checks and skip controls - #578
feat(serverless): add safe early fitness checks and skip controls#578justinwlin wants to merge 9 commits into
Conversation
Built-in GPU/system fitness checks ran in run_worker, which a handler module only reaches after loading its model. Run them when runpod.serverless is imported instead, so a broken environment fails in seconds. User-registered checks still run at start(); checks that already passed are not repeated. Adds RUNPOD_SKIP_FITNESS_CHECKS to disable all checks and RUNPOD_DEFER_FITNESS_CHECKS to restore the previous start()-only timing.
_cuda_init_check and _benchmark_check import torch and allocate on the device. Running them at import would leave a CUDA context in a process the handler may later fork, which CUDA does not support and vLLM/DeepSpeed trip over. Mark them @defer_to_worker_start so only subprocess-based and non-GPU checks run early.
- run startup pass on a dedicated event loop instead of asyncio.run, which resets the loop policy and breaks asyncio.get_event_loop() in handler code on Python 3.10+ - set RUNPOD_FITNESS_CHECKS_DONE after the startup pass so children re-importing this module under multiprocessing 'spawn' skip the checks - latch check auto-registration state only on success, so a malformed RUNPOD_MIN_*/GPU timeout value re-raises loudly in run_worker instead of silently disabling all system checks - compare completed checks by identity, not equality, so distinct registrations that compare equal (bound methods) are not skipped - bound the nvidia-smi call in rp_cuda.is_available with a 5s timeout - accept 1/true/yes/on for RUNPOD_SKIP_GPU_CHECK and RUNPOD_SKIP_AUTO_SYSTEM_CHECKS, matching the new flags - tests: pin the worker.py and import-time wiring, the full defer behavior, the done marker, the real auto-registration path (guard: no torch import), and bound-method re-registration; fix an orphaned coroutine in test_unexpected_error_does_not_propagate - docs: thresholds/skip flags must be set before import runpod, realtime API mode runs only the import-time checks, refresh stale ARCHITECTURE.md execution flow
…touch-ups - regression test: malformed RUNPOD_MIN_* must re-raise in run_worker, never fail open (latch-on-success) - fix dormant called/calls typo in the done-marker test - README: checks run once per check, not once at startup - ARCHITECTURE.md: failure path is os._exit(1), not sys.exit(1) - docs: GPU benchmark default timeout is 2s, not 100ms - rp_gpu_fitness docstring: lazy registration + truthy flag values
The import-time pass consumes RUNPOD_MIN_*/RUNPOD_SKIP_*/RUNPOD_GPU_* at import; values set from the handler afterwards were silently ignored. run_fitness_checks now diffs the current env against the values snapshot at the startup pass and warns with the exact fix (set before import, or RUNPOD_DEFER_FITNESS_CHECKS=true).
There was a problem hiding this comment.
🟡 Changes recommended
Auto-registration failures can currently escape the hard-exit failure path, undermining the “must not hang” operational guarantee during worker startup.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Moves most serverless worker fitness checks earlier (at import runpod.serverless) to fail unhealthy workers before model load, while keeping CUDA-context-creating checks deferred to start() and adding env flags to skip/defer behavior.
Changes:
- Add an import-time startup pass (
run_startup_fitness_checks) gated by worker env, with once-per-process deduping and deferred-check support. - Introduce global skip/defer env flags and more consistent “truthy” env parsing; add
nvidia-smitimeout for CUDA detection. - Add/expand tests and docs to cover startup timing, deduping, deferred checks, and late-config warnings.
File summaries
| File | Description |
|---|---|
| tests/test_serverless/test_worker.py | Asserts worker loop still runs fitness checks. |
| tests/test_serverless/test_utils/test_cuda.py | Updates CUDA availability test expectations for timeout=5. |
| tests/test_serverless/test_modules/test_fitness/test_startup.py | New test suite validating import/start timing, deferral, dedupe, and config warnings. |
| tests/test_serverless/test_modules/test_fitness/conftest.py | Resets new startup/dedupe global state between tests. |
| runpod/serverless/utils/rp_cuda.py | Adds bounded nvidia-smi probe with timeout to avoid hangs. |
| runpod/serverless/modules/rp_system_fitness.py | Marks CUDA-init and benchmark checks as deferred-to-worker-start. |
| runpod/serverless/modules/rp_gpu_fitness.py | Uses shared truthy env flag parsing for skip behavior. |
| runpod/serverless/modules/rp_fitness.py | Implements startup pass, deduping, skip/defer env flags, and late-config warnings. |
| runpod/serverless/init.py | Triggers startup checks at import (worker-only no-op otherwise). |
| README.md | Updates high-level behavior summary for new timing model. |
| docs/serverless/worker_fitness_checks.md | Documents import-time checks, deferral, and new env flags. |
| ARCHITECTURE.md | Updates architecture docs for new timing and hard-exit behavior. |
Review details
- Files reviewed: 12/12 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # 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() |
deanq
left a comment
deanq
left a comment
There was a problem hiding this comment.
Review from /code-review (correctness + cleanup pass). Four findings, most centered on moving os._exit(1)-capable checks to import time. Lines re-anchored to the current diff.
| if _env_flag(SKIP_FITNESS_CHECKS_ENV) or _env_flag(DEFER_FITNESS_CHECKS_ENV): | ||
| return | ||
|
|
||
| if not os.environ.get("RUNPOD_WEBHOOK_GET_JOB"): |
There was a problem hiding this comment.
The import-time gate keys only on RUNPOD_WEBHOOK_GET_JOB, but that misses the _is_local / test_input local-mode guard that previously protected these checks. A handler run with rp_args.test_input used to skip all fitness checks (local mode -> run_worker never called). Now import runpod (eager on main) runs run_startup_fitness_checks(), executes the built-in memory/disk/network/gpu checks, and any failure hard-kills the process via os._exit(1).
Same hazard for any auxiliary CLI/process that imports runpod only for the API client while RUNPOD_WEBHOOK_GET_JOB is inherited in the environment -- it will now run worker fitness checks and can be killed.
Suggest gating the import-time pass on the same local-mode/test-input signal that run_worker uses, so local and non-worker imports stay exempt.
|
|
||
| # Check the environment here rather than in start(), which a handler module | ||
| # only reaches after loading its model. No-op outside a real worker. | ||
| run_startup_fitness_checks() |
There was a problem hiding this comment.
The whole fail-fast benefit assumes import runpod eagerly imports runpod.serverless. That holds on this PR's declared base (main), but the repo's active feat/apps-sdk line lazy-loads serverless via PEP 562 __getattr__. There, a typical handler -- import runpod -> load model -> runpod.serverless.start(...) -- won't trigger serverless/__init__ until the start() line, i.e. after the multi-minute model load.
So on the apps-sdk line the checks fire no earlier than before, silently regressing, while the README/docs added in this PR assert "built-ins at import" / "any import runpod triggers it."
Which branch does this actually merge into? If it's the lazy-import line, either the docs need correcting or the trigger needs an explicit hook that doesn't depend on eager submodule import.
| # raises RuntimeError on Python 3.10+. | ||
| loop = asyncio.new_event_loop() | ||
| try: | ||
| loop.run_until_complete(run_fitness_checks(include_deferred=False)) |
There was a problem hiding this comment.
Running the network check at import time means a transient failure hard-exits the process mid-import runpod (_terminate_unhealthy -> os._exit(1)). On a cold worker container whose network stack isn't up yet when the handler module is first imported, this turns a recoverable warm-up delay into a boot crash-loop.
Previously this ran in run_worker after start(), giving the container time to become ready. Consider keeping network (and other environment-readiness) checks on the post-start() path, or adding a bounded retry before terminating.
| return | ||
|
|
||
| if _config_snapshot: | ||
| _warn_late_config() |
There was a problem hiding this comment.
run_startup_fitness_checks populates _config_snapshot microseconds before calling run_fitness_checks, where this if _config_snapshot: _warn_late_config() re-compares all nine env vars against the values just captured -- guaranteed no change, no warning. It's pure overhead on the import path; the late-config warning is only meaningful on the later run_worker pass. Consider skipping _warn_late_config() when invoked from the import-time pass.
justinwlin
changed the title
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Closes DR-1409.
Workers can spend minutes loading a model before discovering an unhealthy machine. This adds early hardware checks for explicitly identified handler processes while preserving worker-start checks for existing launchers, local tests, and helper scripts.
Behavior
runpod-worker handler.py(orrunpod-worker -m package.handler) runs memory, disk, CUDA-version, and native GPU health checks before executing the handler. No handler edits are needed.RUNPOD_FITNESS_WORKER_PIDto the handler's PID before exec. The top-level import hook requires that exact process identity and excludes local test/API arguments. The webhook environment alone never authorizes import-time checks.serverless, so the hook also works with lazy submodule loading. Legacy import paths remain compatibility aliases.fitness_check_setupand force-exit, even with live background threads.Configuration and rollout
RUNPOD_SKIP_FITNESS_CHECKS=truedisables all checks, including custom checks.RUNPOD_DEFER_FITNESS_CHECKS=truekeeps worker-start-only timing, including with the new launcher.python handler.pydeployments keep their previous worker-start timing until their launcher adopts the entrypoint or process hook. This PR does not change deployed platform launch configuration.Validation
feat/apps-sdkinitializer (eef269e76e3e07267cd84c123b3adbf261450df9) tested in an isolated package with the top-level hook: early checks pass without loading serverless, torch, or cupy.runpod-worker -mentrypoint smoke-tested.secret_name=calls; short, empty, and object-valued secrets are covered.