[CI][Test] Deflake test_mem.py sleep-mode asserts via allocator bookeeping (#54312)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Oxana Korzh committed
Aug 29, 2026 at 13:05 UTC
6cddad414ee46796f21aaf7b8643a6e7a00c09b5
1 file changed
+28
-9
tests/basic_correctness/test_mem.py
+28
-9
@@ -18,6 +18,21 @@ from ..utils import create_new_process_for_each_test, requires_fp8
18
DEVICE_TYPE = current_platform.device_type
19
20
21
+def mapped_usage(allocator) -> int:
22
+ """Bytes the allocator still has mapped on the device.
23
+
24
+ `torch.accelerator.get_memory_info()` is a device-wide `cudaMemGetInfo` /
25
+ `hipMemGetInfo` query, so its delta also carries whatever else on the GPU
26
+ allocates or frees while we measure. The allocator's own bookkeeping is
27
+ process-local and exact.
28
+ """
29
+ return sum(
30
+ data.handle[1]
31
+ for data in allocator.pointer_to_data.values()
32
+ if not data.is_asleep
33
+ )
34
+
35
+
36
def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None:
37
"""Wake discarded allocations with deterministic nonzero contents."""
38
original_create_and_map = cumem.create_and_map
@@ -82,11 +97,13 @@ def test_basic_cumem():
97
output = x + y + z
98
assert torch.allclose(output, torch.ones_like(output) * 3)
99
85
- free_bytes = torch.accelerator.get_memory_info()[0]
100
+ # Track the allocator's own mapped bytes, not device-wide free memory.
101
+ mapped_bytes = mapped_usage(allocator)
102
+ assert mapped_bytes >= y.nbytes + z.nbytes
103
allocator.sleep()
87
- free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]
88
- assert free_bytes_after_sleep > free_bytes
104
+ assert mapped_usage(allocator) == 0
105
allocator.wake_up()
106
+ assert mapped_usage(allocator) == mapped_bytes
107
108
# they can be used together
109
output = x + y + z
@@ -105,13 +122,14 @@ def test_discard_tags():
122
with allocator.use_memory_pool("kv_cache"):
123
kv = torch.ones(512, 512, device=DEVICE_TYPE)
124
108
- free_bytes = torch.accelerator.get_memory_info()[0]
125
+ mapped_bytes = mapped_usage(allocator)
126
127
# Discard kv_cache only — weights should remain valid
128
allocator.discard("kv_cache")
129
113
- free_bytes_after_discard = torch.accelerator.get_memory_info()[0]
114
- assert free_bytes_after_discard > free_bytes
130
+ mapped_bytes_after_discard = mapped_usage(allocator)
131
+ assert mapped_bytes - mapped_bytes_after_discard >= kv.nbytes
132
+ assert mapped_bytes_after_discard >= weights.nbytes
133
134
# Weights are still usable
135
assert torch.allclose(weights, torch.ones_like(weights))
@@ -172,11 +190,12 @@ def test_cumem_with_cudagraph():
190
with torch.cuda.graph(model_graph):
191
y = model(x)
192
175
- free_bytes = torch.accelerator.get_memory_info()[0]
193
+ mapped_bytes = mapped_usage(allocator)
194
+ assert mapped_bytes >= weight.nbytes + cache.nbytes
195
allocator.sleep()
177
- free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]
178
- assert free_bytes_after_sleep > free_bytes
196
+ assert mapped_usage(allocator) == 0
197
allocator.wake_up()
198
+ assert mapped_usage(allocator) == mapped_bytes
199
200
# after waking up, the content in the weight tensor
201
# should be restored, but the content in the cache tensor