feat(serverless): run fitness checks at startup, add skip env var - #578
Draft
justinwlin wants to merge 4 commits into
Draft
feat(serverless): run fitness checks at startup, add skip env var#578justinwlin wants to merge 4 commits into
justinwlin wants to merge 4 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes DR-1409.
Problem
Fitness (health) checks run in
run_worker, insiderunpod.serverless.start()— the last line of a typical handler module. A worker that loads a large model at import spends minutes loading before the SDK checks whether the GPU works or the disk is full. The ticket also asks for an env var to skip the checks.Change
import runpod.serverless, before model load: memory, disk, network, CUDA version (nvidia-smi), and the nativegpu_testbinary.@defer_to_worker_startand still run atstart(). Both importtorchand 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.run_fitness_checks()records what passed, so therun_workerpass only runs checks registered since — the user's@register_fitness_checkfunctions, which cannot exist before that import.RUNPOD_SKIP_FITNESS_CHECKS=truedisables every check. (RUNPOD_SKIP_GPU_CHECK/RUNPOD_SKIP_AUTO_SYSTEM_CHECKSonly ever covered built-ins.)RUNPOD_DEFER_FITNESS_CHECKS=truerestores the previousstart()-only timing.The import-time pass no-ops outside a real worker (no
RUNPOD_WEBHOOK_GET_JOB), so local dev, tests, and therunpodCLI are untouched, and it never raises — failing to run the checks must not stop a worker booting. A failing check still force-exits as before.Verification
pytest tests— 645 passed, including 12 new tests intest_fitness/test_startup.py.Behavior change worth flagging
The memory check now measures a fresh container instead of one with the model already loaded, so
RUNPOD_MIN_MEMORY_GBvalidates the environment the worker was given rather than the headroom left after loading. Conversely thegpu_testVRAM allocation no longer contends with a resident model, which should remove a class of spurious failures.For review
This works by calling a function at module top level in
runpod/serverless/__init__.py, soimport runpodnow spawns thegpu_testsubprocess on a worker. TheRUNPOD_WEBHOOK_GET_JOBgate keeps that off laptops and CI, but it does mean an import behaves differently based on an env var the reader can't see. The alternative is an explicit opt-in env var, at the cost of nobody getting the fix by default.