diff --git a/tests/conftest.py b/tests/conftest.py index 3c698c666d5..4fc43ef04b1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1294,6 +1294,7 @@ class VllmRunner: # Ignore shutdown errors as cleanup will still proceed pass del self.llm + torch._dynamo.reset() cleanup_dist_env_and_memory() self._wait_for_rocm_memory_release(gpu_memory_utilization) diff --git a/tests/v1/shutdown/test_delete.py b/tests/v1/shutdown/test_delete.py index adf99fb922d..39386f3fd63 100644 --- a/tests/v1/shutdown/test_delete.py +++ b/tests/v1/shutdown/test_delete.py @@ -4,7 +4,8 @@ import pytest -from tests.utils import wait_for_gpu_memory_to_clear +from tests.conftest import VllmRunner +from tests.utils import create_new_process_for_each_test, wait_for_gpu_memory_to_clear from tests.v1.shutdown.utils import ( SHUTDOWN_TEST_THRESHOLD_BYTES, SHUTDOWN_TEST_TIMEOUT_SEC, @@ -106,3 +107,29 @@ def test_llm_delete( devices=list(range(tensor_parallel_size)), threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES, ) + + +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC) +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("send_one_request", [False, True]) +def test_llm_delete_inprocess( + monkeypatch, + model: str, + send_one_request: bool, +) -> None: + """Test that VllmRunner frees GPU memory in in-process (no MP) mode.""" + with monkeypatch.context() as m: + m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + with VllmRunner(model) as vllm_model: + if send_one_request: + vllm_model.generate( + ["Hello my name is"], + SamplingParams(max_tokens=1), + ) + + wait_for_gpu_memory_to_clear( + devices=[0], + threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES, + ) diff --git a/vllm/compilation/wrapper.py b/vllm/compilation/wrapper.py index 5635fe03ae2..4821daf9dde 100644 --- a/vllm/compilation/wrapper.py +++ b/vllm/compilation/wrapper.py @@ -154,7 +154,9 @@ class TorchCompileWithNoGuardsWrapper: ) if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE: - torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook) + self._bytecode_hook_handle = ( + torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook) + ) self._compiled_bytecode: CodeType | None = None def aot_compile(self, *args: Any, **kwargs: Any) -> Any: @@ -261,6 +263,12 @@ class TorchCompileWithNoGuardsWrapper: ) raise RuntimeError(msg) + def cleanup(self) -> None: + """Remove the bytecode hook registered by this instance.""" + handle = getattr(self, "_bytecode_hook_handle", None) + if handle is not None: + handle.remove() + @contextmanager def _dispatch_to_compiled_code(self) -> Generator[None, None, None]: # noqa: E501 diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index f3e8a95b0d6..ff86a1dffd9 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time +import weakref from collections.abc import Callable, Mapping from copy import copy from typing import Any @@ -123,6 +124,14 @@ class LLMEngine: # for v0 compatibility self.model_executor = self.engine_core.engine_core.model_executor # type: ignore + # Capture the model while reachable so the finalizer can drop the + # bytecode hooks pinning it (frees GPU memory on engine deletion). + model = self._get_driver_model_for_cleanup() + if model is not None: + self._finalizer = weakref.finalize( + self, LLMEngine._cleanup_instance_caches, model + ) + if self.external_launcher_dp: # If we use DP in external launcher mode, we reuse the # existing DP group used for data communication. @@ -419,6 +428,20 @@ class LLMEngine: def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]: return self.collective_rpc("apply_model", args=(func,)) + def _get_driver_model_for_cleanup(self) -> nn.Module | None: + driver_worker = getattr(self.model_executor, "driver_worker", None) + model_runner = getattr(driver_worker, "model_runner", None) + return getattr(model_runner, "model", None) + + @staticmethod + def _cleanup_instance_caches(model) -> None: + """Remove the bytecode hooks that pin the compiled model.""" + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + + for module in model.modules(): + if isinstance(module, TorchCompileWithNoGuardsWrapper): + module.cleanup() + def __del__(self): dp_group = getattr(self, "dp_group", None) if dp_group is not None and not self.external_launcher_dp: