|
1 | | -"""GPU memory-allocator smoke test — allocate, free, verify no large leak.""" |
| 1 | +"""GPU memory-allocator smoke test. |
| 2 | +
|
| 3 | +TF's BFC allocator deliberately retains pool memory across `clear_session()` / |
| 4 | +`gc.collect()` — `get_memory_info(device)["current"]` reflects allocator-held |
| 5 | +bytes, not user-held bytes. Unlike PyTorch's `torch.cuda.empty_cache()`, TF |
| 6 | +exposes no public API that forces the pool back to the OS at runtime, so a |
| 7 | +"memory after free == memory before alloc" assertion is unreliable. |
| 8 | +
|
| 9 | +Instead, assert the property we actually care about: when a tensor is freed, |
| 10 | +the next allocation of the same size REUSES the existing pool rather than |
| 11 | +growing it. Peak memory across two alloc/free cycles should be flat. |
| 12 | +""" |
2 | 13 |
|
3 | 14 | import gc |
4 | 15 |
|
|
7 | 18 |
|
8 | 19 | pytestmark = pytest.mark.skipif(not tf.config.list_physical_devices("GPU"), reason="GPU only") |
9 | 20 |
|
10 | | -# 50 MB tolerance for runtime overhead; TF caches GPU memory so an exact match isn't expected. |
11 | | -LEAK_TOLERANCE_BYTES = 50 * 1024 * 1024 |
12 | | - |
| 21 | +# 5% growth tolerance between cycles for minor allocator bookkeeping. |
| 22 | +PEAK_GROWTH_TOLERANCE = 1.05 |
13 | 23 |
|
14 | | -def test_gpu_memory_returns_to_baseline(): |
15 | | - device = "GPU:0" |
16 | 24 |
|
17 | | - # Warm up so steady-state allocator overhead is included in the baseline. |
18 | | - warmup = tf.zeros([1024, 1024]) |
19 | | - del warmup |
| 25 | +def _peak_for_cycle(device: str, shape) -> int: |
| 26 | + """Reset peak stats, alloc + free, return the cycle's peak in bytes.""" |
| 27 | + tf.config.experimental.reset_memory_stats(device) |
| 28 | + t = tf.zeros(shape) |
| 29 | + peak = tf.config.experimental.get_memory_info(device)["peak"] |
| 30 | + del t |
20 | 31 | gc.collect() |
| 32 | + return peak |
21 | 33 |
|
22 | | - baseline = tf.config.experimental.get_memory_info(device)["current"] |
23 | 34 |
|
24 | | - # Allocate ~400 MB: 1000 * 1000 * 100 float32 elements. |
25 | | - big = tf.zeros([1000, 1000, 100]) |
26 | | - peak = tf.config.experimental.get_memory_info(device)["current"] |
27 | | - assert peak > baseline, f"allocation did not raise current memory: {baseline} -> {peak}" |
| 35 | +def test_gpu_memory_reuses_pool_across_cycles(): |
| 36 | + device = "GPU:0" |
| 37 | + shape = [1000, 1000, 100] # ~400 MB float32 |
28 | 38 |
|
29 | | - del big |
30 | | - tf.keras.backend.clear_session() |
31 | | - gc.collect() |
| 39 | + peak1 = _peak_for_cycle(device, shape) |
| 40 | + assert peak1 > 0, "allocation did not register in peak memory" |
32 | 41 |
|
33 | | - after = tf.config.experimental.get_memory_info(device)["current"] |
34 | | - leaked = after - baseline |
35 | | - assert leaked <= LEAK_TOLERANCE_BYTES, ( |
36 | | - f"memory not released: baseline={baseline}, after={after}, leaked={leaked} bytes" |
| 42 | + peak2 = _peak_for_cycle(device, shape) |
| 43 | + assert peak2 <= peak1 * PEAK_GROWTH_TOLERANCE, ( |
| 44 | + f"BFC pool grew on identical re-allocation: peak1={peak1}, peak2={peak2} " |
| 45 | + f"(tolerance {PEAK_GROWTH_TOLERANCE}x). Suggests freed memory was not reused." |
37 | 46 | ) |
0 commit comments