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 12, 08:30 UTC
Add gpu memory wait before test_async_tp (#28893)
Signed-off-by: angelayi <yiangela7@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Angela Yi committed
Nov 28, 2025 at 20:19 UTC
4b17ce6815ea3b648a15383e7354f0a43d2fe3f6
2 files changed
+32
-3
.buildkite/test-pipeline.yaml
+3
-3
@@ -1313,11 +1313,11 @@ steps:
1313
working_dir: "/vllm-workspace/"
1314
num_gpus: 2
1315
commands:
1316
- - pytest -v -s tests/compile/distributed/test_async_tp.py
1316
+ - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
1317
- pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
1318
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
1319
- - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1320
- - pytest -v -s tests/distributed/test_sequence_parallel.py
1319
+ - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
1320
+ - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
1321
- pytest -v -s tests/distributed/test_context_parallel.py
1322
- CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048
1323
- pytest -v -s tests/v1/distributed/test_dbo.py
tests/conftest.py
+29
@@ -1424,3 +1424,32 @@ def disable_deepgemm_ue8m0(monkeypatch):
1424
# Clear cache so the next time it is used it is processed with the
1425
# default VLLM_USE_DEEP_GEMM_E8M0 setting.
1426
is_deep_gemm_e8m0_used.cache_clear()
1427
+
1428
+
1429
+@pytest.fixture(autouse=True)
1430
+def clean_gpu_memory_between_tests():
1431
+ if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
1432
+ yield
1433
+ return
1434
+
1435
+ # Wait for GPU memory to be cleared before starting the test
1436
+ import gc
1437
+
1438
+ from tests.utils import wait_for_gpu_memory_to_clear
1439
+
1440
+ num_gpus = torch.cuda.device_count()
1441
+ if num_gpus > 0:
1442
+ try:
1443
+ wait_for_gpu_memory_to_clear(
1444
+ devices=list(range(num_gpus)),
1445
+ threshold_ratio=0.1,
1446
+ )
1447
+ except ValueError as e:
1448
+ logger.info("Failed to clean GPU memory: %s", e)
1449
+
1450
+ yield
1451
+
1452
+ # Clean up GPU memory after the test
1453
+ if torch.cuda.is_available():
1454
+ torch.cuda.empty_cache()
1455
+ gc.collect()