trajectory_tracker.py — 这个文件实现了轨迹追踪器（Trajectory Tracker）

模块路径: verl.utils.debug.trajectory_tracker

文件概述

这个文件实现了轨迹追踪器（Trajectory Tracker），用于在训练过程中保存中间结果到 HDFS，便于离线比较和调试。每个进程通过 Ray Actor 将数据异步上传到 HDFS。

需要通过环境变量 VERL_ENABLE_TRACKER=1 启用。

关键代码讲解

1. HDFS 保存函数

@ray.remote
def save_to_hdfs(data: io.BytesIO, name, hdfs_dir, verbose):
    """将数据保存到 HDFS（Ray remote 任务）"""
    filename = name + ".pth"
    with tempfile.TemporaryDirectory() as tmpdirname:
        local_filepath = os.path.join(tmpdirname, filename)
        with open(local_filepath, "wb") as f:
            f.write(data.getbuffer())
        copy(local_filepath, hdfs_dir)  # 上传到 HDFS

2. TrajectoryTracker Actor

@ray.remote
class TrajectoryTracker:
    """Ray Actor：管理异步数据上传"""
    def __init__(self, hdfs_dir, verbose):
        self.hdfs_dir = hdfs_dir
        makedirs(hdfs_dir)
        self.handle = deque()  # 未完成的上传任务队列

    def dump(self, data, name):
        self.handle.append(save_to_hdfs.remote(data, name, self.hdfs_dir, self.verbose))

    def wait_for_hdfs(self):
        """等待所有上传任务完成"""
        while len(self.handle) != 0:
            future = self.handle.popleft()
            ray.get(future)

3. 便捷接口

def dump_data(data, name):
    """保存数据的便捷函数（需要环境变量启用）"""
    enable = os.getenv("VERL_ENABLE_TRACKER", "0") == "1"
    if not enable:
        return
    buffer = io.BytesIO()
    torch.save(data, buffer)
    tracker = get_trajectory_tracker()
    ray.get(tracker.dump.remote(buffer, name))

def get_trajectory_tracker():
    """获取或创建全局 Tracker Actor（单例模式）"""
    tracker = TrajectoryTracker.options(
        name="global_tracker", get_if_exists=True, lifetime="detached"
    ).remote(hdfs_dir, verbose)
    return tracker

get_if_exists=True 确保多次调用获取同一个 Actor 实例。

核心类/函数列表

类/函数名	作用
TrajectoryTracker	Ray Actor，管理异步数据上传
dump_data	保存数据的便捷函数
get_trajectory_tracker	获取/创建全局 Tracker
save_to_hdfs	HDFS 上传的 Ray task

与其他模块的关系

• 使用 verl.utils.hdfs_io 的 copy 和 makedirs 操作 HDFS
• 使用 Ray Actor 实现异步上传
• 通过环境变量 VERL_ENABLE_TRACKER 控制开关

小结

这是一个调试/分析工具，在训练中异步保存中间数据到 HDFS。Ray Actor 的使用使得上传不会阻塞训练。deque 管理未完成的任务，wait_for_hdfs 确保训练结束前所有数据都已上传。默认关闭，需要环境变量显式开启。