[Core] Add lifecycle hooks for worker extensions (#59463)

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>

tests/v1/worker/test_worker_base.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.v1.worker import worker_base
+from vllm.v1.worker.worker_base import WorkerBase, WorkerWrapperBase
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class ExtensionA:
+    events: list[str]
+
+    def init_worker_extension(self):
+        self.events.append("init-a")
+
+    def destroy_worker_extension(self):
+        self.events.append("destroy-a")
+
+
+class ExtensionB: ...
+
+
+class FailingExtension:
+    def destroy_worker_extension(self):
+        raise RuntimeError("destroy failed")
+
+
+@pytest.fixture
+def make_wrapper(monkeypatch):
+    # Fresh worker class per test, since extension injection mutates its bases.
+    class Worker(WorkerBase):
+        def __init__(self, **_kwargs):
+            self.events: list[str] = []
+
+        def init_device(self):
+            self.events.append("device")
+
+        def shutdown(self):
+            self.events.append("shutdown")
+
+    classes = {
+        cls.__name__: cls for cls in (Worker, ExtensionA, ExtensionB, FailingExtension)
+    }
+    monkeypatch.setattr(worker_base, "resolve_obj_by_qualname", classes.__getitem__)
+    monkeypatch.setattr(
+        worker_base,
+        "worker_receiver_cache_from_config",
+        lambda *_args, **_kwargs: None,
+    )
+
+    def make(extension: str) -> WorkerWrapperBase:
+        config = SimpleNamespace(
+            parallel_config=SimpleNamespace(
+                worker_cls="Worker", worker_extension_cls=extension
+            ),
+            enable_trace_function_call_for_thread=lambda: None,
+        )
+        wrapper = WorkerWrapperBase()
+        wrapper.init_worker([{"vllm_config": config}])
+        wrapper.init_device()
+        return wrapper
+
+    return make
+
+
+def test_worker_extension_hooks_follow_configured_class(make_wrapper):
+    for extension, expected in (
+        ("ExtensionA", ["device", "init-a", "destroy-a", "shutdown"]),
+        ("ExtensionB", ["device", "shutdown"]),
+        ("ExtensionA", ["device", "init-a", "destroy-a", "shutdown"]),
+    ):
+        wrapper = make_wrapper(extension)
+        wrapper.shutdown()
+        assert wrapper.worker.events == expected
+
+
+def test_failing_destroy_hook_is_logged_not_raised(make_wrapper, caplog_vllm):
+    wrapper = make_wrapper("FailingExtension")
+    wrapper.shutdown()
+    assert wrapper.worker.events == ["device", "shutdown"]
+    assert "Worker extension destroy hook failed" in caplog_vllm.text

vllm/config/parallel.py

@@ -309,7 +309,10 @@ class ParallelConfig:
     """The full name of the worker extension class to use. The worker extension
     class is dynamically inherited by the worker class. This is used to inject
     new attributes and methods to the worker class for use in collective_rpc
-    calls."""
+    calls. It may also define `init_worker_extension(self)`, called after
+    device and distributed initialization but before model loading, and
+    `destroy_worker_extension(self)`, called (best-effort) before the worker
+    shuts down."""
     master_addr: str = "127.0.0.1"
     """distributed master address for multi-node distributed 
     inference when distributed_executor_backend is mp."""

vllm/v1/worker/worker_base.py

@@ -244,9 +244,17 @@ class WorkerWrapperBase:
         # Initialized after init_worker is called
         self.worker: WorkerBase
         self.vllm_config: VllmConfig
+        self._worker_extension_cls: type | None = None
 
     def shutdown(self) -> None:
         if self.worker is not None:
+            hook = getattr(self._worker_extension_cls, "destroy_worker_extension", None)
+            if callable(hook):
+                try:
+                    hook(self.worker)
+                except Exception:
+                    # Don't raise: callers still need to run their own teardown.
+                    logger.exception("Worker extension destroy hook failed.")
             self.worker.shutdown()
 
     def update_environment_variables(
@@ -293,6 +301,7 @@ class WorkerWrapperBase:
                 "and pass the qualified name of the class as a string."
             )
 
+        worker_extension_cls: type | None = None
         if parallel_config.worker_extension_cls:
             worker_extension_cls = resolve_obj_by_qualname(
                 parallel_config.worker_extension_cls
@@ -336,6 +345,7 @@ class WorkerWrapperBase:
         with set_current_vllm_config(self.vllm_config):
             # To make vLLM config available during worker initialization
             self.worker = worker_class(**kwargs)
+            self._worker_extension_cls = worker_extension_cls
 
     def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
         kv_cache_config = kv_cache_configs[self.global_rank]
@@ -348,6 +358,9 @@ class WorkerWrapperBase:
         with set_current_vllm_config(self.vllm_config):
             # To make vLLM config available during device initialization
             self.worker.init_device()  # type: ignore
+            hook = getattr(self._worker_extension_cls, "init_worker_extension", None)
+            if callable(hook):
+                hook(self.worker)
 
     def __getattr__(self, attr: str):
         return getattr(self.worker, attr)