A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 23:15 UTC
[ROCm][CI] Stabilize runner teardown between sampler tests (#43023)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 22, 2026 at 23:19 UTC
6a4723a2e0d808ce33bfe378f761e585a7685d58
2 files changed
+32
-1
tests/conftest.py
+25
@@ -1247,11 +1247,35 @@ class VllmRunner:
1247
def __enter__(self):
1248
return self
1249
1250
+ def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None:
1251
+ from tests.utils import wait_for_gpu_memory_to_clear
1252
+ from vllm.platforms import current_platform
1253
+
1254
+ if not current_platform.is_rocm():
1255
+ return
1256
+
1257
+ num_gpus = torch.accelerator.device_count()
1258
+ if num_gpus == 0:
1259
+ return
1260
+
1261
+ # V1 startup requires free_memory >= total * gpu_memory_utilization.
1262
+ # Wait for the complementary used-memory ratio so the next runner does
1263
+ # not fail the startup guard immediately after this runner exits. Bound
1264
+ # the wait so cleanup failures fail this test instead of hanging.
1265
+ wait_for_gpu_memory_to_clear(
1266
+ devices=list(range(num_gpus)),
1267
+ threshold_ratio=1.0 - gpu_memory_utilization,
1268
+ timeout_s=120,
1269
+ )
1270
+
1271
def __exit__(self, exc_type, exc_value, traceback):
1272
# Explicitly shutdown the engine core to release GPU resources
1273
# This is needed because when executing consecutive tests, the GC
1274
# might not be fast enough in shutting down the llm engine. This can lead to OOMs
1275
# because when the next test starts some GPU memory is still in use.
1276
+ gpu_memory_utilization = (
1277
+ self.llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
1278
+ )
1279
try:
1280
self.llm.llm_engine.engine_core.shutdown()
1281
except Exception:
@@ -1259,6 +1283,7 @@ class VllmRunner:
1283
pass
1284
del self.llm
1285
cleanup_dist_env_and_memory()
1286
+ self._wait_for_rocm_memory_release(gpu_memory_utilization)
1287
1288
1289
@pytest.fixture(scope="session")
vllm/v1/engine/core.py
+7
-1
@@ -22,7 +22,10 @@ import zmq
22
23
import vllm.envs as envs
24
from vllm.config import ParallelConfig, VllmConfig
25
-from vllm.distributed import stateless_destroy_torch_distributed_process_group
25
+from vllm.distributed import (
26
+ cleanup_dist_env_and_memory,
27
+ stateless_destroy_torch_distributed_process_group,
28
+)
29
from vllm.envs import enable_envs_cache
30
from vllm.logger import init_logger
31
from vllm.logging_utils.dump_input import dump_engine_exception
@@ -603,6 +606,9 @@ class EngineCore:
606
# visible to the garbage collector again. Without this, deleting
607
# the engine in-process (e.g. unit tests) leaks GPU memory.
608
gc.unfreeze()
609
+ # Tear down distributed state initialized in this EngineCore process
610
+ # before it exits and release cached memory.
611
+ cleanup_dist_env_and_memory()
612
613
def profile(self, is_start: bool = True, profile_prefix: str | None = None):
614
self.model_executor.profile(is_start, profile_prefix)