[ROCm][CI] Make memory sampling less racy in tests and sleep mode (#45490)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Codex <codex@example.invalid> Co-authored-by: Codex <codex@example.invalid>
Andreas Karatzas committed
Jun 29, 2026 at 16:26 UTC
4eb227992aa2231ad538b8c90bc8191397ba3697
2 files changed
+59
-7
tests/utils.py
+46
-4
@@ -1501,6 +1501,9 @@ def wait_for_gpu_memory_to_clear(
1501
threshold_bytes: int | dict[int, int] | None = None,
1502
threshold_ratio: float | dict[int, float] | None = None,
1503
timeout_s: float = 120,
1504
+ stable_duration_s: float = 0,
1505
+ stable_tolerance_bytes: int = 512 * 1024**2,
1506
+ poll_interval_s: float = 5,
1507
) -> None:
1508
assert threshold_bytes is not None or threshold_ratio is not None
1509
devices = get_physical_device_indices(devices)
@@ -1528,8 +1531,13 @@ def wait_for_gpu_memory_to_clear(
1531
# Use nvml instead of pytorch to reduce measurement error from torch cuda
1532
# context.
1533
start_time = time.time()
1534
+ stable_since: float | None = None
1535
+ stable_used_bytes: dict[int, int] | None = None
1536
while True:
1537
output_raw = record_gpu_memory_usage_stats(devices=devices)
1538
+ used_bytes_by_device = {
1539
+ device: int(gb_used * 2**30) for device, (gb_used, _) in output_raw.items()
1540
+ }
1541
output = {
1542
device: f"{gb_used:.02f}/{gb_total:.02f}"
1543
for device, (gb_used, gb_total) in output_raw.items()
@@ -1577,15 +1585,45 @@ def wait_for_gpu_memory_to_clear(
1585
1586
dur_s = time.time() - start_time
1587
if all_free:
1580
- print(f"Done waiting for free GPU memory on ({threshold=}) {dur_s=:.02f}")
1581
- break
1588
+ if stable_duration_s <= 0:
1589
+ print(
1590
+ f"Done waiting for free GPU memory on devices {devices=} "
1591
+ f"({threshold=}) {dur_s=:.02f}"
1592
+ )
1593
+ break
1594
+
1595
+ now = time.time()
1596
+ if stable_used_bytes is None:
1597
+ stable_since = now
1598
+ stable_used_bytes = used_bytes_by_device
1599
+ else:
1600
+ memory_changed = any(
1601
+ abs(used_bytes_by_device[device] - stable_used_bytes[device])
1602
+ > stable_tolerance_bytes
1603
+ for device in devices
1604
+ )
1605
+ if memory_changed:
1606
+ stable_since = now
1607
+ stable_used_bytes = used_bytes_by_device
1608
+ elif (
1609
+ stable_since is not None and now - stable_since >= stable_duration_s
1610
+ ):
1611
+ print(
1612
+ f"Done waiting for stable free GPU memory on devices "
1613
+ f"{devices=} ({threshold=}) {dur_s=:.02f}"
1614
+ )
1615
+ break
1616
+ else:
1617
+ stable_since = None
1618
+ stable_used_bytes = None
1619
1620
if dur_s >= timeout_s:
1621
raise ValueError(
1585
- f"Memory of devices not free after {dur_s=:.02f} ({threshold=})"
1622
+ f"Memory of devices {devices=} not free after "
1623
+ f"{dur_s=:.02f} ({threshold=})"
1624
)
1625
1588
- time.sleep(5)
1626
+ time.sleep(poll_interval_s)
1627
1628
1629
def wait_for_rocm_memory_to_settle(
@@ -1606,11 +1644,15 @@ def wait_for_rocm_memory_to_settle(
1644
num_gpus = current_platform.device_count()
1645
if num_gpus == 0:
1646
return
1647
+ if threshold_ratio is None:
1648
+ threshold_ratio = 0.1
1649
1650
wait_for_gpu_memory_to_clear(
1651
devices=list(range(num_gpus)),
1652
threshold_ratio=threshold_ratio,
1653
timeout_s=timeout_s,
1654
+ stable_duration_s=2.0,
1655
+ poll_interval_s=1.0,
1656
)
1657
1658
vllm/v1/worker/gpu_worker.py
+13
-3
@@ -4,6 +4,7 @@
4
5
import gc
6
import os
7
+import time
8
from collections.abc import Callable
9
from contextlib import AbstractContextManager, contextmanager, nullcontext
10
from datetime import timedelta
@@ -170,7 +171,8 @@ class Worker(WorkerBase):
171
self._pp_send_work: list[Handle] = []
172
173
def sleep(self, level: int = 1) -> None:
173
- free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
174
+ torch.accelerator.synchronize()
175
+ free_bytes_before_sleep = current_platform.mem_get_info()[0]
176
177
# Save the buffers before level 2 sleep
178
if level == 2:
@@ -181,8 +183,16 @@ class Worker(WorkerBase):
183
184
allocator = get_mem_allocator_instance()
185
allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
184
- free_bytes_after_sleep, total = torch.cuda.mem_get_info()
185
- freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
186
+
187
+ torch.accelerator.synchronize()
188
+ deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
189
+ while True:
190
+ free_bytes_after_sleep, total = current_platform.mem_get_info()
191
+ freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
192
+ if freed_bytes >= 0 or time.monotonic() >= deadline:
193
+ break
194
+ time.sleep(0.1)
195
+
196
used_bytes = total - free_bytes_after_sleep
197
assert freed_bytes >= 0, "Memory usage increased after sleeping."
198
logger.info(