MemSnapDump is a toolkit for replaying, slicing, and exporting memory snapshots collected from torch / torch_npu.
- Replay snapshot event history to reconstruct allocator state changes
- Slice large snapshot files into smaller pieces for focused inspection
- Export snapshot data into SQLite for downstream analysis and tooling
pip install memsnapdumpFor development:
python -m pip install -e .[dev]Check the available commands:
memsnapdump -h
memsnapdump --help
memsnapdump --version
memsnapdump split -h
memsnapdump dump2db -hmemsnapdump split /data/snapshot.pickle --slices 4memsnapdump dump2db /data/snapshot.pickle -o /data/outputFor memory snapshot replay, the Python API is intended primarily for extension and custom development rather than quick CLI-style usage. You can register custom hooks into the replay process to add your own statistics, validation, export, or observability logic while allocator state is reconstructed.
from pathlib import Path
from memsnapdump.simulate import SimulateDeviceSnapshot, SimulateHooker
from memsnapdump.util.file_util import load_pickle_to_dict
class EventCounterHooker(SimulateHooker):
def __init__(self):
self.count = 0
def pre_undo_event(self, wait4undo_event, current_snapshot) -> bool:
self.count += 1
return True
def post_undo_event(self, already_undo_event, current_snapshot) -> bool:
return True
snapshot_dict = load_pickle_to_dict(Path("tests/test_data/snapshot_expandable.pkl"))
snapshot = SimulateDeviceSnapshot(snapshot_dict, 0)
hooker = EventCounterHooker()
snapshot.register_hooker(hooker)
snapshot.replay()
print(f"replayed events: {hooker.count}")- Getting started
- Replay guide
- Split guide
- Dump-to-database guide
- SQLite schema reference
- Development guide
ruff check .
black --check .
pytest --cov=memsnapdump --cov-fail-under=85Contributions are welcome via issues and pull requests.