[GG] perf: preload CUDA-clean forkserver for local GPU workers - #194
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughForkserver support is added to multiprocessing configuration and lifecycle management. Serve entrypoints and executors initialize a clean forkserver with preloaded modules, while GPU warmup imports are deferred until call time. Tests cover configuration, lifecycle, startup, fallback, and lazy loading. ChangesCUDA-clean forkserver support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI_main
participant ensure_cuda_clean_forkserver
participant multiprocessing_forkserver
CLI_main->>CLI_main: inspect serve command and forkserver configuration
CLI_main->>ensure_cuda_clean_forkserver: initialize preloaded forkserver
ensure_cuda_clean_forkserver->>multiprocessing_forkserver: start forkserver
multiprocessing_forkserver-->>CLI_main: return before command imports
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Updated the CUDA-clean warmup boundary to compose cleanly with #172. |
7804621 to
2f3dfb5
Compare
|
E2E follow-up: the original revision still allowed the normal The latest commit starts the CUDA-clean API forkserver at the beginning of the Validation on the same GLM-5.2 image/config/cache, TP8/DCP1/MTP0, GPUs 0-7:
Forkserver improved full warm startup by 12.11 s (9.3%). The process tree showed separate API and EngineCore forkservers with the expected preload lists, no spawn fallback was logged, InstantTensor loading completed, Tests: 19 targeted tests passed, including CLI-before-import bootstrap, idempotence, CUDA fail-closed behavior, and lazy worker warmup. The direct preload probe kept |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
vllm/utils/system_utils.py (2)
122-153: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winIdempotency silently drops later
preload_moduleslists.Once
_cuda_clean_forkserver_owner_pidmatches the current PID,ensure_cuda_clean_forkserverreturns immediately regardless of thepreload_modulespassed in. If a future call site needs modules the first caller didn't preload, that need is silently unmet with no log/warning to reveal the mismatch.♻️ Proposed fix: warn on preload mismatch
if _cuda_clean_forkserver_is_running(): + missing = set(preload_modules) - set(_cuda_clean_forkserver_preload_modules) + if missing: + logger.warning( + "CUDA-clean forkserver already running without preload " + "modules %s; they will not be preloaded.", sorted(missing) + ) return(requires tracking the applied preload list in a new module-level variable)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/utils/system_utils.py` around lines 122 - 153, Update ensure_cuda_clean_forkserver to track the preload modules applied by the current forkserver, and when it is already running compare them with the later preload_modules request. Warn about any mismatch before returning, while preserving idempotent behavior and existing validation for starting a new forkserver.
122-133: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring missing
Args/Raisessections.
ensure_cuda_clean_forkservertakespreload_modules/set_start_methodand raisesRuntimeError/ValueError, but the docstring documents neither.As per coding guidelines, "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields such as:param:,:return:, and:rtype:."📝 Proposed docstring fix
"""Start a forkserver before this process initializes an accelerator. A forkserver child inherits the server's runtime state. Starting that server after CUDA/XPU initialization would reproduce the same unsafe state inheritance as a direct fork, so fail closed instead of creating workers that can fail later during device initialization. + + Args: + preload_modules: Module names to import inside the forkserver before + it starts accepting fork requests. + set_start_method: Whether to also force the global multiprocessing + start method to "forkserver". + + Raises: + RuntimeError: If CUDA or XPU is already initialized in this process. + ValueError: If ``preload_modules`` is empty. """🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/utils/system_utils.py` around lines 122 - 133, Update the docstring for ensure_cuda_clean_forkserver to use Google-style sections: document preload_modules and set_start_method under Args, and document the RuntimeError and ValueError conditions under Raises. Keep the existing behavior and explanatory summary unchanged.Source: Coding guidelines
vllm/entrypoints/cli/main.py (1)
27-34: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate CUDA-clean-forkserver preload list across two files. The exact same three-module preload list is hardcoded in both
vllm/entrypoints/cli/main.pyandvllm/entrypoints/openai/api_server.py. Sinceensure_cuda_clean_forkserverapplies whichever list arrives first (idempotent, first-caller-wins), extracting this into one shared constant avoids the two literals silently drifting apart if only one is updated later.
vllm/entrypoints/cli/main.py#L27-L34: import a shared_SERVE_FORKSERVER_PRELOAD_MODULESconstant (e.g. fromvllm.utils.system_utils) instead of the inline list literal.vllm/entrypoints/openai/api_server.py#L130-L136: use the same shared constant instead of re-declaring the identical list literal.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/entrypoints/cli/main.py` around lines 27 - 34, The CUDA-clean forkserver preload modules are duplicated across both entry points. Define or reuse a shared _SERVE_FORKSERVER_PRELOAD_MODULES constant, import it in vllm/entrypoints/cli/main.py lines 27-34, and pass it to ensure_cuda_clean_forkserver; update vllm/entrypoints/openai/api_server.py lines 130-136 to use the same constant instead of its inline list.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vllm/v1/worker/gpu_worker.py`:
- Around line 95-101: Update the kernel_warmup function docstring to include
Google-style Args and Returns sections documenting the worker parameter and its
None return value, while preserving the existing summary and lazy-import
behavior.
---
Nitpick comments:
In `@vllm/entrypoints/cli/main.py`:
- Around line 27-34: The CUDA-clean forkserver preload modules are duplicated
across both entry points. Define or reuse a shared
_SERVE_FORKSERVER_PRELOAD_MODULES constant, import it in
vllm/entrypoints/cli/main.py lines 27-34, and pass it to
ensure_cuda_clean_forkserver; update vllm/entrypoints/openai/api_server.py lines
130-136 to use the same constant instead of its inline list.
In `@vllm/utils/system_utils.py`:
- Around line 122-153: Update ensure_cuda_clean_forkserver to track the preload
modules applied by the current forkserver, and when it is already running
compare them with the later preload_modules request. Warn about any mismatch
before returning, while preserving idempotent behavior and existing validation
for starting a new forkserver.
- Around line 122-133: Update the docstring for ensure_cuda_clean_forkserver to
use Google-style sections: document preload_modules and set_start_method under
Args, and document the RuntimeError and ValueError conditions under Raises. Keep
the existing behavior and explanatory summary unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: fd912023-7bb3-4ff9-8881-2052c07b0466
📒 Files selected for processing (9)
tests/entrypoints/test_cli_main.pytests/utils_/test_system_utils.pytests/v1/worker/test_gpu_worker.pyvllm/entrypoints/cli/main.pyvllm/entrypoints/openai/api_server.pyvllm/envs.pyvllm/utils/system_utils.pyvllm/v1/executor/multiproc_executor.pyvllm/v1/worker/gpu_worker.py
c1acc51
into
local-inference-lab:dev/gilded-gnosis
2 participants
Summary
Add an opt-in, CUDA-clean
forkserverpath for local multi-GPU workers. This is a clean extraction of the startup optimization validated on Kimi K3 TP16; it contains none of the Kimi loader or release-overlay changes from the experimental branch.When
VLLM_WORKER_MULTIPROC_METHOD=forkserver:gpu_workerimportskernel_warmuponly at warmup time, keeping forkserver preload CUDA-clean;This extends upstream vllm-project#45483. That PR makes the surviving forkserver path reachable but intentionally avoids preload. Preloading the second, EngineCore-owned worker forkserver is the part that removes repeated per-rank imports at high TP. The previous preload regression is addressed here by moving the identified CUDA-initializing import out of module scope and testing that contract.
Measured result
Matched Kimi K3 TP16 startup,
Initializing a V1 LLM enginetoStarting to load model:spawnforkserverAll 16 ranks entered distributed initialization within one second. Full EngineCore-to-ready time decreased from 407 s to 271 s. A repeat run completed model loading, CUDA graph capture,
/health, and a correctness request.Validation
tests/utils_/test_system_utils.pytests/v1/worker/test_gpu_worker.py14 passedtorch.cuda.is_initialized() == Falsemultiproc_executor:Falsegpu_worker:FalseThe release helper should keep this opt-in until GLM TP8 and the canonical image complete matched spawn/forkserver E2E validation.
Summary by CodeRabbit
New Features
forkservermultiprocessing method viaVLLM_WORKER_MULTIPROC_METHOD.Tests