mem_snap.py

"""Quick CUDA memory snapshot after window 1 encode_video to find the 44GB holder."""
import sys, os, gc, json, pickle, gzip
sys.path.insert(0, '/workspace/LTX-2/packages/ltx-pipelines/src')
sys.path.insert(0, '/workspace/LTX-2/packages/ltx-core/src')

import torch

# Find all GPU tensors >= 1GB, print shape + data_ptr + storage_size
def dump_big_gpu_tensors(label):
    print(f"\n=== {label} === VRAM={torch.cuda.memory_allocated()//1024**3}GB ===", flush=True)
    big = []
    for obj in gc.get_objects():
        try:
            if isinstance(obj, torch.Tensor) and obj.device.type == 'cuda':
                gb = obj.untyped_storage().nbytes() / 1024**3
                if gb > 1.0:
                    big.append((gb, obj.shape, hex(obj.data_ptr()), obj.dtype))
        except Exception:
            pass
    big.sort(reverse=True)
    total = sum(g for g, *_ in big)
    print(f"  {len(big)} tensors >= 1GB, total={total:.1f}GB", flush=True)
    for gb, shape, ptr, dt in big[:10]:
        print(f"  {gb:.2f}GB {list(shape)} {dt} ptr={ptr}", flush=True)
    return total

# Use torch's memory snapshot
def take_snapshot(path):
    torch.cuda.synchronize()
    snap = torch.cuda.memory._snapshot()
    # find largest active segments
    segments = snap.get('segments', [])
    active = []
    for seg in segments:
        for blk in seg.get('blocks', []):
            if blk.get('state') == 'active_allocated':
                sz = blk.get('size', 0)
                if sz > 500 * 1024**2:  # > 500MB
                    trace = blk.get('frames', [])
                    trace_str = ' -> '.join(f"{f.get('name','?')}({f.get('filename','?')}:{f.get('line','?')})" for f in trace[-5:])
                    active.append((sz / 1024**3, trace_str))
    active.sort(reverse=True)
    print(f"\n=== SNAPSHOT {path}: {len(active)} blocks >= 500MB ===", flush=True)
    for gb, trace in active[:20]:
        print(f"  {gb:.2f}GB | {trace}", flush=True)