@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d9ff7e4e9a

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[ROCm][CI] Stabilizing teardown and timeout of flaky tests to prevent rare OOMs (#44761)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Jun 8, 2026 at 01:11 UTC d9ff7e4e9a397c5c3ac6337116cc71969c966a23

4 files changed +130 -65

tests/conftest.py

+18 -17

@@ -854,8 +854,13 @@ class HfRunner:

854 return self

855

856 def __exit__(self, exc_type, exc_value, traceback):

857 + from tests.utils import wait_for_rocm_memory_to_settle

858 +

859 del self.model

860 cleanup_dist_env_and_memory()

861 + # ROCm frees VRAM lazily; wait so a runner started right after this HF

862 + # model exits does not OOM on its startup memory guard.

863 + wait_for_rocm_memory_to_settle()

864

865

866 @pytest.fixture(scope="session")

@@ -1248,25 +1253,13 @@ class VllmRunner:

1253 return self

1254

1255 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

1256 + from tests.utils import wait_for_rocm_memory_to_settle

1257

1258 # V1 startup requires free_memory >= total * gpu_memory_utilization.

1259 # 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 - )

1260 + # not fail the startup guard immediately after this runner exits. The

1261 + # wait is bounded so cleanup failures fail this test instead of hanging.

1262 + wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)

1263

1264 def __exit__(self, exc_type, exc_value, traceback):

1265 # Explicitly shutdown the engine core to release GPU resources

@@ -1276,8 +1269,16 @@ class VllmRunner:

1269 gpu_memory_utilization = (

1270 self.llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization

1271 )

1272 + from vllm.platforms import current_platform

1273 +

1274 try:

1280 - self.llm.llm_engine.engine_core.shutdown()

1275 + # Give the engine core time to run its own graceful shutdown

1276 + # (model_executor teardown + empty_cache + process-group destroy)

1277 + # before the process manager SIGKILLs it at the default 5s. On ROCm

1278 + # a hard kill leaves the whole allocation for the driver's slow async

1279 + # VRAM reclamation, which starves the next test's startup.

1280 + shutdown_timeout = 60.0 if current_platform.is_rocm() else None

1281 + self.llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)

1282 except Exception:

1283 # Ignore shutdown errors as cleanup will still proceed

1284 pass

tests/models/language/generation/test_hybrid.py

+8 -26

@@ -7,7 +7,7 @@ from contextlib import contextmanager, nullcontext

7 import pytest

9 from tests.models.registry import HF_EXAMPLE_MODELS

10 -from tests.utils import multi_gpu_test, wait_for_gpu_memory_to_clear

10 +from tests.utils import multi_gpu_test

11 from vllm import LLM

12 from vllm.engine.arg_utils import EngineArgs

13 from vllm.platforms import current_platform

@@ -405,28 +405,10 @@ def _get_vllm_runner_params(

405 }

406

407

408 -def _wait_for_rocm_memory_to_settle() -> None:

409 - if not current_platform.is_rocm():

410 - return

411 -

412 - num_gpus = current_platform.device_count()

413 - if num_gpus == 0:

414 - return

415 -

416 - wait_for_gpu_memory_to_clear(

417 - devices=list(range(num_gpus)),

418 - threshold_ratio=0.01,

419 - timeout_s=120,

420 - )

421 -

422 -

408 @contextmanager

424 -def _owned_vLLM_runner(vllm_runner, kwargs):

425 - try:

426 - with vllm_runner(**kwargs) as runner:

427 - yield runner

428 - finally:

429 - _wait_for_rocm_memory_to_settle()

409 +def _owned_vllm_runner(vllm_runner, kwargs):

410 + with vllm_runner(**kwargs) as runner:

411 + yield runner

412

413

414 def _get_vLLM_output(

@@ -439,7 +421,7 @@ def _get_vLLM_output(

421 vllm_model=None,

422 ):

423 runner_context = (

442 - _owned_vLLM_runner(vllm_runner, kwargs)

424 + _owned_vllm_runner(vllm_runner, kwargs)

425 if vllm_model is None

426 else nullcontext(vllm_model)

427 )

@@ -801,7 +783,7 @@ def test_apc_multiple_prompts_partial_cached_outputs(

783

784 # Cache only part of all the prompts

785 vllm_runner_kwargs["enable_prefix_caching"] = True

804 - with _owned_vLLM_runner(vllm_runner, vllm_runner_kwargs) as vllm_model:

786 + with _owned_vllm_runner(vllm_runner, vllm_runner_kwargs) as vllm_model:

787 vllm_outputs_partial_cache, _ = _get_vLLM_output(

788 vllm_runner,

789 vllm_runner_kwargs,

@@ -861,7 +843,7 @@ def test_same_mamba_output_apc_on_vs_off(

843

844 # No prefix caching

845 kwargs_no_apc = {**base_kwargs, "enable_prefix_caching": False}

864 - with _owned_vLLM_runner(vllm_runner, kwargs_no_apc) as vllm_model:

846 + with vllm_runner(**kwargs_no_apc) as vllm_model:

847 outputs_no_apc, _ = _get_vLLM_output(

848 vllm_runner,

849 kwargs_no_apc,

@@ -876,7 +858,7 @@ def test_same_mamba_output_apc_on_vs_off(

858 "enable_prefix_caching": True,

859 "mamba_block_size": 16,

860 }

879 - with _owned_vLLM_runner(vllm_runner, kwargs_with_apc) as vllm_model:

861 + with vllm_runner(**kwargs_with_apc) as vllm_model:

862 outputs_with_apc, _ = _get_vLLM_output(

863 vllm_runner,

864 kwargs_with_apc,

tests/models/language/pooling/test_colbert.py

+34 -12

@@ -6,15 +6,19 @@ Tests are parametrized across multiple ColBERT backbones to ensure the

6 generic ColBERT support works with different encoder architectures.

7 """

9 +from contextlib import contextmanager

10 +

11 import pytest

12 import torch

14 +from tests.utils import wait_for_rocm_memory_to_settle

15 +from vllm.distributed import cleanup_dist_env_and_memory

16 from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score

18 # -----------------------------------------------------------------------

19 # Model definitions: (model_name, colbert_dim, extra vllm_runner kwargs)

20 # -----------------------------------------------------------------------

17 -COLBERT_MODELS: dict[str, dict] = {

21 +COLBERT_MODELS = {

22 "bert": {

23 "model": "answerdotai/answerai-colbert-small-v1",

24 "colbert_dim": 96,

@@ -145,6 +149,24 @@ def _compute_hf_colbert_embeddings(model, tokenizer, linear_weight, texts, devic

149 return embeddings

150

151

152 +@contextmanager

153 +def _hf_colbert_model(model_name: str, hf_spec: dict, device: torch.device):

154 + """Load the HF backbone + ColBERT projection, freeing the GPU on exit.

155 +

156 + These live outside any runner context, so without explicit cleanup ROCm

157 + keeps the VRAM resident and the next backend parametrization (or test)

158 + OOMs on startup.

159 + """

160 + hf_model = _load_hf_model(model_name, hf_spec, device)

161 + linear_weight = _load_projection_weight(model_name, hf_spec, device)

162 + try:

163 + yield hf_model, linear_weight

164 + finally:

165 + del hf_model, linear_weight

166 + cleanup_dist_env_and_memory()

167 + wait_for_rocm_memory_to_settle()

168 +

169 +

170 def _assert_embeddings_close(vllm_outputs, hf_embeddings):

171 """Assert that vLLM and HuggingFace embeddings match."""

172 for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):

@@ -363,9 +385,11 @@ def test_colbert_hf_comparison(vllm_runner, backend):

385

386 spec = COLBERT_MODELS[backend]

387 hf_spec = spec["hf_comparison"]

388 + extra_kwargs = spec["extra_kwargs"]

389 model_name = spec["model"]

390 assert isinstance(model_name, str)

391 assert isinstance(hf_spec, dict)

392 + assert isinstance(extra_kwargs, dict)

393 test_texts = [TEXTS_1[0], TEXTS_2[0]]

394

395 with vllm_runner(

@@ -374,7 +398,7 @@ def test_colbert_hf_comparison(vllm_runner, backend):

398 dtype="float32",

399 max_model_len=spec["max_model_len"],

400 enforce_eager=True,

377 - **spec["extra_kwargs"],

401 + **extra_kwargs,

402 ) as vllm_model:

403 vllm_outputs = vllm_model.token_embed(test_texts)

404

@@ -384,15 +408,13 @@ def test_colbert_hf_comparison(vllm_runner, backend):

408 model_name,

409 trust_remote_code=hf_spec.get("trust_remote_code", False),

410 )

387 - hf_model = _load_hf_model(model_name, hf_spec, device)

388 - linear_weight = _load_projection_weight(model_name, hf_spec, device)

389 -

390 - hf_embeddings = _compute_hf_colbert_embeddings(

391 - hf_model,

392 - hf_tokenizer,

393 - linear_weight,

394 - test_texts,

395 - device,

396 - )

411 + with _hf_colbert_model(model_name, hf_spec, device) as (hf_model, linear_weight):

412 + hf_embeddings = _compute_hf_colbert_embeddings(

413 + hf_model,

414 + hf_tokenizer,

415 + linear_weight,

416 + test_texts,

417 + device,

418 + )

419

420 _assert_embeddings_close(vllm_outputs, hf_embeddings)

tests/utils.py

+70 -10

@@ -538,11 +538,22 @@ class RemoteVLLMServer:

538 if current_platform.is_rocm():

539 with _nvml():

540 handles = amdsmi_get_processor_handles()

541 - total_used = 0

542 - for handle in handles:

541 + devices = get_physical_device_indices(

542 + list(range(current_platform.device_count()))

543 + )

544 + total_used_mib = 0

545 + for device in devices:

546 + handle = handles[device]

547 vram_info = amdsmi_get_gpu_vram_usage(handle)

544 - total_used += vram_info["vram_used"]

545 - return total_used

548 + total_used_mib += vram_info["vram_used"]

549 + # amdsmi reports VRAM in MiB; convert to bytes so this

550 + # matches the CUDA/nvml branch (already bytes) and the

551 + # byte-based target in _wait_for_gpu_memory_release. Without

552 + # this, that wait compares MiB against a ~2e9-byte target,

553 + # is always satisfied instantly, and returns "released to

554 + # 0.00 GB" while the previous server's VRAM is still

555 + # resident -- OOMing the next server's startup on ROCm.

556 + return total_used_mib * 1024 * 1024

557 elif current_platform.is_cuda():

558 with _nvml():

559 total_used = 0

@@ -1428,6 +1439,18 @@ def wait_for_gpu_memory_to_clear(

1439 timeout_s: float = 120,

1440 ) -> None:

1441 assert threshold_bytes is not None or threshold_ratio is not None

1442 + if (

1443 + current_platform.is_rocm()

1444 + and threshold_ratio is not None

1445 + and threshold_ratio < 0.05

1446 + ):

1447 + # ROCm can keep a small runtime/driver footprint resident even after

1448 + # all model allocations are gone. On MI300 this has been observed

1449 + # around 2.5 GiB, which is above a strict 1% idle threshold but nowhere

1450 + # near the amount of free memory needed by the next vLLM runner.

1451 + min_threshold_bytes = 4 * 1024**3

1452 + threshold_bytes = max(threshold_bytes or 0, min_threshold_bytes)

1453 +

1454 # Use nvml instead of pytorch to reduce measurement error from torch cuda

1455 # context.

1456 devices = get_physical_device_indices(devices)

@@ -1454,15 +1477,26 @@ def wait_for_gpu_memory_to_clear(

1477 print(f"{k}={v}; ", end="")

1478 print("")

1479

1457 - if threshold_bytes is not None:

1458 - is_free = lambda used, total: used <= threshold_bytes / 2**30

1459 - threshold = f"{threshold_bytes / 2**30} GiB"

1480 + if threshold_bytes is not None and threshold_ratio is not None:

1481 + threshold_gib = threshold_bytes / 2**30

1482 + threshold = f"max({threshold_gib:.2f} GiB, {threshold_ratio:.3f})"

1483 + all_free = all(

1484 + used <= max(threshold_gib, total * threshold_ratio)

1485 + for used, total in output_raw.values()

1486 + )

1487 + elif threshold_bytes is not None:

1488 + threshold_gib = threshold_bytes / 2**30

1489 + threshold = f"{threshold_gib} GiB"

1490 + all_free = all(used <= threshold_gib for used, _ in output_raw.values())

1491 else:

1461 - is_free = lambda used, total: used / total <= threshold_ratio

1462 - threshold = f"{threshold_ratio:.2f}"

1492 + assert threshold_ratio is not None

1493 + threshold = f"{threshold_ratio:.3f}"

1494 + all_free = all(

1495 + used / total <= threshold_ratio for used, total in output_raw.values()

1496 + )

1497

1498 dur_s = time.time() - start_time

1465 - if all(is_free(used, total) for used, total in output_raw.values()):

1499 + if all_free:

1500 print(

1501 f"Done waiting for free GPU memory on devices {devices=} "

1502 f"({threshold=}) {dur_s=:.02f}"

@@ -1478,6 +1512,32 @@ def wait_for_gpu_memory_to_clear(

1512 time.sleep(5)

1513

1514

1515 +def wait_for_rocm_memory_to_settle(

1516 + *,

1517 + threshold_ratio: float = 0.1,

1518 + timeout_s: float = 240,

1519 +) -> None:

1520 + """Block until ROCm device VRAM usage drops below ``threshold_ratio``.

1521 +

1522 + ROCm reclaims GPU memory more lazily than CUDA, so back-to-back model

1523 + loads in a single test process can OOM the *next* engine/model startup

1524 + even after ``cleanup_dist_env_and_memory``. This gives the driver time to

1525 + actually release VRAM before the next allocation. No-op off ROCm.

1526 + """

1527 + if not current_platform.is_rocm():

1528 + return

1529 +

1530 + num_gpus = current_platform.device_count()

1531 + if num_gpus == 0:

1532 + return

1533 +

1534 + wait_for_gpu_memory_to_clear(

1535 + devices=list(range(num_gpus)),

1536 + threshold_ratio=threshold_ratio,

1537 + timeout_s=timeout_s,

1538 + )

1539 +

1540 +

1541 _P = ParamSpec("_P")

1542

1543