[Bugfix][V1] Clean up compiled-model bytecode hooks on VllmRunner exit (#45195)

Signed-off-by: Ting Sun <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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)
 

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,
+        )

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

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: