# Memory Parser 开发者指南

本文面向 Memory Parser 的开发者，介绍其设计思路、内部接口约定以及扩展开发指南。用户侧快速入门见 [Memory Parser Quickstart](../overview/memory_parser_quickstart.md)。

## 1. 设计概述

### 1.1 背景与目标

RL-Insight 已有时序分析 Parser（`mstx` / `torch`），基于 `EventRow` 数据模型输出算子级时间线事件。Memory Parser 新增**内存分析**能力，从 Ascend Profiler 输出的内存相关文件中提取内存分配信息，为 RL 训练的内存瓶颈分析提供数据支撑。

**核心目标**：在每条内存申请记录中补全以下信息：

- 调用栈（Call Stack）—— 定位内存申请源头
- 内存申请大小（Size）—— 量化内存消耗
- 内存申请时间（Allocation Time）—— 时序关联
- 内存占用时长（Duration）—— 生命周期分析

### 1.2 数据流

```
InputData (ASCEND_MEMORY)
    │
    ├── DataChecker 校验
    │
    ▼
MemoryClusterParser
    │
    ├── allocate_prof_data()     ← 扫描目录，构建 DataMap
    │
    ├── mapper_func()            ← 多进程并行调度
    │   └── parse_analysis_data()
    │       ├── _build_call_stack_index()  ← 流式解析 trace_view.json
    │       └── _parse_operator_memory()   ← 解析 operator_memory.csv
    │           └── _match_call_stack()    ← 二分查找调用栈
    │
    ├── reducer_func()           ← 汇总排序
    │
    ▼
OutputData (SUMMARY_MEMORY_EVENT) → pd.DataFrame
```

### 1.3 类结构

```
BaseClusterParser (rl_insight/parser/parser.py)
  └── MemoryClusterParser (rl_insight/parser/memory_parser.py)
        ├── input_type = DataEnum.ASCEND_MEMORY
        ├── allocate_prof_data()          → 扫描目录，构建 DataMap
        ├── parse_analysis_data()         → 主解析流程
        ├── _build_call_stack_index()     → 流式解析 trace_view.json，构建调用栈索引
        ├── _parse_operator_memory()      → 解析 operator_memory.csv，输出 dict[str, Any]
        ├── _match_call_stack()           → name + ts 匹配调用栈
        ├── _get_data_map()               → 构建 (role, rank_id) → [path] 映射
        ├── _get_rank_path_with_role()    → 生成 DataMap 列表
        ├── _get_profiler_data_path()     → 拼接 ASCEND_PROFILER_OUTPUT 路径
        ├── _get_rank_id()                → 从 profiler_info_*.json 提取 rank_id
        ├── _get_task_role()              → 从 profiler_metadata.json 提取 role
        └── _extract_timestamp_key()      → 提取目录名中的时间戳排序键
```

---

## 2. 内部接口

### 2.1 注册与数据类型

```python
@register_cluster_parser("memory")
class MemoryClusterParser(BaseClusterParser):
    input_type: DataEnum = DataEnum.ASCEND_MEMORY
```

| 属性 | 值 | 说明 |
|------|-----|------|
| 注册名 | `"memory"` | CLI `--profiler-type memory` |
| `input_type` | `DataEnum.ASCEND_MEMORY` | 输入数据类型 |
| 输出类型 | `DataEnum.SUMMARY_MEMORY_EVENT` | Parser 输出 / Visualizer 输入 |

`DataEnum` 新增值（定义在 `rl_insight/data/data_checker.py`）：

```python
class DataEnum(Enum):
    ASCEND_MEMORY = "ascend_memory"              # Memory Parser 输入
    SUMMARY_MEMORY_EVENT = "summary_memory_event" # Memory Parser 输出
```

### 2.2 MemoryEventRow

定义在 `rl_insight/utils/schema.py`：

```python
class MemoryEventRow(TypedDict):
    name: str                    # 算子名称
    role: str                    # RL 角色
    rank_id: int                 # 进程 rank
    call_stack: str              # 完整调用栈（";\r\n" 分隔）
    call_stack_top: str          # 调用栈顶层入口
    size_kb: float               # 内存大小（KB），正数=申请，负数=释放
    start_time_ms: float       # 内存申请时间（ms），与 BaseClusterParser.reducer_func 的排序键对齐
    duration_ms: float           # 内存占用时长（ms），0 表示未释放
    total_allocated_mb: float    # 申请时刻累计已分配内存（MB）
    total_reserved_mb: float     # 申请时刻累计预留内存（MB）
    total_active_mb: float       # 申请时刻累计活跃内存（MB）
    device_type: str             # 设备类型
```

**设计决策**：

- `size_kb`：正数表示内存申请，负数表示内存释放，与 `operator_memory.csv` 中的 `Size(KB)` 语义一致
- `duration_ms`：若 `Duration(us)` 有值则转换，无值则为 `0.0`（表示内存尚未释放）
- `call_stack_top`：取调用栈第一行（用户代码入口），便于快速定位，无需解析完整调用栈
- 时间单位统一为毫秒（ms），与现有 `EventRow` 保持一致

### 2.3 allocate_prof_data()

复用 MstxClusterParser 的目录扫描逻辑：

1. 遍历 `input_path`，找到所有 `<date>_<time>_ascend_pt` 目录
2. 从 `profiler_metadata.json` 提取 `role`，从 `profiler_info_*.json` 提取 `rank_id`
3. 按 `_extract_timestamp_key` 提取的时间戳对同 (role, rank_id) 下的目录排序
4. `profiler_data_path` 指向 `ASCEND_PROFILER_OUTPUT` 目录（注意：与 MstxClusterParser 指向 `trace_view.json` 文件不同，Memory Parser 指向目录，因为需要同时访问 `trace_view.json` 和 `operator_memory.csv`）

**关键差异**：

```python
# MstxClusterParser
def _get_profiler_data_path(self, rank_id, data_path):
    return os.path.join(data_path, Constant.ASCEND_PROFILER_OUTPUT, "trace_view.json")

# MemoryClusterParser
def _get_profiler_data_path(self, rank_id, data_path):
    return os.path.join(data_path, Constant.ASCEND_PROFILER_OUTPUT)
```

### 2.4 parse_analysis_data()

主解析流程，接收单个 Rank 的数据路径，返回 `list[dict[str, Any]]`：

```
输入: profiler_data_path (ASCEND_PROFILER_OUTPUT 目录), rank_id, role

步骤1: 构建调用栈索引
       _build_call_stack_index(trace_view_path)
       → 流式解析 trace_view.json (ijson)
       → 筛选 cat=="cpu_op" 且含 "Call stack" 的事件
       → 按 name 分组，组内按 ts 排序
       → 返回 dict[str, list[{ts, dur, call_stack}]]

步骤2: 解析 operator_memory.csv
       _parse_operator_memory(csv_path, call_stack_index, rank_id, role)
       → 逐行构建 MemoryEventRow:
          a. 时间转换: us → ms (÷ 1000)
          b. 调用栈匹配: _match_call_stack(name, allocation_time, index)
             未命中 → call_stack = "", call_stack_top = ""
          c. 提取 call_stack_top: 取调用栈第一行
          d. duration_ms: Duration(us) 有值则转换，无值则为 0

步骤3: 返回 list[dict[str, Any]]
```

### 2.5 _build_call_stack_index()

流式解析 `trace_view.json`，构建调用栈索引：

- **输入**：`trace_view.json` 文件路径
- **输出**：`dict[str, dict]`，每个值包含 `"entries"`（`list[{ts, dur, call_stack}]`，按 `ts` 升序排序）和 `"ts_list"`（预提取的 `ts` 列表，供二分查找直接使用，避免每次调用重建列表）
- **过滤条件**：仅保留 `cat=="cpu_op"` 且 `args` 中含 `"Call stack"` 的事件
- **流式解析**：使用 `ijson.items(f, "item")` 逐条读取，避免将整个 JSON 加载到内存
- **排序**：组内按 `ts` 升序排序，确保后续二分查找的正确性

### 2.6 _match_call_stack()

调用栈匹配算法，基于二分查找：

- **输入**：算子名 `name`、分配时间 `allocation_time_us`、调用栈索引
- **输出**：`(call_stack, call_stack_top)`，未命中返回 `("", "")`
- **匹配策略**：在同名算子组中，查找 `ts ≤ allocation_time` 的最近一条记录
- **匹配语义**：`ts` 是算子开始执行时间，`Allocation Time` 是算子内触发内存分配的时间，因此 `Allocation Time ≥ ts`；取 `ts` 最接近 `Allocation Time` 的一条即为触发该次内存分配的算子调用
- **允许多对一**：一个算子可能分配多次内存，因此多条 `operator_memory` 记录可以匹配到同一条 `trace_view` 记录

```python
import bisect

idx = bisect.bisect_right(ts_list, allocation_time_us) - 1
if idx < 0:
    return "", ""
```

### 2.7 DataChecker 校验

`DataEnum.ASCEND_MEMORY` 当前未注册校验规则（`DataChecker.rules` 中为空列表），输入校验由 Parser 内部的文件存在性检查承担。`DataEnum.SUMMARY_MEMORY_EVENT` 同样未注册校验规则。

如需增加校验，参考 [DataRule 扩展说明](./rule_extending_guide.md)。

---

## 3. 扩展指南

### 3.1 新增内存数据校验规则

适用于：为 `ASCEND_MEMORY` 或 `SUMMARY_MEMORY_EVENT` 增加输入/输出校验。

1. 在 `rl_insight/data/rules.py` 中继承 `ValidationRule`，实现 `check()` 和 `error_message`
2. 在 `DataChecker.rules` 中为新类型挂载规则

示例——为 `ASCEND_MEMORY` 增加 `operator_memory.csv` 存在性校验：

```python
# rl_insight/data/rules.py
class AscendMemoryFileExistsRule(ValidationRule):
    """校验 ASCEND_PROFILER_OUTPUT 下存在 operator_memory.csv"""

    def check(self, data) -> bool:
        if not isinstance(data, str):
            self._error_message = "Data object is not a path"
            return False
        root_path = Path(data)
        if not root_path.exists():
            self._error_message = f"Source path does not exist: {data}"
            return False
        ascend_pt_pattern = str(root_path / "*" / "*_ascend_pt")
        ascend_pt_folders = glob.glob(ascend_pt_pattern)
        for folder in ascend_pt_folders:
            csv_path = Path(folder) / "ASCEND_PROFILER_OUTPUT" / "operator_memory.csv"
            if not csv_path.exists():
                self._error_message = f"operator_memory.csv not found in {folder}"
                return False
        return True

    @property
    def error_message(self) -> str:
        return self._error_message
```

```python
# rl_insight/data/data_checker.py
class DataChecker:
    rules: dict[DataEnum, List[ValidationRule]] = {
        DataEnum.ASCEND_MEMORY: [PathExistsRule(), AscendMemoryFileExistsRule()],
        ...
    }
```

### 3.2 新增 Memory Visualizer

适用于：为 `MemoryEventRow` 输出增加可视化能力。

1. 新增模块 `rl_insight/visualizer/memory_visualizer.py`
2. 继承 `BaseVisualizer`，设置 `input_type = DataEnum.SUMMARY_MEMORY_EVENT`
3. 实现 `run()` 方法，消费 `pd.DataFrame` 生成图表
4. 使用 `@register_cluster_visualizer("<name>")` 注册
5. 更新 `main.py` 中 `--vis-type` 的 help 文本
6. 在 `rl_insight/visualizer/__init__.py` 中导出新类

示例：

```python
# rl_insight/visualizer/memory_visualizer.py
from .visualizer import BaseVisualizer, register_cluster_visualizer
from rl_insight.data import DataEnum

@register_cluster_visualizer("memory_heatmap")
class MemoryHeatmapVisualizer(BaseVisualizer):
    input_type: DataEnum = DataEnum.SUMMARY_MEMORY_EVENT

    def __init__(self, config: dict):
        super().__init__(config)
        self.output_path = config.get("output_path", "output")

    def run(self, data):
        # 实现内存热力图可视化
        # data 为 pd.DataFrame，包含 MemoryEventRow 字段
        ...
```

### 3.3 支持 GPU（CUDA）内存分析

适用于：将 Memory Parser 扩展至支持 CUDA 内存 Profiling 数据。

1. 在 `DataEnum` 中新增 `CUDA_MEMORY = "cuda_memory"`
2. 新增 `CudaMemoryParser`，继承 `BaseClusterParser`，实现 `allocate_prof_data()` 和 `parse_analysis_data()`
3. 使用 `@register_cluster_parser("cuda_memory")` 注册
4. 定义 `CudaMemoryEventRow`（若字段与 `MemoryEventRow` 不同）或复用 `MemoryEventRow`
5. 在 `DataChecker.rules` 中为新类型挂载校验规则
6. 更新 `main.py` 中 `--profiler-type` 的 help 文本
7. 在 `docs/data/data_specification.md` 中补充数据形态说明

### 3.4 改进调用栈匹配精度

当前匹配策略基于算子名 + 时间戳二分查找，在极短时间内同一算子多次调用时可能匹配不精确。可能的改进方向：

1. **增加 `dur` 约束**：匹配时额外要求 `allocation_time ≤ ts + dur`，即分配时间必须在算子执行时间范围内
2. **增加 `tid` 约束**：若 `operator_memory.csv` 可提供线程信息，可按 `name + tid` 分组匹配
3. **最近邻匹配**：在 `ts ≤ allocation_time` 的候选中，选择 `allocation_time - ts` 最小的一条

---

## 4. 依赖

| 库 | 用途 | 说明 |
|----|------|------|
| `ijson` | 流式解析大 JSON | 需安装: `pip install ijson` |
| `csv` | 解析 CSV | 标准库 |
| `bisect` | 二分查找调用栈 | 标准库 |

## 5. 测试

测试文件位于 `tests/parser/test_memory_parser.py`，覆盖以下场景：

- Parser 注册（`"memory"` 在 `CLUSTER_PARSER_REGISTRY` 中）
- `_build_call_stack_index`：过滤 `cpu_op`、过滤无调用栈事件、按 `name` 分组、按 `ts` 排序、空 JSON
- `_match_call_stack`：匹配成功、多条记录取最近、未命中、所有 `ts` 大于 `allocation_time`
- `_parse_operator_memory`：正常解析、空 CSV、调用栈匹配/未匹配
- `allocate_prof_data`：目录扫描、rank_id 提取、role 提取
- E2E：完整解析流程