AI导哈

PyTorch内存快照分析工具

MemSnap MCP Server 是一个用于分析存储在SQLite数据库中的PyTorch内存快照的工具，支持从pickle快照文件转换并提供多种内存分析功能。

  • by codeasier
  • Python
  • 本地部署
  • 更新于
在 GitHub 查看 codeasier/memsnap-mcp

使用说明

MemSnap MCP Server

MCP server for PyTorch memory snapshot analysis.

English | 中文

Overview

This server provides tools to analyze PyTorch memory snapshots stored in SQLite databases. It works with data converted from pickle snapshot files using the MemSnapDump toolkit.

Installation

pip install -e .

Usage

1. Convert Pickle to SQLite

First, convert your PyTorch memory snapshot pickle file to SQLite using MemSnapDump:

python dump2db.py snapshot.pkl -o ./output/

This will generate a .db file containing:

  • trace_entry_{device} - Memory events (alloc, free, segment_alloc, etc.)
  • block_{device} - Memory block states

2. Configure MCP Client

Add to your MCP client configuration:

{
  "mcpServers": {
    "memsnap-analysis": {
      "command": "python",
      "args": ["-m", "memsnap_analysis_server.server"],
      "env": {
        "PATH": "/path/to/your/python/env/bin"
      }
    }
  }
}

Available Tools

Tool Description
load_database Load a MemSnap SQLite database file
list_devices List all available devices in the loaded database
get_memory_summary Get memory allocation summary (max reserved/allocated/active)
get_event_stats Get statistics of events by action type
query_events Query events with SQL WHERE clause
query_blocks Query blocks with SQL WHERE clause
get_top_allocations Get top N largest memory allocations
get_callstack_analysis Analyze memory by callstack
get_memory_timeline Get memory usage timeline from events

Event Types

Value Action
0 segment_map
1 segment_unmap
2 segment_alloc
3 segment_free
4 alloc
5 free_requested
6 free_completed
7 workspace_snapshot

Block States

Value State
-1 inactive
0 active_pending_free
1 active_allocated

Example Queries

-- Query large allocations
SELECT * FROM trace_entry_0 WHERE size > 1048576

-- Query active blocks
SELECT * FROM block_0 WHERE state = 1

-- Query specific event type
SELECT * FROM trace_entry_0 WHERE action = 4

Database Schema

trace_entry_{device}

Column Type Description
id int Event ID (primary key)
action int Action type
address int Memory address
size int Size in bytes
stream int Stream ID
allocated int Total allocated memory
active int Active memory
reserved int Reserved memory
callstack text Call stack string

block_{device}

Column Type Description
id int Block ID (primary key)
address int Memory address
size int Size in bytes
requestedSize int Requested size
state int Block state
allocEventId int Allocation event ID
freeEventId int Free event ID

License

MIT

相关 MCP