[Core] Add lifecycle hooks for worker extensions

tests/v1/worker/test_worker_base.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+
+import vllm.platforms
+import vllm.plugins
+from vllm.v1.executor import multiproc_executor
+from vllm.v1.executor.multiproc_executor import WorkerProc
+from vllm.v1.worker import worker_base
+from vllm.v1.worker.worker_base import WorkerBase, WorkerWrapperBase
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class Worker(WorkerBase):
+    def __init__(self, vllm_config, **_kwargs):
+        self.events = []
+        self.fail_destroy = vllm_config.fail_destroy
+
+    def init_device(self):
+        self.events.append("device")
+
+    def shutdown(self):
+        self.events.append("shutdown")
+
+
+class ExtensionA:
+    events: list[str]
+    fail_destroy: bool
+
+    def init_worker_extension(self):
+        self.events.append("init-a")
+
+    def destroy_worker_extension(self):
+        self.events.append("destroy-a")
+        if self.fail_destroy:
+            raise RuntimeError("destroy failed")
+
+
+class ExtensionB: ...
+
+
+@pytest.fixture
+def make_wrapper(monkeypatch):
+    classes = {
+        "Worker": Worker,
+        "ExtensionA": ExtensionA,
+        "ExtensionB": ExtensionB,
+    }
+    monkeypatch.setattr(worker_base, "resolve_obj_by_qualname", classes.__getitem__)
+    monkeypatch.setattr(
+        worker_base,
+        "worker_receiver_cache_from_config",
+        lambda *_args, **_kwargs: None,
+    )
+    monkeypatch.setattr(
+        worker_base, "set_current_vllm_config", lambda _config: nullcontext()
+    )
+    monkeypatch.setattr(vllm.plugins, "load_general_plugins", lambda: None)
+    monkeypatch.setattr(
+        vllm.platforms.current_platform,
+        "register_triton_kernel_overrides",
+        lambda: None,
+    )
+
+    def make(extension: str, fail_destroy: bool = False) -> WorkerWrapperBase:
+        config = SimpleNamespace(
+            parallel_config=SimpleNamespace(
+                worker_cls="Worker", worker_extension_cls=extension
+            ),
+            enable_trace_function_call_for_thread=lambda: None,
+            fail_destroy=fail_destroy,
+        )
+        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, fail_destroy, expected in (
+        ("ExtensionA", False, ["device", "init-a", "destroy-a", "shutdown"]),
+        ("ExtensionB", False, ["device", "shutdown"]),
+        ("ExtensionA", True, ["device", "init-a", "destroy-a", "shutdown"]),
+    ):
+        wrapper = make_wrapper(extension, fail_destroy)
+        # A failing destroy hook is logged, not raised.
+        wrapper.shutdown()
+        assert wrapper.worker.events == expected
+
+
+def test_failing_destroy_hook_does_not_skip_outer_teardown(make_wrapper, monkeypatch):
+    events = []
+    monkeypatch.setattr(
+        multiproc_executor,
+        "destroy_model_parallel",
+        lambda: events.append("destroy_model_parallel"),
+    )
+    monkeypatch.setattr(
+        multiproc_executor,
+        "destroy_distributed_environment",
+        lambda: events.append("destroy_distributed_environment"),
+    )
+    proc = object.__new__(WorkerProc)
+    proc.rpc_broadcast_mq = None
+    proc.worker_response_mq = None
+    proc.worker = make_wrapper("ExtensionA", fail_destroy=True)
+
+    proc.shutdown()
+
+    assert proc.worker.worker.events[-2:] == ["destroy-a", "shutdown"]
+    assert events == ["destroy_model_parallel", "destroy_distributed_environment"]

vllm/v1/worker/worker_base.py

@@ -244,10 +244,22 @@ def __init__(
         # 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:
-            self.worker.shutdown()
+        worker = self.__dict__.get("worker")
+        if worker is None:
+            return
+        hook = getattr(self._worker_extension_cls, "destroy_worker_extension", None)
+        try:
+            if callable(hook):
+                hook(worker)
+        except Exception:
+            # Log instead of raising so that callers still run their own
+            # teardown (e.g. destroying the distributed environment).
+            logger.exception("Worker extension destroy hook failed.")
+        finally:
+            worker.shutdown()
 
     def update_environment_variables(
         self,
@@ -293,6 +305,7 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
                 "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 +349,7 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None:
         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 +362,9 @@ def init_device(self):
         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)