primus-lens-wandb-exporter 0.1.2

自动劫持 WandB 上报的工具，无需修改用户代码

Primus Lens WandB Exporter

自动劫持 WandB 上报的工具 - 零代码侵入，pip install 即用

✨ 核心特性

• 🚀 零代码修改：pip install 后自动生效，训练代码无需任何改动
• 📊 自动增强：自动为 wandb.log() 添加系统级指标（CPU、内存、GPU等）
• 💾 本地备份：自动保存所有 wandb 指标到本地文件
• 🌐 API 异步上报：自动上报框架检测和训练指标到 telemetry-processor（不阻塞训练）
• 🎯 智能框架检测：采集环境变量、WandB 配置等多源证据，生成预判断 hints
• 🔧 灵活配置：通过环境变量控制所有功能
• 🎯 分布式支持：完美支持多节点、多卡训练场景

🎯 快速开始

一键安装（推荐）

使用一键安装脚本，自动完成包安装和配置：

# 方法1: 直接通过 curl 安装
curl -fsSL https://raw.githubusercontent.com/your-repo/main/install.sh | bash

# 方法2: 通过 wget 安装
wget -qO- https://raw.githubusercontent.com/your-repo/main/install.sh | bash

# 方法3: 下载后运行
wget https://raw.githubusercontent.com/your-repo/main/install.sh
bash install.sh

在训练启动脚本中使用：

#!/bin/bash
# train_launch.sh

# 一键安装（包含 pip install + .pth 文件创建 + 验证）
curl -fsSL https://raw.githubusercontent.com/your-repo/main/install.sh | bash

# 设置环境变量
export PRIMUS_LENS_WANDB_DEBUG=true
export PRIMUS_LENS_WANDB_API_REPORTING=true
export PRIMUS_LENS_API_BASE_URL=http://your-api-endpoint

# 启动训练（hook 自动生效）
python3 train.py

这个脚本会：

1. ✅ 检测 Python 环境
2. ✅ 安装 primus-lens-wandb-exporter 包
3. ✅ 自动创建 .pth 文件
4. ✅ 验证安装成功
5. ✅ 提供详细的使用说明

手动安装（备用）

如果需要手动控制安装过程：

# 1. 安装包
pip install primus-lens-wandb-exporter

# 2. 创建 .pth 文件
python3 -c "import site,os; open(os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth'), 'w').write('import primus_lens_wandb_exporter.wandb_hook\n')"

# 如果权限不足，使用 sudo
sudo python3 -c "import site,os; open(os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth'), 'w').write('import primus_lens_wandb_exporter.wandb_hook\n')"

使用

你的训练代码完全不需要修改：

import wandb

# 正常使用 wandb - Primus Lens 会自动劫持
wandb.init(project="my-project")

for epoch in range(10):
    # wandb.log 会被自动增强
    wandb.log({"loss": 0.5, "accuracy": 0.9})

运行训练脚本：

# 设置环境变量（可选）
export PRIMUS_LENS_WANDB_HOOK=true                # 启用劫持（默认启用）
export PRIMUS_LENS_WANDB_ENHANCE_METRICS=true     # 添加系统指标
export PRIMUS_LENS_WANDB_OUTPUT_PATH=/tmp/metrics # 本地保存路径
export PRIMUS_LENS_WANDB_SAVE_LOCAL=true          # 保存到本地

# API 上报配置（可选，在 Kubernetes 环境中自动注入）
export WORKLOAD_UID="your-workload-uid"           # Workload 标识
export POD_UID="your-pod-uid"                     # Pod 标识
export PRIMUS_LENS_WANDB_API_REPORTING=true       # 启用 API 上报
export PRIMUS_LENS_API_BASE_URL="http://primus-lens-telemetry-processor:8080/api/v1"

# 运行训练脚本（无需修改）
python train.py

🔧 配置选项

通过环境变量控制行为：

基础配置

环境变量	默认值	说明
PRIMUS_LENS_WANDB_HOOK	true	是否启用 WandB 劫持
PRIMUS_LENS_WANDB_DEBUG	false	是否启用 debug 日志输出
PRIMUS_LENS_WANDB_ENHANCE_METRICS	false	是否自动添加系统指标（CPU、内存、GPU）
PRIMUS_LENS_WANDB_SAVE_LOCAL	true	是否保存指标到本地文件
PRIMUS_LENS_WANDB_OUTPUT_PATH	None	本地指标输出路径

API 上报配置 (NEW!)

环境变量	默认值	说明
PRIMUS_LENS_WANDB_API_REPORTING	true	是否启用 API 异步上报
PRIMUS_LENS_API_BASE_URL	http://primus-lens-telemetry-processor:8080/api/v1	API 基础 URL
WORKLOAD_UID	必需	Workload 唯一标识（由 Adapter 注入）
POD_UID	必需	Pod 唯一标识

详细配置请参考：API 上报文档

分布式训练配置

支持以下环境变量自动识别节点和进程：

• NODE_RANK / SLURM_NODEID / PET_NODE_RANK - 节点 rank
• RANK / LOCAL_RANK / WORLD_SIZE - 进程 rank 信息

🛠️ 工作原理

1. 安装时：pip install 会在 site-packages 创建 primus_lens_wandb_hook.pth 文件
2. Python 启动时：自动读取并执行 .pth 文件，注册 Import Hook
3. 导入劫持：当用户 import wandb 时，Import Hook 触发
4. Monkey Patching：自动 patch wandb.init() 和 wandb.log() 方法
5. 透明增强：原有功能保持不变，同时添加额外能力

pip install 
    ↓
创建 .pth 文件
    ↓
Python 启动时自动导入
    ↓
注册 Import Hook
    ↓
import wandb 触发劫持
    ↓
自动 Patch wandb
    ↓
透明增强 wandb 功能

📁 输出格式

当设置 PRIMUS_LENS_WANDB_OUTPUT_PATH 时，指标会保存到：

{OUTPUT_PATH}/
  └── node_0/
      ├── rank_0/
      │   └── wandb_metrics.jsonl
      ├── rank_1/
      │   └── wandb_metrics.jsonl
      └── ...

每行 JSON 格式：

{
  "timestamp": 1700000000.123,
  "step": 10,
  "data": {
    "loss": 0.5,
    "accuracy": 0.9,
    "_primus_lens_enabled": true,
    "_primus_sys_cpu_percent": 45.2,
    "_primus_sys_memory_percent": 60.5
  }
}

🧪 验证安装

检查 .pth 文件是否已安装：

# 检查 .pth 文件
python3 -c "import site,os; pth=os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth'); print(f'PTH 文件: {pth}'); print(f'存在: {os.path.exists(pth)}')"

# 使用提供的检查脚本
python install_hook.py check

测试劫持是否生效：

import os
os.environ['PRIMUS_LENS_WANDB_DEBUG'] = 'true'

import wandb

# 检查是否被 patch
print(f"WandB patched: {hasattr(wandb, '_primus_lens_patched')}")

🚀 高级用法

仅保存到本地，不上报到 wandb

# 设置 wandb 为 offline 模式
export WANDB_MODE=offline
export PRIMUS_LENS_WANDB_SAVE_LOCAL=true
export PRIMUS_LENS_WANDB_OUTPUT_PATH=/data/metrics

python train.py

多节点训练示例

# 节点 0
export NODE_RANK=0
export PRIMUS_LENS_WANDB_OUTPUT_PATH=/shared/metrics
python -m torch.distributed.launch train.py

# 节点 1
export NODE_RANK=1
export PRIMUS_LENS_WANDB_OUTPUT_PATH=/shared/metrics
python -m torch.distributed.launch train.py

临时禁用劫持

export PRIMUS_LENS_WANDB_HOOK=false
python train.py  # wandb 正常工作，不会被劫持

启用 Debug 日志

# 启用 debug 日志以查看详细的运行信息
export PRIMUS_LENS_WANDB_DEBUG=true
python train.py

# 禁用所有日志（默认，保持训练输出清爽）
export PRIMUS_LENS_WANDB_DEBUG=false
python train.py

🗑️ 卸载

# 方法 1：使用脚本
python install_hook.py uninstall

# 方法 2：手动删除
python -c "import site, os; pth=os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth'); os.remove(pth) if os.path.exists(pth) else None"

# 方法 3：临时禁用（不删除）
export PRIMUS_LENS_WANDB_HOOK=false

📊 性能影响

• 启动开销：< 10ms（加载 .pth 文件）
• 每次 log 开销：< 1ms（dict.copy() + 指标收集）
• 对训练影响：< 0.1%

🤝 与其他 Primus Lens 模块集成

WandB Exporter 可以与其他 Primus Lens 模块配合使用：

• workload-exporter：收集 GPU/RDMA 底层指标
• node-exporter：收集节点级别资源
• network-exporter：收集网络流量

🔧 故障排除

问题1：Hook 未生效，没有看到日志

症状：训练代码运行正常，但看不到任何 Primus Lens 的日志输出。

原因：.pth 文件未创建。

解决方法：

# 重新运行一键安装脚本
curl -fsSL https://raw.githubusercontent.com/your-repo/main/install.sh | bash

# 或手动创建 .pth 文件
python install_hook.py install

# 或使用一行命令
python3 -c "import site,os; open(os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth'), 'w').write('import primus_lens_wandb_exporter.wandb_hook\n')"

验证修复：

# 检查 .pth 文件
python3 -c "import site, os; print(os.path.exists(os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth')))"
# 应该输出：True

问题2：权限不足

症状：看到 "Permission denied" 错误。

解决方法：

# 使用 sudo
curl -fsSL https://raw.githubusercontent.com/your-repo/main/install.sh | sudo bash

# 或手动创建（使用 sudo）
sudo python3 -c "import site,os; open(os.path.join(site.getsitepackages()[0], 'primus_lens_wandb_hook.pth'), 'w').write('import primus_lens_wandb_exporter.wandb_hook\n')"

问题3：修改 .pth 文件后仍不生效

原因：.pth 文件只在 Python 启动时读取。

解决方法：重新启动 Python 进程（重新运行训练脚本）。

问题4：如何确认 Hook 已加载？

# 启用 debug 日志
export PRIMUS_LENS_WANDB_DEBUG=true

# 运行训练脚本，应该看到：
# [Primus Lens WandB] Installing WandB hook...
# [Primus Lens WandB] WandB successfully patched!
# [Primus Lens WandB] Intercepted wandb.init()

⚠️ 注意事项

1. 虚拟环境：需要在每个虚拟环境中单独安装和创建 .pth 文件
2. 权限问题：可能需要 sudo/管理员权限来创建 .pth 文件
3. Import 顺序：必须在导入 wandb 之前完成 hook 注册（通过 .pth 自动完成）
4. Python 重启：创建/修改 .pth 文件后必须重启 Python 进程
5. 兼容性：已测试 wandb >= 0.12.0

📖 更多文档

• 安装指南 - 详细安装说明和故障排除
• API 上报文档 - 异步 API 上报功能详解（NEW!）
• 使用指南 - 中文详细使用说明
• 示例脚本 - API 上报示例代码
• Megatron 集群测试 - 集群环境上报能力测试指南（NEW!）

🧪 集群测试示例

我们提供了完整的 Megatron-LM 分布式训练模拟脚本，用于在实际集群中测试 API 上报能力：

快速测试（5 分钟）

# 快速验证功能
bash quickstart_megatron.sh

单节点测试

# 单机多卡测试
export WORKLOAD_UID="your-workload-uid"
export POD_UID="your-pod-uid"
export PRIMUS_LENS_API_BASE_URL="http://telemetry-processor:8080/api/v1"
bash launch_megatron_single_node.sh

多节点集群测试

# 在每个节点上运行（注意设置不同的 NODE_RANK）
export WORKLOAD_UID="your-workload-uid"
export NODE_RANK=0  # 节点 0, 1, 2, ...
export MASTER_ADDR="node-0"
export NNODES=4
bash launch_megatron_cluster.sh

Kubernetes 部署

# 修改 megatron_k8s.yaml 中的配置后部署
kubectl apply -f megatron_k8s.yaml
kubectl logs -f megatron-master

详细说明请参考：Megatron 集群测试指南

📄 许可证

Apache-2.0