diff --git a/tests/v1/worker/test_worker_base.py b/tests/v1/worker/test_worker_base.py new file mode 100644 index 0000000000..079b643dc8 --- /dev/null +++ b/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 diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index e103c9ffd7..73d83442ba 100644 --- a/vllm/config/parallel.py +++ b/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.""" diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 95a0c3475d..e1a4eaf291 100644 --- a/vllm/v1/worker/worker_base.py +++ b/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)