Repository navigation
[Core] Add lifecycle hooks for worker extensions - #59463
Merged
aoshen02 merged 10 commits intoOct 13, 2026
Merged
Conversation
Assisted-by: OpenAI Codex Signed-off-by: Dakai An <dakaian108@gmail.com>
aoshen02
reviewed
Oct 2, 2026
aoshen02
left a comment
aoshen02
left a comment
Contributor
There was a problem hiding this comment.
Reviewed commit 3461bd74f68ed2498b4c07e6a019aee2f3ad4311 with the current Codex model, Claude Opus 5.5, and Kimi K3. Both external reviewers completed their reports; the conclusions below were cross-checked against the implementation.
One P2 shutdown robustness issue is recorded inline: a failing extension destroy hook still allows worker.shutdown(), but aborts the remaining engine/executor teardown. No P0/P1 issue was confirmed.
Before finalizing the hook API, please document the contract beside ParallelConfig.worker_extension_cls or in the worker-extension documentation:
- Define the supported signature as an ordinary instance method, e.g.
def init_worker_extension(self), with no extra arguments. The current class lookup plus explicit worker argument is not normal descriptor binding for a no-argument static/class method. - Init runs after device/distributed setup and before model loading/KV-cache initialization. Destroy is best-effort; Ray force-kill does not guarantee it, and failed startup currently does not guarantee a matching destroy call.
- State whether repeated shutdown and partial initialization must be tolerated by the extension. Avoid promising exactly-once paired lifecycle callbacks without implementing that guarantee.
Simplicity/test suggestions (non-blocking): the production change is already small, so no generic hook dispatcher or framework is needed. Keep configured-class lookup and self.__dict__.get("worker"); replacing these with unconditional instance lookup or ordinary getattr(self, "worker", None) can reintroduce stale-extension dispatch or __getattr__ recursion. Preserve the existing sequential A/B/A test because it intentionally exercises persistent class injection. Add a focused failure-path test that checks outer scheduler/distributed cleanup, rather than asserting only the inner worker shutdown. The pre-existing conflict between two injected classes that both define the same methods is a limitation to document/test, not a new independent regression to attribute to this patch.
Validation: the new tests/v1/worker/test_worker_base.py passed (1 test). CPU probes using the actual wrapper and unchanged shutdown-method bodies with mocked collaborators confirmed skipped outer cleanup on a hook exception. No GPU, real distributed-executor, or Ray end-to-end run was performed; broader worker tests are limited by missing CUDA extensions. This is a COMMENT review, not an approval or a claim of complete runtime validation.
| try: | ||
| hook = getattr(self._worker_extension_cls, "destroy_worker_extension", None) | ||
| if callable(hook): | ||
| hook(worker) |
Contributor
There was a problem hiding this comment.
[P2] Finish outer teardown when an extension destroy hook fails
The finally guarantees worker.shutdown(), but the hook exception still escapes afterward. On the in-process path, UniProcExecutor.shutdown() propagates it through EngineCore.shutdown() before scheduler.shutdown(), gc.unfreeze(), and cleanup_dist_env_and_memory() can run. On the multiprocess path, it also skips WorkerProc.shutdown()'s destroy_model_parallel() and destroy_distributed_environment(). This is especially relevant to colocated training/inference that keeps the hosting process alive. CPU probes with the unchanged shutdown methods reproduce the skipped cleanup; the new test only checks that the inner worker shutdown ran. Please catch and log ordinary extension teardown errors while continuing cleanup, or ensure every outer cleanup step runs before re-raising. Add a regression test for that outer failure path.
penguin-wwy
reviewed
Oct 8, 2026
| self.worker.init_device() # type: ignore | ||
| hook = getattr(self._worker_extension_cls, "init_worker_extension", None) | ||
| if callable(hook): | ||
| hook(self.worker) |
There was a problem hiding this comment.
Should init_worker_extension, as the base class initialization function, be placed at an earlier point—for example, after the call to worker_class(**kwargs)?
Contributor
Author
There was a problem hiding this comment.
Thanks. I didn't follow the design in #59120. I put it after init_device here. The device binding and the distributed setup both happen there, and the use cases in #59120 (weight sync, communicators, buffers) need both. The hook still runs before load_model() and any collective_rpc, so I think it still covers what you previously asked for.
Do you have a case that needs setup before device init? If so, I'm happy to add.
A failing destroy_worker_extension() hook raised out of WorkerWrapperBase.shutdown(), which skipped the caller's remaining teardown (e.g. destroy_model_parallel() and destroy_distributed_environment() in WorkerProc.shutdown()). Log the error and continue so the rest of shutdown always runs. Signed-off-by: Dakai An <dakaian108@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
andakai
marked this pull request as ready for review
October 12, 2026 13:56
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Document init_worker_extension() and destroy_worker_extension() on ParallelConfig.worker_extension_cls. Keep main's original self.worker check in WorkerWrapperBase.shutdown() and only wrap the destroy hook call in try/except. Signed-off-by: Dakai An <dakaian108@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
andakai
requested review from
ProExpertProg,
WoosukKwon,
houseroad,
mgoin,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256 and
youkaichao
as code owners
October 12, 2026 14:29
Drop monkeypatches the tests don't need and leave the destroy failure case to the dedicated outer-teardown test. Signed-off-by: Dakai An <dakaian108@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
✅ @andakai, CI is now available for this PR.
|
aoshen02
approved these changes
Oct 13, 2026
Create a fresh worker class per test so extension injection does not leak between tests, and check the failing destroy hook at the wrapper level instead of driving WorkerProc internals. Signed-off-by: Dakai An <dakaian108@gmail.com>
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94408 for commit |
Collaborator
CI selector (shadow): 132 test steps (171 jobs) instead of 70 (90 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (132)
Would skip (today's rules run them) (17)
Would add (today's rules do not run them) (79)
AMD mirrors: would skip (12)
AMD mirrors: would add (69)
3 changed files · base |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #94460 for commit |
AndreasKaratzas
pushed a commit
to AndreasKaratzas/vllm
that referenced
this pull request
Oct 16, 2026
Signed-off-by: Dakai An <dakaian108@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>
michael-gpu
pushed a commit
to practical-inference/vllm-ampere-extended
that referenced
this pull request
Oct 20, 2026
Signed-off-by: Dakai An <dakaian108@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>
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
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.
Purpose
Related to #59120. Worker extensions currently require callers to initialize their state manually. This change adds optional
init_worker_extension()anddestroy_worker_extension()hooks: initialization runs after device and distributed setup, and cleanup runs before worker shutdown. Hooks are selected from the configured extension class, so a previously injected extension cannot be called by mistake. Ray executors currently force-kill worker actors, so cleanup is not guaranteed on Ray shutdown.Currently the implementation is a rough initial draft. Will need further discussion in #59120.
AI assistance: OpenAI Codex assisted with implementation and review.