开源大数据平台E-MapReduce：基于 EMR Serverless Spark Ray 打造具身智能数据管线

1. 准备 OSS Bucket 和输入视频

准备一个阿里云 OSS Bucket，用于存放输入视频、模型权重、编译缓存和输出数据集。流水线通过 scan_input 扫描指定前缀下的所有 .mp4 文件，也支持本地目录或 .txt 文件列表。OSS 访问走 EMR Ray 集群预装的 pyjindo（JindoOssFileSystem），通过节点 RAM 角色免 AK/SK 认证，与 Ray runtime_env 对 oss:// 协议的处理一致。提交机不再需要配置明文 OSS 凭证。

2. 准备模型文件与第三方源码

下载 HaWoR、mega-sam、GeoCalib、MoGe、pytorch3d 和 utils3d 源码及其子模块，按以下目录结构组织。模型权重不随源码打包，而是在运行时由 oss_model.ensure_model() 从 OSS 按需下载到 worker 本地 weights/ 目录。

vitra/
├── pipeline.py                        # CLI 入口
├── config.yaml                        # 全部配置
├── setup_pip.py                       # 依赖安装脚本（集群冷启动后跑一次）
├── core/orchestrator.py               # 编排逻辑
├── stages/                            # 流水线阶段实现
│   ├── s0_preprocess.py               #   预处理
│   ├── s1_motion_labeling/            #   3D 运动标注 (Ray GPU Actor)
│   ├── s2_segmentation.py             #   动作分割
│   ├── lerobot_output.py              #   LeRobot v3.0 写入器
│   └── visualization/                 #   可视化渲染器
├── ray_pipeline/                      # Ray 运行时层
│   ├── runtime/cluster.py             #   Ray 初始化与生命周期
│   ├── runtime/oss_io.py              #   OSS 读写
│   ├── runtime/oss_model.py           #   模型权重下载 + CUDA 扩展编译
│   ├── manifest/scanner.py            #   输入扫描
│   └── steps/gpu/                     #   GPU 推理步骤
├── modules/                           # 业务逻辑模块
│   ├── hawor_no_filler.py             #   HaWoR 封装
│   ├── megasam.py                     #   MegaSAM 封装
│   └── long_video.py                  #   长视频分段
├── utils/                             # 几何工具
│   └── geometry.py                    #   旋转/速度/投影
├── weights/                           # 模型权重（运行时从 OSS 下载）
│   ├── hawor/                         #   HaWoR 权重 + MANO 资产
│   ├── moge2/                         #   MoGe-2 权重
│   ├── megasam/                       #   MegaSAM 权重
│   └── geocalib/                      #   GeoCalib 权重
└── thirdparty/                        # 第三方源码
    ├── HaWoR/                         #   双手 MANO + DROID-SLAM
    ├── MoGe/                          #   MoGe-2 深度估计
    ├── mega-sam/                      #   MegaSAM 相机追踪
    ├── GeoCalib/                      #   相机内参校准
    ├── pytorch3d/                     #   MANO 渲染（Stage 1 HaWoR）
    └── utils3d/                       #   MoGe 的传递依赖

模型权重在 OSS 上的存放结构（models/ 复数前缀，由 OSS_KEY_PREFIX 控制，默认 vitra/）：

oss://<bucket>/vitra/
├── models/
│   ├── geocalib/pinhole.tar
│   ├── hawor/{hawor.ckpt, detector.pt, model_config.yaml, hawor_DATA.zip}
│   ├── moge2/model.pt
│   └── megasam/megasam_final.pth
├── compiled/
│   └── droid_backends-cp312-cu118-x86_64.tar.gz   # CUDA 扩展编译缓存
└── videos/                                          # 输入视频

模型权重和编译缓存在 worker 上通过节点级持久缓存（/tmp/vitra_cache/）跨 Ray session 复用。setup_persistent_weights() 在每次 Actor 启动时把 repo_root/weights symlink 到 /tmp/vitra_cache/weights/，同节点旧 session 的 weights 目录自动清理，避免约 4.5 GB 模型重复下载占满磁盘。编译产物同理缓存在 /tmp/vitra_cache/compiled/，下次 session 直接解压，不走 OSS。

3. 准备 Python 依赖

Ray 集群镜像已自带 Ray 2.47.1 和 Python 3.12。由于 Ray runtime_env 的 pip 方式会在 worker 上创建 virtualenv（overlay 磁盘容易爆，且屏蔽系统已装的 conda 包），项目改为通过 setup_pip.py 将依赖直接安装到系统 Python。集群冷启动后跑一次即可。

setup_pip.py 安装的核心依赖包括 ultralytics（HaWoR 手部检测）、scipy（动作分割的 find_peaks）、pyarrow/pandas（LeRobot 写出）、opencv-python、imageio-ffmpeg（MP4 编码）。HaWoR、MoGe-2 和 MegaSAM 的 PyTorch/torchvision 等大件由集群 conda 环境提供，不在 pip 列表中。OSS 读写依赖的 pyjindo 由集群镜像预装，无需手动安装。

若干依赖包含 C++/CUDA 扩展，需要针对 Ray worker 的运行环境编译：

必须在 GPU worker 上编译，因为 CUDA 扩展要链接 worker 上确切的 PyTorch / Python / CUDA 版本。编译产物缓存在两层：节点级持久缓存（/tmp/vitra_cache/compiled/，跨 session 复用）和 OSS（compiled/ 前缀，跨集群冷启动复用）。worker 命中节点级缓存时直接解压，秒级完成；不命中则从 OSS 下载并写入持久缓存；OSS 也无缓存则在 worker 上现场编译并回传 OSS。

4. 配置集群

使用阿里云 EMR Serverless Spark Ray 集群，CPU 头节点运行 driver 进程和 CPU 阶段，GPU worker 节点运行三维重建 Actor。源码（vitra/ 目录）通过 ray job submit 的 --working-dir 上传到集群各节点，模型权重和编译缓存在运行时从 OSS 按需下载。

OSS 读写通过 EMR Ray 集群预装的 pyjindo（JindoOssFileSystem）走节点 RAM 角色免 AK/SK 认证，提交机无需配置明文凭证，只需确保 OSS_BUCKET 环境变量正确。

集群入口（RAY_ADDRESS、RAY_TOKEN）和 OSS Bucket（OSS_BUCKET）在 submit_vitra.py 中通过环境变量配置，集群重启或重建后更新即可。其中 RAY_ADDRESS 和 RAY_TOKEN 可在 EMR Serverless Spark 控制台的 Ray 集群详情页获取：进入目标工作空间后，在 Ray 集群列表中单击目标集群，在集群详情页的连接信息中即可查看集群访问地址（对应 RAY_ADDRESS）和访问令牌（对应 RAY_TOKEN），将其填入 submit_vitra.py 对应的环境变量即可。

1. 读取 OSS 视频元数据

流水线通过 scan_input 扫描 OSS 前缀或本地目录，发现待处理的视频文件：

class VideoDiscovery:
    def __init__(self, config: dict):
        self.config = config

    def __call__(self) -> list[str]:
        """Discover input videos from config.

        Supports: directory of .mp4, .txt list, single .mp4, oss:// prefix.
        """
        from manifest.scanner import scan_input

        specific = self.config.get('pipeline', {}).get('specific_videos')
        if specific:
            return specific

        input_path = self.config.get('paths', {}).get('input_videos', '')
        if not input_path:
            raise ValueError('paths.input_videos not set in config')

        _, videos = scan_input(input_path)
        return [str(p) for p in videos]

scan_input 支持 OSS 前缀扫描、本地目录和 .txt 文件列表三种输入方式，返回视频路径列表。OSS 上的视频会在 Stage 0 按需下载到本地。

2. 裁切并规范化视频

预处理算子下载源视频后，按 20 秒进行裁切，相邻片段重叠 2 秒；随后按需去畸变，并将每个片段重采样为 30 FPS。OpenCV 负责逐帧去畸变和缩放，JPEG 帧直接写入磁盘供后续 GPU 阶段复用，避免重复解码视频。长视频会先提取全部帧到共享目录，再按 clip 边界创建符号链接。

一段源视频会产生多个视频片段，因此该阶段按 clip 列表展开。下面为示意代码：

class PreprocessVideo:
    def __init__(self, config: dict):
        self.config = config

    def __call__(self, video_path: str, work_dir: str) -> list[dict]:
        """Run Stage 0 preprocessing for a single video."""
        from runtime.oss_io import ensure_local_video, scene_from_url

        pp = self.config.get('pipeline', {})
        clip_duration_s = pp.get('clip_duration_s', 20.0)
        clip_overlap_s = pp.get('clip_overlap_s', 2.0)
        target_fps = pp.get('target_fps', 30)
        undistort = pp.get('undistort', True)

        local_video = ensure_local_video(video_path)
        scene = scene_from_url(video_path)
        work_dir = Path(work_dir)
        work_dir.mkdir(parents=True, exist_ok=True)

        fps, duration_s, total_frames = _get_video_info(local_video)
        clips = _plan_clips(fps, total_frames, clip_duration_s, clip_overlap_s)

        if len(clips) <= 1:
            # Short video: extract all frames into one directory
            frames_dir = work_dir / 'frames'
            n = _extract_frames(local_video, frames_dir, target_fps, undistort)
            clips[0]['frames_dir'] = str(frames_dir)
            clips[0]['video_path'] = local_video
            clips[0]['fps'] = target_fps
        else:
            # Long video: extract all frames, then create per-clip symlink dirs
            all_frames_dir = work_dir / '_all_frames'
            n_all = _extract_frames(all_frames_dir, target_fps, undistort)
            all_frames_p = Path(all_frames_dir)
            for clip in clips:
                clip_frames_dir = work_dir / f'clip_{clip["clip_idx"]:03d}' / 'frames'
                clip_frames_dir.mkdir(parents=True, exist_ok=True)
                saved = 0
                for src_idx in range(clip['start_frame'], clip['end_frame']):
                    src = all_frames_p / f'{src_idx:06d}.jpg'
                    dst = clip_frames_dir / f'{saved:06d}.jpg'
                    dst.symlink_to(os.path.relpath(src, clip_frames_dir))
                    saved += 1
                clip['frames_dir'] = str(clip_frames_dir)
                clip['video_path'] = local_video
                clip['fps'] = target_fps

        return clips

预处理后的 JPEG 帧保存在 cache/{clip_id}/*.jpg，供 Stage 1 直接读取。

3. 恢复深度、手部姿态和相机轨迹

Stage 1 运行在一个 Ray GPU remote function 中（@ray.remote(num_gpus=1)），为适配 GPU 显存，三个模型串行加载——每个加载下一个前都会卸载（del + torch.cuda.empty_cache()）。driver 在 CPU 头节点上调度，GPU 推理在 worker 上执行：

import ray

@ray.remote(num_gpus=1)
def _remote_stage1(video_path, clips, work_dir, config):
    from stages.s1_motion_labeling.actor import run_stage1
    return run_stage1(video_path, clips, work_dir, config)

motion_label = ray.get(_remote_stage1.remote(
    video_path, clips, str(work_dir), config))

MotionLabelingActor 是一个普通 Python 类（非 Ray Actor），在 worker 进程内实例化后串行执行三个阶段：

class MotionLabelingActor:
    def __init__(self, config: dict):
        self.config = config
        _ensure_models(config)  # 下载模型权重 + 编译 CUDA 扩展
        self.focal = None

    def run(self, video_path, clips, work_dir) -> dict | None:
        local_video = _ensure_local_video(video_path)
        self.focal = _get_focal(local_video, work_dir, self.config)
        if len(clips) == 1:
            return self._run_single_clip(
                local_video, clips[0], self.focal, work_dir)
        else:
            return self._run_multi_clip(
                local_video, clips, self.focal, work_dir)

    def _run_single_clip(self, video_path, clip, focal, work_dir):
        # Phase 1: HaWoR (detect, track, motion estimation)
        detect_meta, fc_map, cs_map, args = _run_hawor_phase(
            video_path, clip['frames_dir'], focal, work_dir)
        # Phase 2: MoGe-2 (depth)
        moge_result = _run_moge2_phase(
            clip['frames_dir'], clip['scene'], work_dir, focal)
        # Phase 3: MegaSAM (camera tracking)
        megasam_result = _run_megasam_phase(
            clip['frames_dir'], clip['scene'], work_dir,
            moge_result['moge_base_dir'], focal, self.config)
        # Finalize: inject camera poses into HaWoR
        save_path = str(Path(work_dir) / 'result' / f'{clip["scene"]}_cleaned.pth')
        pred_result = _finalize_hawor(
            video_path=video_path, image_dir=clip['frames_dir'], focal=focal,
            cam_c2w=megasam_result['cam_c2w'],
            detect_meta=detect_meta, fc_map=fc_map,
            cs_map=cs_map, save_path=save_path)
        return _compose_result(
            pred_result=pred_result,
            cam_c2w=megasam_result['cam_c2w'],
            fov=moge_result['avg_fov'], fps=clip.get('fps', 30),
            frames_dir=clip['frames_dir'], video_path=video_path,
            scene=clip['scene'])

Actor 内部按三个阶段串行执行：

• Phase 1 — HaWoR：在 Actor 进程内运行（modules/hawor_no_filler 通过 load_hawor_model() 加载模型后直接调用，PyTorch3D 由集群 conda 环境提供）。内部链路为 YOLO 手部检测（detector.pt）→ MANO → Masked DROID-SLAM → 深度先验，输出左右手各 T 帧的 HandFrameResult，包含手腕旋转（轴角）、平移、手部姿态和手型参数。

• Phase 2 — MoGe-2：逐帧推理输出深度图、点云图、法线图和 FoV。变通处理包括把 MemEffAttention patch 成 SDPA（规避 xformers），以及 monkey-patch np.savez 将 pyjindo 不支持的 seek() 写入重定向到 /tmp。

• Phase 3 — MegaSAM：在 Actor 进程内三步串行——Depth-Anything 单目视差 → MoGe-2 深度转 UniDepth 格式 → DROID-SLAM 相机追踪，输出 camera_to_world [T,4,4]。pyjindo 变通处理包括把 .so 拷到 /tmp 本地磁盘（pyjindo 无法 mmap/dlopen），并扩展 LD_LIBRARY_PATH。

结果合成为 MotionLabel3D 字典，双手数组补齐到 T = max(left, right, depth_count, 1)，缺失处填零并标记 valid=False：

三个模型分别作为串行加载的独立阶段，便于在 GPU 上运行。显存预算大致为 HaWoR 约 10 GB、MoGe-2 约 6 GB、MegaSAM 约 4 GB。

4. 从视频片段生成原子动作片段

动作分割算子接收 Stage 1 的 MotionLabel3D，计算双手手腕 3D 速度，取 max(left_speed, right_speed) 作为联合信号，用移动平均平滑后，通过 scipy.signal.find_peaks 检测速度极小值（即动作间的停顿），将视频切分为原子动作片段。默认片段长度由最小间距（0.5 秒）和峰值显著度约束。

class ActionSegmentation:
    def __init__(self, config: dict):
        self.config = config

    def __call__(self, motion_label: dict) -> list[dict[str, Any]]:
        """Segment the motion into atomic action clips."""
        from utils.geometry import compute_wrist_speed
        from scipy.signal import find_peaks

        T = motion_label['T']
        fps = motion_label.get('fps', 30)

        left_speed = compute_wrist_speed(
            motion_label['left_wrist_pos'],
            motion_label['left_valid'],
        )
        right_speed = compute_wrist_speed(
            motion_label['right_wrist_pos'],
            motion_label['right_valid'],
        )

        # Joint signal: max of left and right hand speeds
        joint_speed = np.maximum(left_speed, right_speed)

        # Smooth the speed signal with a moving average
        window = max(1, int(fps * 0.3))  # 300ms window
        kernel = np.ones(window) / window
        smoothed = np.convolve(joint_speed, kernel, mode='same')

        # Find speed minima (pauses between actions) using find_peaks on negated signal
        min_distance = max(1, int(fps * 0.5))
        max_speed = float(np.max(smoothed)) if len(smoothed) > 0 else 0.0
        prominence = max(max_speed * 0.05, 1e-6)

        peaks, _ = find_peaks(-smoothed, distance=min_distance, prominence=prominence)

        # Build segment boundaries: [0, peak1, peak2, ..., T-1]
        boundaries = [0] + list(peaks) + [T]
        segments = []

        for i in range(len(boundaries) - 1):
            start = boundaries[i]
            end = boundaries[i + 1]
            if end <= start:
                continue

            # Determine main hand: which hand has higher average speed in this segment
            left_avg = float(np.mean(left_speed[start:end])) if end > start else 0
            right_avg = float(np.mean(right_speed[start:end])) if end > start else 0

            if left_avg > right_avg and left_avg > 0:
                main_hand = 'left'
            elif right_avg > 0:
                main_hand = 'right'
            else:
                main_hand = None

            seg = {
                'start': int(start),
                'end': int(end),
                'main_hand': main_hand,
                'left_avg_speed': left_avg,
                'right_avg_speed': right_avg,
                'duration_s': float(end - start) / fps if fps > 0 else 0,
            }
            segments.append(seg)

        # Fallback: if no segments found, use the whole video as one segment
        if not segments:
            segments = [{
                'start': 0, 'end': T, 'main_hand': None,
                'left_avg_speed': float(np.mean(left_speed)) if T > 0 else 0,
                'right_avg_speed': float(np.mean(right_speed)) if T > 0 else 0,
                'duration_s': float(T) / fps if fps > 0 else 0,
            }]

        return segments

每个片段携带双手子轨迹、主手判断（哪只手速度更高）和持续时间。

5. 按标准格式上传数据集

LeRobot v3.0 使用 Parquet 保存帧级数据和动作片段元数据，并使用 MP4 保存视频。写出算子将每个动作片段构造为 108 维 state 向量 [left_hand(54), right_hand(54)]，其中每只手 = wrist_pos(3) + wrist_rot_6d(6) + mano_joints(45)，action 为下一帧的 state。帧通过 ffmpeg 编码为 H.264 MP4（1280x720@30fps），再按 LeRobotDataset 目录结构生成文件：

class WriteLeRobotDataset:
    def __init__(self, config: dict):
        self.config = config
        self.dataset_dir = None

    def __call__(self, scene: str, clips: list[dict], motion_label: dict,
                 segments: list[dict[str, Any]], instructions: list[str],
                 output_dir: str) -> Path:
        """Write LeRobot v3.0 dataset for a scene. Each segment becomes one episode."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        self.dataset_dir = Path(output_dir) / self.config.get('output', {}).get(
            'dataset_dir', 'vla_dataset')
        T = motion_label['T']
        fps = motion_label.get('fps', 30.0)
        frames_dir = motion_label.get('frames_dir')

        # Build state array for all frames
        all_states = np.zeros((T, STATE_DIM), dtype=np.float32)
        for t in range(T):
            all_states[t] = _build_state_vector(motion_label, t)

        episode_index = 0
        total_frames = 0
        tasks: dict[str, int] = {}

        for seg_idx, seg in enumerate(segments):
            start, end = seg['start'], seg['end']
            task_text = instructions[seg_idx] if seg_idx < len(instructions) else 'null'
            task_index = tasks.setdefault(task_text, len(tasks))

            rows = []
            for t in range(start, end):
                local_t = t - start
                action_t = min(t + 1, T - 1)
                rows.append({
                    'frame_index': local_t,
                    'episode_index': episode_index,
                    'index': total_frames + local_t,
                    'task_index': task_index,
                    'observation.state': all_states[t].tolist(),
                    'observation.action': all_states[action_t].tolist(),
                    'timestamp': float(local_t) / fps if fps > 0 else 0.0,
                })

            # Write parquet
            pq_path = self.dataset_dir / 'data' / 'chunk-000' / f'file-{episode_index:03d}.parquet'
            pq.write_table(pa.Table.from_pylist(rows, schema=schema), pq_path)

            # Encode H.264 MP4 video
            if frames_dir:
                frame_list = _frame_files(frames_dir, start, end)
                mp4_path = (self.dataset_dir / 'videos' / 'observation.images.ego'
                            / 'chunk-000' / f'file-{episode_index:03d}.mp4')
                _encode_video(frame_list, mp4_path, fps,
                    target_w=1280, target_h=720)

            total_frames += (end - start)
            episode_index += 1

最终训练数据目录如下：

dataset/
├── meta/
│   ├── info.json
│   ├── stats.json
│   ├── tasks.parquet
│   └── episodes/chunk-000/file-000.parquet
├── data/chunk-000/file-000.parquet
└── videos/observation.images.ego/chunk-000/file-000.mp4

1. 首次部署：安装依赖

集群冷启动后（autoscaler 销毁过一轮 worker），先通过 ray job submit 把 setup_pip.py 提交到集群，在 driver 和所有 worker 上把 pip 依赖装到系统 Python：

python3 submit_vitra.py  # submit_vitra.py 首次提交会跑 pipeline.py；
                          # 如需单独装依赖，手动提交 setup_pip.py

2. 提交作业

从能访问阿里云 EMR Ray 集群入口的机器上运行 submit_vitra.py。脚本会：

1. 把 vitra/ 目录作为 working_dir 上传到集群（约 12 MB，含第三方源码）。

2. 通过 runtime_env.env_vars 把 Bucket、前缀等注入到所有 worker（不注入 AK/SK）。

3. 以 --no-wait 提交，拿到 job_id 后自动 follow 日志。

# 提交作业（自动 follow 日志，Ctrl+C 不断 job，只断开本地输出）
python3 submit_vitra.py

# 用独立 venv 的 ray CLI（推荐，避免版本不匹配）
RAY_CLI= python3 submit_vitra.py

submit_vitra.py 的关键环境变量：

3. pipeline.py CLI 选项

submit_vitra.py 默认以 pipeline.py --config config.yaml 为入口。pipeline.py 支持的选项：

submit_vitra.py 默认传 --upload-to-oss 和 --delete-after-upload，即跑完即推 OSS 并清理 worker 本地磁盘。

1. 深度估计叠加

MoGe-2 输出的逐帧深度图经 cv2.applyColorMap 转为伪彩色后叠加到原视频上，近距离区域偏暖、远距离区域偏冷，能直观分辨人手、物体和背景的空间层次。

2. 腕部运动轨迹

双手手腕的 3D 位置经相机内参投影到像素坐标后，以折线形式绘制在每帧画面上。绿色折线代表右手轨迹，蓝色折线代表左手轨迹，可清晰追踪手腕的运动路径和当前位置。

3. 相机运动轨迹

MegaSAM 恢复的 camera_to_world [T,4,4] 平移分量以 XY 俯视图绘制为轨迹曲线，起点标为绿色、终点标为红色，用于判断第一视角视频中的整体相机运动幅度和方向。