A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][V1] Clean up compiled-model bytecode hooks on VllmRunner exit (#45195)
Signed-off-by: Ting Sun <suntcrick@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ting SUN committed
Jun 17, 2026 at 11:31 UTC
7b5d60cc3733025fc97ee03a91b5f46188aefe7d
4 files changed
+61
-2
tests/conftest.py
+1
@@ -1294,6 +1294,7 @@ class VllmRunner:
1294
# Ignore shutdown errors as cleanup will still proceed
1295
pass
1296
del self.llm
1297
+ torch._dynamo.reset()
1298
cleanup_dist_env_and_memory()
1299
self._wait_for_rocm_memory_release(gpu_memory_utilization)
1300
tests/v1/shutdown/test_delete.py
+28
-1
@@ -4,7 +4,8 @@
4
5
import pytest
6
7
-from tests.utils import wait_for_gpu_memory_to_clear
7
+from tests.conftest import VllmRunner
8
+from tests.utils import create_new_process_for_each_test, wait_for_gpu_memory_to_clear
9
from tests.v1.shutdown.utils import (
10
SHUTDOWN_TEST_THRESHOLD_BYTES,
11
SHUTDOWN_TEST_TIMEOUT_SEC,
@@ -106,3 +107,29 @@ def test_llm_delete(
107
devices=list(range(tensor_parallel_size)),
108
threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
109
)
110
+
111
+
112
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
113
+@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
114
+@pytest.mark.parametrize("model", MODELS)
115
+@pytest.mark.parametrize("send_one_request", [False, True])
116
+def test_llm_delete_inprocess(
117
+ monkeypatch,
118
+ model: str,
119
+ send_one_request: bool,
120
+) -> None:
121
+ """Test that VllmRunner frees GPU memory in in-process (no MP) mode."""
122
+ with monkeypatch.context() as m:
123
+ m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
124
+
125
+ with VllmRunner(model) as vllm_model:
126
+ if send_one_request:
127
+ vllm_model.generate(
128
+ ["Hello my name is"],
129
+ SamplingParams(max_tokens=1),
130
+ )
131
+
132
+ wait_for_gpu_memory_to_clear(
133
+ devices=[0],
134
+ threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
135
+ )
vllm/compilation/wrapper.py
+9
-1
@@ -154,7 +154,9 @@ class TorchCompileWithNoGuardsWrapper:
154
)
155
156
if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE:
157
- torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)
157
+ self._bytecode_hook_handle = (
158
+ torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)
159
+ )
160
self._compiled_bytecode: CodeType | None = None
161
162
def aot_compile(self, *args: Any, **kwargs: Any) -> Any:
@@ -261,6 +263,12 @@ class TorchCompileWithNoGuardsWrapper:
263
)
264
raise RuntimeError(msg)
265
266
+ def cleanup(self) -> None:
267
+ """Remove the bytecode hook registered by this instance."""
268
+ handle = getattr(self, "_bytecode_hook_handle", None)
269
+ if handle is not None:
270
+ handle.remove()
271
+
272
@contextmanager
273
def _dispatch_to_compiled_code(self) -> Generator[None, None, None]:
274
# noqa: E501
vllm/v1/engine/llm_engine.py
+23
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import time
5
+import weakref
6
from collections.abc import Callable, Mapping
7
from copy import copy
8
from typing import Any
@@ -123,6 +124,14 @@ class LLMEngine:
124
# for v0 compatibility
125
self.model_executor = self.engine_core.engine_core.model_executor # type: ignore
126
127
+ # Capture the model while reachable so the finalizer can drop the
128
+ # bytecode hooks pinning it (frees GPU memory on engine deletion).
129
+ model = self._get_driver_model_for_cleanup()
130
+ if model is not None:
131
+ self._finalizer = weakref.finalize(
132
+ self, LLMEngine._cleanup_instance_caches, model
133
+ )
134
+
135
if self.external_launcher_dp:
136
# If we use DP in external launcher mode, we reuse the
137
# existing DP group used for data communication.
@@ -419,6 +428,20 @@ class LLMEngine:
428
def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
429
return self.collective_rpc("apply_model", args=(func,))
430
431
+ def _get_driver_model_for_cleanup(self) -> nn.Module | None:
432
+ driver_worker = getattr(self.model_executor, "driver_worker", None)
433
+ model_runner = getattr(driver_worker, "model_runner", None)
434
+ return getattr(model_runner, "model", None)
435
+
436
+ @staticmethod
437
+ def _cleanup_instance_caches(model) -> None:
438
+ """Remove the bytecode hooks that pin the compiled model."""
439
+ from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
440
+
441
+ for module in model.modules():
442
+ if isinstance(module, TorchCompileWithNoGuardsWrapper):
443
+ module.cleanup()
444
+
445
def __del__(self):
446
dp_group = getattr(self, "dp_group", None)
447
if dp_group is not None and not self.external_launcher_dp: