Add multi-GPU Ego4D local feature extraction without symlinks

README.md

@@ -294,6 +294,12 @@ without a dev set. `run_ego4d_local_baselines.py` preflights all train/test file
 and encoder contracts before training; the summary reports matched-seed deltas,
 including AV minus audio. No new accuracy claim is implied by adding the code.
 
+The global `train/<uid>.pt` and `test/<uid>.pt` layout is now read directly by
+extraction, training and inference; no symlink index is needed.
+`extract_ego4d_local_multigpu.py` distributes full clips across eight independent
+GPU workers, with per-GPU logs, validated resume, and foreground supervision.
+See the guide's L20 commands for the `/data1/yfzx/wangx` paths.
+
 ## Current Ego4D context-length results
 
 The following single-seed snapshot was recorded on 2026-08-22. Both variants

docs/ego4d_local_baselines.md

@@ -50,6 +50,9 @@ LOCAL_ARGS=(
 ```
 
 原始输入应存在：`v1/clips/<uid>.mp4`、`v1/annotations/av_train.json` 和 `av_val.json`。
+全局特征既支持 `<feature-root>/<uid>.pt`，也直接支持原批量提取器的
+`<feature-root>/train/<uid>.pt`、`<feature-root>/test/<uid>.pt`；无需复制或建立软链接。
+提取、训练、推理使用同一个路径解析器；同一 UID 在这几个位置存在多份文件会报错，避免误读旧特征。
 这里 `av_val.json` 是原 AVD 标注的文件名，不意味着把它直接当作本实验的 test 划分；
 train/test 始终按 EgoSpeak manifest。非标准导出路径可加 `--ego4d-annotations <path1> <path2>`。
 局部编码器支持现有 Transformers 4.57.6，也支持 4.52.3；同一次实验所有 clips 应统一版本。
@@ -117,6 +120,70 @@ checkpoint 的 ID 是用户声明，不是权重文件校验和：不要在同
 索引是**先按 split 过滤，再按 manifest 顺序的 1-based 闭区间**；不传 split 时保留整个文件的行顺序。
 如果自行在多卡不同终端运行，必须选择互不重叠的区间和同一套编码器配置；默认不启动后台任务。
 
+### L20 机器：8 卡并行提取，不使用软链接
+
+下面使用已经检查过框的原始视频，以及原 8s 全局特征目录。不要再传旧的占位符 `LOCAL_ARGS`，
+不要带 `--preview-only`、`--preview-with-boxes` 或 `--max-steps`。新入口会自动添加校验后续跑。
+
+```bash
+cd /data1/yfzx/wangx/Code/icassp2027-experiment
+
+MULTIGPU_ARGS=(
+  --gpus 0 1 2 3 4 5 6 7
+  --cpu-threads 4
+  --data-root /data1/yfzx/wangx/Datasets/ego4d
+  --manifest /data1/yfzx/wangx/Datasets/ego4d/ego4d_egospeak.tsv
+  --feature-root /data1/yfzx/wangx/Datasets/ego4d/omni_features/open_stream_2fps_8s_decord
+  --target-root /data1/yfzx/wangx/Datasets/ego4d/ego4d_target_frame
+  --output-root /data1/yfzx/wangx/Datasets/ego4d/omni_features/local_qwen_1s
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B
+  --dtype bfloat16 --flash-attention-2
+  --video-fps 2 --visual-window-seconds 1 --audio-window-seconds 1
+  --decoder-backend decord --decoder-threads 1
+  --log-root logs/ego4d_local_8gpu
+)
+
+# 只检查路径和分片，不加载 Qwen、启动 GPU 任务或写日志。
+python extract_ego4d_local_multigpu.py "${MULTIGPU_ARGS[@]}" --dry-run
+
+# 确认分片和路径后，前台管理 8 个独立 worker。
+python extract_ego4d_local_multigpu.py "${MULTIGPU_ARGS[@]}"
+```
+
+默认 433 clips 分配为 55/54/54/54/54/54/54/54，按 manifest 顺序轮转，每个 UID 只属于一个 worker。
+每个 worker 用 `CUDA_VISIBLE_DEVICES=<指定卡号>` 隔离到一张卡，内部 `--device cuda:0`；
+会覆盖父 shell 遗留的 GPU mask。GPU 编号必须是当前容器暴露的卡，而不是另一个容器的编号。
+这不是 DDP；每张卡独立处理 clips，不需要 GPU 间同步。全量数据默认每个 worker 至多加载一次编码器。
+如使用 `--limit`，少于 GPU 数量时只启动有任务的卡。
+
+主进程打印每卡状态，详细的逐帧/逐 state 进度位于启动时打印的
+`logs/ego4d_local_8gpu/<时间_PID>/gpu0.log` 至 `gpu7.log`。每次运行另建日志目录，不覆盖旧日志；
+同目录的 `launch_plan.json` 保存每卡 UID 和命令。可在另一个终端查看指定卡：
+
+```bash
+tail -f logs/ego4d_local_8gpu/<本次时间_PID>/gpu0.log
+```
+
+不用 `nohup` 或 shell 的 `&`，主命令会一直等待；Ctrl+C 或一个 worker 失败时，会停止本次启动的其他
+worker 及其进程组，不杀无关训练。已完成的 `.pt` 保留，重跑同一命令会校验后跳过；
+失败时先读对应日志。不要同时另开同输出根的重复提取任务。
+
+最终输出是实体文件 `local_qwen_1s/<uid>.pt`，每个文件同时含 A/V，不需要第二个音频任务。
+无框帧仍保留在 5Hz 特征/标签网格中，不能按预览筛选删除。
+8 workers 也会增加 CPU、主机 RAM、磁盘 I/O 和标注 JSON 的解析开销；不要预设能达到 8 倍速度。
+主机内存不足时改为 `--gpus 0 1 2 3` 并重跑，完成的文件仍可复用；不同次分片也不会改写已完成特征。
+
+提取完成后，原训练入口也直接用同一个原始全局特征根，无需 flat 索引：
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 \
+python run_ego4d_local_baselines.py \
+  --data-root /data1/yfzx/wangx/Datasets/ego4d \
+  --feature-root /data1/yfzx/wangx/Datasets/ego4d/omni_features/open_stream_2fps_8s_decord \
+  --local-feature-root /data1/yfzx/wangx/Datasets/ego4d/omni_features/local_qwen_1s \
+  --seeds 0 --dry-run
+```
+
 ## 4. 跑 seed0 四组
 
 ```bash
@@ -180,9 +247,10 @@ python summarize_local_baselines.py \
 ## 回归测试
 
 ```bash
-python -m unittest -q test_easycom_head_objectives test_local_turn_features test_ego4d_local_workflow
+python -m unittest -q test_easycom_head_objectives test_local_turn_features test_ego4d_local_workflow test_ego4d_local_multigpu
 ```
 
 CPU 回归包括 manifest 顺序/防交叉重复、完整批量提取的模型复用、断点文件校验、
 单层/多层全局特征、测试集输入预检，以及合成 Ego4D 上四组各 4 epochs 的训练/推理与增量汇总。
+另测试原 train/test 目录、433 clips 的 8 卡完整互斥分片、GPU mask、独立日志、失败/中断时清理自有进程。
 真实 Ego4D 像素、提供框和 3B/L20 提取仍需按上面的 preview/smoke 步骤验证。

easycom_head.py

@@ -638,6 +638,30 @@ def normalize_sample_name(
         raise ValueError(f"Invalid Ego4D clip UUID: {value}") from error
 
 
+def resolve_feature_path(
+    feature_root: str | Path,
+    sample: str | int,
+    *,
+    dataset: DatasetName = "easycom",
+) -> Path:
+    """Read flat features or Ego4D train/test subdirectories without link farms.
+
+    Manifest membership, not directory discovery, determines the data split.
+    Multiple copies are rejected rather than silently choosing stale features.
+    """
+    name = normalize_sample_name(sample, dataset)
+    root = Path(feature_root)
+    candidates = [root / f"{name}.pt"]
+    if dataset == "ego4d":
+        candidates += [root / split / f"{name}.pt" for split in ("train", "test")]
+    found = [path for path in candidates if path.is_file()]
+    if len(found) > 1:
+        raise ValueError(f"Ambiguous feature copies for {name}: {found}; select one feature root")
+    if not found:
+        raise FileNotFoundError(f"No features for {name}; checked: " + ", ".join(map(str, candidates)))
+    return found[0]
+
+
 def load_ego4d_manifest_split(
     manifest: str | Path,
     split: Literal["train", "test", "all"],
@@ -727,7 +751,7 @@ def load_session_data(
     """Load and validate one saved Omni feature payload and target array."""
 
     name = normalize_sample_name(session, dataset)
-    feature_path = Path(feature_root) / f"{name}.pt"
+    feature_path = resolve_feature_path(feature_root, name, dataset=dataset)
     target_path = Path(target_root) / f"{name}.npy"
     if not feature_path.is_file():
         raise FileNotFoundError(feature_path)
@@ -804,7 +828,7 @@ def load_session_data(
         from local_turn_features import load_local_features
 
         local, contract = load_local_features(
-            Path(local_feature_root) / f"{name}.pt",
+            resolve_feature_path(local_feature_root, name, dataset=dataset),
             payload,
             sample_id=name,
             dataset=dataset,

extract_ego4d_local_multigpu.py

@@ -0,0 +1,226 @@
+"""Foreground supervisor for independent, disjoint Ego4D local-feature workers.
+
+Launcher options are parsed here; all other arguments go to
+extract_turn_local_features.py. No DDP, feature copying or symlinks are needed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+from pathlib import Path
+import shlex
+import signal
+import subprocess
+import sys
+import time
+from datetime import datetime
+
+import extract_turn_local_features as extractor
+from easycom_head import resolve_feature_path
+
+
+def parse_args(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False,
+                                     epilog="Pass --data-root, --feature-root, --target-root, --manifest, "
+                                            "--output-root and encoder options as for the single-GPU extractor.")
+    parser.add_argument("--gpus", nargs="+", type=int, default=list(range(8)),
+                        help="GPU indices exposed by the container; overrides inherited CUDA_VISIBLE_DEVICES")
+    parser.add_argument("--cpu-threads", type=int, default=4, help="OMP/MKL/OpenBLAS threads per worker")
+    parser.add_argument("--log-root", type=Path, default=Path("logs/ego4d_local_multigpu"))
+    parser.add_argument("--status-interval", type=float, default=10., help="Seconds between console status checks")
+    parser.add_argument("--startup-delay", type=float, default=1., help="Seconds between worker launches")
+    parser.add_argument("--dry-run", action="store_true", help="Check input paths and print shard commands; no GPU or writes")
+    args, forwarded = parser.parse_known_args(argv)
+    if forwarded[:1] == ["--"]:
+        forwarded = forwarded[1:]
+    if len(set(args.gpus)) != len(args.gpus) or any(gpu < 0 for gpu in args.gpus):
+        parser.error("GPU indices must be unique and nonnegative")
+    if args.cpu_threads <= 0:
+        parser.error("--cpu-threads must be positive")
+    if not math.isfinite(args.status_interval) or args.status_interval <= 0:
+        parser.error("--status-interval must be finite and positive")
+    if not math.isfinite(args.startup_delay) or args.startup_delay < 0:
+        parser.error("--startup-delay must be finite and nonnegative")
+    forwarded = ["--dataset", "ego4d", *forwarded]
+    extraction = extractor.parse_args(forwarded)
+    if extraction.dataset != "ego4d" or extraction.manifest is None:
+        parser.error("This launcher requires Ego4D --manifest selection")
+    if extraction.preview_only or extraction.max_steps is not None:
+        parser.error("Use the single-GPU extractor for preview/smoke; this launcher writes complete features")
+    if extraction.num_shards != 1 or extraction.shard_index != 0:
+        parser.error("The launcher manages --num-shards/--shard-index; do not set them manually")
+    if extraction.overwrite:
+        parser.error("Multi-GPU extraction uses validated reuse; use a new output root to recompute features")
+    if not extraction.reuse_existing:
+        forwarded.append("--reuse-existing")
+    return args, extraction, forwarded
+
+
+def plan_workers(args, extraction, forwarded):
+    names = extractor.select_samples(extraction)
+    # Fewer clips than GPUs is allowed for a one-clip, full-length pilot.
+    count = min(len(args.gpus), len(names))
+    script = Path(extractor.__file__).resolve()
+    return [{
+        "rank": rank, "gpu": gpu, "samples": names[rank::count],
+        "command": [sys.executable, "-u", str(script), *forwarded,
+                    "--num-shards", str(count), "--shard-index", str(rank),
+                    "--device", "cuda:0", "--decoder-gpu-id", "0"],
+    } for rank, gpu in enumerate(args.gpus[:count])]
+
+
+def check_input_paths(extraction, plans):
+    """Lightweight fail-fast check; workers validate tensors and reuse contracts."""
+    missing = []
+    names = [name for plan in plans for name in plan["samples"]]
+    for name in names:
+        required = [extraction.target_root / f"{name}.npy",
+                    extraction.data_root / "v1/clips" / f"{name}.mp4"]
+        try:
+            required.append(resolve_feature_path(extraction.feature_root, name, dataset="ego4d"))
+        except FileNotFoundError as error:
+            missing.append(str(error))
+        for path in required:
+            if not path.is_file() or path.stat().st_size == 0:
+                missing.append(str(path))
+    annotations = extraction.ego4d_annotations or [extraction.data_root / "v1/annotations" / f"av_{part}.json"
+                                                  for part in ("train", "val")]
+    for path in annotations:
+        if not path.is_file() or path.stat().st_size == 0:
+            missing.append(str(path))
+    model = Path(extraction.model_id).expanduser()
+    if model.is_absolute() and not model.is_dir():
+        missing.append(str(model))
+    if missing:
+        raise FileNotFoundError(f"Missing/empty inputs ({len(missing)}; first 20):\n" + "\n".join(missing[:20]))
+
+
+def worker_environment(gpu, cpu_threads):
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = str(gpu)
+    env["PYTHONUNBUFFERED"] = "1"
+    for name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
+        env[name] = str(cpu_threads)
+    return env
+
+
+def log_tail(path):
+    try:
+        with path.open("rb") as handle:
+            handle.seek(max(0, path.stat().st_size - 4096))
+            lines = handle.read().decode("utf-8", errors="replace").replace("\r", "\n").splitlines()
+        return next((line.strip() for line in reversed(lines) if line.strip()), "starting")[-300:]
+    except FileNotFoundError:
+        return "starting"
+
+
+def signal_worker(process, sig):
+    if process.poll() is not None:
+        return
+    try:
+        if os.name == "posix":
+            os.killpg(process.pid, sig)  # Each owned child starts its own session.
+        elif sig == signal.SIGTERM:
+            process.terminate()
+        else:
+            process.kill()
+    except ProcessLookupError:
+        pass
+
+
+def stop_workers(processes):
+    for process in processes:
+        signal_worker(process, signal.SIGTERM)
+    deadline = time.monotonic() + 10.
+    for process in processes:
+        try:
+            process.wait(timeout=max(0., deadline - time.monotonic()))
+        except subprocess.TimeoutExpired:
+            signal_worker(process, signal.SIGKILL)
+            process.wait(timeout=5.)
+
+
+def run_workers(plans, args, log_dir):
+    running, handles = [], []
+    previous_term = signal.getsignal(signal.SIGTERM)
+
+    def interrupted(signum, frame):
+        raise KeyboardInterrupt
+
+    signal.signal(signal.SIGTERM, interrupted)
+    try:
+        for plan in plans:
+            log = log_dir / f"gpu{plan['gpu']}.log"
+            handle = log.open("xb")
+            handles.append(handle)
+            process = subprocess.Popen(plan["command"],
+                                       env=worker_environment(plan["gpu"], args.cpu_threads),
+                                       stdin=subprocess.DEVNULL, stdout=handle, stderr=subprocess.STDOUT,
+                                       start_new_session=True)
+            running.append((plan, process, log))
+            print(f"GPU {plan['gpu']}: {len(plan['samples'])} clips, pid={process.pid}, log={log}", flush=True)
+            if args.startup_delay:
+                time.sleep(args.startup_delay)
+            for started, child, child_log in running:
+                if child.poll() not in (None, 0):
+                    raise RuntimeError(f"GPU {started['gpu']} failed (exit={child.returncode}); see {child_log}")
+        next_status, last_status = 0., {}
+        while True:
+            completed = 0
+            for plan, process, log in running:
+                code = process.poll()
+                if code not in (None, 0):
+                    raise RuntimeError(f"GPU {plan['gpu']} failed (exit={code}); see {log}\n{log_tail(log)}")
+                completed += code == 0
+            if time.monotonic() >= next_status or completed == len(running):
+                for plan, process, log in running:
+                    state = "done" if process.poll() == 0 else "running"
+                    status = f"GPU {plan['gpu']} {state}: {log_tail(log)}"
+                    if last_status.get(plan["gpu"]) != status:
+                        print(status, flush=True)
+                        last_status[plan["gpu"]] = status
+                next_status = time.monotonic() + args.status_interval
+            if completed == len(running):
+                print(f"All {len(running)} workers completed successfully.", flush=True)
+                return
+            time.sleep(.2)
+    finally:
+        try:
+            stop_workers([process for _, process, _ in running])
+        finally:
+            for handle in handles:
+                handle.close()
+            signal.signal(signal.SIGTERM, previous_term)
+
+
+def main(argv=None):
+    args, extraction, forwarded = parse_args(argv)
+    plans = plan_workers(args, extraction, forwarded)
+    check_input_paths(extraction, plans)
+    print(f"Selected {sum(len(plan['samples']) for plan in plans)} clips across {len(plans)} GPUs; "
+          "each worker loads its encoder once, with fresh media per clip.")
+    print(f"Global input: {extraction.feature_root}; flat or train/test layout is read directly.")
+    print(f"Local output: {extraction.output_root}/<uid>.pt; full A/V, validated reuse, no symlinks.")
+    print("Eight workers also multiply host RAM, annotation parsing and disk I/O; speedup is not guaranteed to be 8x.")
+    for plan in plans:
+        print(f"GPU {plan['gpu']}: rank {plan['rank']}, {len(plan['samples'])} clips")
+        print(f"CUDA_VISIBLE_DEVICES={plan['gpu']} " + shlex.join(plan["command"]))
+    if args.dry_run:
+        print("Dry run: paths/shards checked; tensors, encoder runtime and GPU availability are checked by workers.")
+        return
+    log_dir = args.log_root / (datetime.now().strftime("%Y%m%d_%H%M%S") + f"_{os.getpid()}")
+    log_dir.mkdir(parents=True, exist_ok=False)
+    (log_dir / "launch_plan.json").write_text(json.dumps(plans, indent=2) + "\n")
+    print(f"Logs: {log_dir.resolve()}; foreground supervisor, Ctrl+C stops its workers.", flush=True)
+    try:
+        run_workers(plans, args, log_dir)
+    except KeyboardInterrupt:
+        print("Stopped this launch's workers. Completed .pt files are retained; rerun the same command to resume.", flush=True)
+        raise SystemExit(130)
+
+
+if __name__ == "__main__":
+    main()

extract_turn_local_features.py

@@ -19,7 +19,7 @@ import numpy as np
 import torch
 from tqdm.auto import tqdm
 
-from easycom_head import load_ego4d_manifest_split, normalize_sample_name
+from easycom_head import load_ego4d_manifest_split, normalize_sample_name, resolve_feature_path
 from local_turn_features import (
     LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes,
     load_local_features, pooled_recent_visual, reference_alignment,
@@ -406,6 +406,9 @@ def select_samples(args: argparse.Namespace) -> list[str]:
         raise ValueError("No samples selected")
     if len(set(names)) != len(names):
         raise ValueError("Duplicate samples")
+    num_shards = getattr(args, "num_shards", 1)
+    shard_index = getattr(args, "shard_index", 0)
+    names = names[shard_index::num_shards]
     return names
 
 
@@ -445,7 +448,8 @@ def parse_args(argv=None):
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom")
     p.add_argument("--data-root", required=True, type=Path)
-    p.add_argument("--feature-root", required=True, type=Path)
+    p.add_argument("--feature-root", required=True, type=Path,
+                   help="Flat .pt directory or Ego4D parent containing train/ and test/")
     p.add_argument("--target-root", required=True, type=Path)
     p.add_argument("--output-root", required=True, type=Path)
     selection = p.add_mutually_exclusive_group(required=True)
@@ -455,6 +459,8 @@ def parse_args(argv=None):
     p.add_argument("--start-index", type=int, default=1, help="One-based index AFTER split filtering")
     p.add_argument("--end-index", type=int, help="Inclusive index AFTER split filtering")
     p.add_argument("--limit", type=int, help="First N clips in the selected manifest range")
+    p.add_argument("--num-shards", type=int, default=1, help="Disjoint round-robin shards after all sample selection")
+    p.add_argument("--shard-index", type=int, default=0, help="Zero-based shard index")
     p.add_argument("--ego4d-annotations", type=Path, nargs="+")
     p.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B")
     p.add_argument("--encoder-id", help="Stable checkpoint identity for relocated copies; default: model-id")
@@ -490,6 +496,8 @@ def parse_args(argv=None):
     p.add_argument("--dry-run", action="store_true", help="Validate references/targets and list work; no model, media decode or writes")
     p.add_argument("--no-progress", action="store_true")
     args = p.parse_args(argv)
+    if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards:
+        p.error("Require --num-shards > 0 and 0 <= --shard-index < --num-shards")
     if args.manifest is not None and args.dataset != "ego4d":
         p.error("--manifest is only supported with --dataset ego4d")
     if args.manifest is None and (args.split != "all" or args.start_index != 1
@@ -527,7 +535,10 @@ def parse_args(argv=None):
         p.error("max-pixels must be >=3136; image-batch-size must be positive")
     if args.max_steps is not None and args.max_steps <= 0:
         p.error("--max-steps must be positive")
-    if args.output_root.resolve() == args.feature_root.resolve():
+    protected_roots = [args.feature_root]
+    if args.dataset == "ego4d":
+        protected_roots += [args.feature_root / split for split in ("train", "test")]
+    if args.output_root.resolve() in [root.resolve() for root in protected_roots]:
         p.error("Local outputs must not overwrite original Omni features")
     return args
 
@@ -541,6 +552,11 @@ def main(argv=None) -> None:
         # Fail before indexing/decoding source media, not after expensive I/O.
         version = validate_local_transformers_version(transformers.__version__)
     names = select_samples(args)
+    if not names:
+        print(f"Shard {args.shard_index}/{args.num_shards}: no assigned clips")
+        return
+    print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} assigned clips", flush=True)
+    reference_paths = {name: resolve_feature_path(args.feature_root, name, dataset=args.dataset) for name in names}
     paths = []
     if args.dataset == "ego4d":
         paths = args.ego4d_annotations or [args.data_root / "v1" / "annotations" / f"av_{split}.json"
@@ -549,7 +565,7 @@ def main(argv=None) -> None:
             if not path.is_file():
                 raise FileNotFoundError(path)
     for name in names:
-        required = [args.feature_root / f"{name}.pt", args.target_root / f"{name}.npy"]
+        required = [reference_paths[name], args.target_root / f"{name}.npy"]
         if args.dataset == "ego4d":
             required.append(args.data_root / "v1" / "clips" / f"{name}.mp4")
         for path in required:
@@ -561,7 +577,7 @@ def main(argv=None) -> None:
     pending, contract = [], None
     # Stream through references on CPU; never keep all global states in memory.
     for name in tqdm(names, desc="local preflight", unit="clip", disable=args.no_progress or args.dry_run):
-        reference = torch.load(args.feature_root / f"{name}.pt", map_location="cpu", weights_only=False)
+        reference = torch.load(reference_paths[name], map_location="cpu", weights_only=False)
         times, _ = reference_alignment(reference)
         target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
         if target.shape != (len(times), 3):
@@ -586,7 +602,7 @@ def main(argv=None) -> None:
     encoder = None
     for position, name in enumerate(pending, 1):
         print(f"[{position}/{len(pending)}] {name}", flush=True)
-        path = args.feature_root / f"{name}.pt"
+        path = reference_paths[name]
         reference = torch.load(path, map_location="cpu", weights_only=False)
         times, indices = reference_alignment(reference)
         target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")

run_ego4d_local_baselines.py

@@ -13,7 +13,7 @@ import shlex
 import subprocess
 import sys
 
-from easycom_head import load_ego4d_manifest_split, load_session_data
+from easycom_head import load_ego4d_manifest_split, load_session_data, resolve_feature_path
 
 
 VARIANTS = ("baseline", "audio", "visual", "av")
@@ -24,7 +24,8 @@ def parse_args(argv=None):
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--data-root", type=Path, default=Path("/data/wangx/Datasets/ego4d"))
     p.add_argument("--manifest", type=Path)
-    p.add_argument("--feature-root", required=True, type=Path, help="Your EXISTING global Omni features")
+    p.add_argument("--feature-root", required=True, type=Path,
+                   help="Existing global Omni features: flat or Ego4D train/test parent")
     p.add_argument("--target-root", type=Path)
     p.add_argument("--local-feature-root", type=Path)
     p.add_argument("--checkpoint-root", type=Path, default=Path("checkpoints/ego4d_local_b1"))
@@ -67,11 +68,19 @@ def preflight(args) -> str:
     use_local = any(variant != "baseline" for variant in args.variants)
     wanted = set(args.variants) - {"baseline"}
     local_modalities = next(iter(wanted)) if len(wanted) == 1 else "av"
-    roots = [(args.feature_root, ".pt"), (args.target_root, ".npy")]
+    roots = [args.feature_root]
     if use_local:
-        roots.append((args.local_feature_root, ".pt"))
-    missing = [str(root / f"{name}{suffix}") for root, suffix in roots for name in names
-               if not (root / f"{name}{suffix}").is_file()]
+        roots.append(args.local_feature_root)
+    missing = []
+    for name in names:
+        target = args.target_root / f"{name}.npy"
+        if not target.is_file():
+            missing.append(str(target))
+        for root in roots:
+            try:
+                resolve_feature_path(root, name, dataset="ego4d")
+            except FileNotFoundError as error:
+                missing.append(str(error))
     if missing:
         raise FileNotFoundError(f"Missing {len(missing)} input files (first 20):\n" + "\n".join(missing[:20]))
     for variant in args.variants:

test_ego4d_local_multigpu.py

@@ -0,0 +1,146 @@
+"""No-GPU tests for shard planning, native split paths and worker supervision."""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import os
+from pathlib import Path
+import signal
+import subprocess
+import sys
+import tempfile
+from types import SimpleNamespace
+import unittest
+from unittest.mock import patch
+
+import extract_ego4d_local_multigpu as launch
+
+
+class MultiGPUExtractionTests(unittest.TestCase):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp.cleanup)
+        self.root = Path(self.temp.name)
+        self.names = [f"00000000-0000-0000-0000-{i:012d}" for i in range(1, 434)]
+        self.manifest = self.root / "manifest.tsv"
+        self.manifest.write_text("split\tuid\n" + "".join(
+            f"{'train' if i < 346 else 'test'}\t{name}\n" for i, name in enumerate(self.names)))
+        self.argv = ["--manifest", str(self.manifest), "--data-root", str(self.root),
+                     "--feature-root", str(self.root / "features"), "--target-root", str(self.root / "targets"),
+                     "--output-root", str(self.root / "local"), "--log-root", str(self.root / "logs")]
+
+    def test_433_clips_have_exactly_one_owner_across_eight_gpus(self):
+        args, extraction, forwarded = launch.parse_args(self.argv)
+        plans = launch.plan_workers(args, extraction, forwarded)
+        self.assertEqual(len(plans), 8)
+        self.assertEqual([len(p["samples"]) for p in plans], [55, 54, 54, 54, 54, 54, 54, 54])
+        combined = [name for p in plans for name in p["samples"]]
+        self.assertCountEqual(combined, self.names)
+        self.assertEqual(len(set(combined)), 433)
+        for rank, plan in enumerate(plans):
+            self.assertEqual(plan["samples"], self.names[rank::8])
+            self.assertIn("--reuse-existing", plan["command"])
+            self.assertEqual(plan["command"][-4:], ["--device", "cuda:0", "--decoder-gpu-id", "0"])
+
+    def test_split_range_limit_then_sharding_and_custom_gpu_mask(self):
+        args, extraction, forwarded = launch.parse_args(self.argv + [
+            "--gpus", "2", "4", "--split", "test", "--start-index", "2", "--limit", "3"])
+        plans = launch.plan_workers(args, extraction, forwarded)
+        self.assertEqual([p["gpu"] for p in plans], [2, 4])
+        self.assertEqual(plans[0]["samples"], [self.names[347], self.names[349]])
+        self.assertEqual(plans[1]["samples"], [self.names[348]])
+        with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0", "OMP_NUM_THREADS": "128"}):
+            env = launch.worker_environment(4, 3)
+        self.assertEqual(env["CUDA_VISIBLE_DEVICES"], "4")
+        for key in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
+            self.assertEqual(env[key], "3")
+
+    def test_short_pilot_only_starts_nonempty_shards(self):
+        args, extraction, forwarded = launch.parse_args(self.argv + ["--limit", "1"])
+        plans = launch.plan_workers(args, extraction, forwarded)
+        self.assertEqual(len(plans), 1)
+        self.assertEqual(plans[0]["samples"], [self.names[0]])
+        command = plans[0]["command"]
+        self.assertEqual(command[command.index("--num-shards") + 1], "1")
+
+    def test_rejects_unsafe_or_ambiguous_launcher_options(self):
+        for extra in (["--gpus", "0", "0"], ["--gpus", "-1"], ["--cpu-threads", "0"],
+                      ["--preview-only"], ["--max-steps", "20"], ["--overwrite"],
+                      ["--num-shards", "8"], ["--status-interval", "nan"]):
+            with self.subTest(extra=extra), contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
+                launch.parse_args(self.argv + extra)
+
+    def test_dry_run_checks_native_split_without_launching_or_writing(self):
+        uid = self.names[0]
+        for relative in (f"features/train/{uid}.pt", f"targets/{uid}.npy", f"v1/clips/{uid}.mp4",
+                         "v1/annotations/av_train.json", "v1/annotations/av_val.json"):
+            path = self.root / relative
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_bytes(b"path-check-only")
+        with patch("extract_ego4d_local_multigpu.subprocess.Popen") as popen, \
+             contextlib.redirect_stdout(io.StringIO()):
+            launch.main(self.argv + ["--limit", "1", "--dry-run"])
+            popen.assert_not_called()
+        self.assertFalse((self.root / "local").exists())
+        self.assertFalse((self.root / "logs").exists())
+        (self.root / "features/train" / f"{uid}.pt").unlink()
+        with patch("extract_ego4d_local_multigpu.subprocess.Popen") as popen, self.assertRaises(FileNotFoundError):
+            launch.main(self.argv + ["--limit", "1", "--dry-run"])
+            popen.assert_not_called()
+
+    def test_supervisor_uses_one_gpu_per_child_and_writes_separate_logs(self):
+        args = SimpleNamespace(cpu_threads=2, startup_delay=0., status_interval=.01)
+        code = "import os; print(os.environ['CUDA_VISIBLE_DEVICES'], os.environ['OMP_NUM_THREADS'], flush=True)"
+        plans = [{"gpu": gpu, "samples": ["dummy"], "command": [sys.executable, "-c", code]} for gpu in (2, 5)]
+        with contextlib.redirect_stdout(io.StringIO()):
+            launch.run_workers(plans, args, self.root)
+        self.assertEqual((self.root / "gpu2.log").read_text().strip(), "2 2")
+        self.assertEqual((self.root / "gpu5.log").read_text().strip(), "5 2")
+
+    def test_failing_worker_stops_only_its_launched_peers(self):
+        args = SimpleNamespace(cpu_threads=1, startup_delay=0., status_interval=.01)
+        plans = [{"gpu": 0, "samples": ["dummy"], "command": [sys.executable, "-c", "import time; time.sleep(60)"]},
+                 {"gpu": 1, "samples": ["dummy"], "command": [sys.executable, "-c", "raise SystemExit(7)"]}]
+        processes = []
+        real_popen = subprocess.Popen
+
+        def record(*a, **kw):
+            process = real_popen(*a, **kw)
+            processes.append(process)
+            return process
+
+        try:
+            with patch("extract_ego4d_local_multigpu.subprocess.Popen", side_effect=record), \
+                 contextlib.redirect_stdout(io.StringIO()), self.assertRaisesRegex(RuntimeError, "exit=7"):
+                launch.run_workers(plans, args, self.root)
+            self.assertEqual(len(processes), 2)
+            self.assertTrue(all(p.poll() is not None for p in processes))
+        finally:
+            launch.stop_workers(processes)
+
+    def test_interrupt_stops_owned_worker_and_restores_signal_handler(self):
+        args = SimpleNamespace(cpu_threads=1, startup_delay=0., status_interval=.01)
+        plans = [{"gpu": 0, "samples": ["dummy"], "command": [sys.executable, "-c", "import time; time.sleep(60)"]}]
+        original_handler = signal.getsignal(signal.SIGTERM)
+        processes = []
+        real_popen = subprocess.Popen
+
+        def record(*a, **kw):
+            process = real_popen(*a, **kw)
+            processes.append(process)
+            return process
+
+        try:
+            with patch("extract_ego4d_local_multigpu.subprocess.Popen", side_effect=record), \
+                 patch("extract_ego4d_local_multigpu.log_tail", side_effect=KeyboardInterrupt), \
+                 contextlib.redirect_stdout(io.StringIO()), self.assertRaises(KeyboardInterrupt):
+                launch.run_workers(plans, args, self.root)
+            self.assertTrue(all(p.poll() is not None for p in processes))
+            self.assertEqual(signal.getsignal(signal.SIGTERM), original_handler)
+        finally:
+            launch.stop_workers(processes)
+
+
+if __name__ == "__main__":
+    unittest.main()

test_ego4d_local_workflow.py

@@ -18,7 +18,7 @@ from unittest.mock import patch
 import numpy as np
 import torch
 
-from easycom_head import load_ego4d_manifest_split
+from easycom_head import load_ego4d_manifest_split, load_session_data, resolve_feature_path
 import extract_turn_local_features as extract
 import run_ego4d_local_baselines as runner
 from test_local_turn_features import fixture
@@ -73,6 +73,45 @@ class Ego4DLocalWorkflowTests(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "Duplicate"):
             load_ego4d_manifest_split(self.manifest, "train")
 
+    def split_feature_layout(self, folder="features"):
+        for name, split in zip(self.names, self.splits):
+            destination = self.root / folder / split
+            destination.mkdir(exist_ok=True)
+            (self.root / folder / f"{name}.pt").rename(destination / f"{name}.pt")
+
+    def test_native_split_paths_and_ambiguous_copies(self):
+        self.split_feature_layout()
+        self.split_feature_layout("local")
+        for name, split in zip(self.names, self.splits):
+            path = resolve_feature_path(self.root / "features", name, dataset="ego4d")
+            self.assertEqual(path, self.root / "features" / split / f"{name}.pt")
+            self.assertFalse(path.is_symlink())
+            data = load_session_data(self.root / "features", self.root / "targets", name,
+                                     dataset="ego4d", expected_mode="open-stream-replay",
+                                     local_feature_root=self.root / "local")
+            self.assertEqual(len(data.states), 20)
+        duplicate = self.root / "features" / f"{self.names[0]}.pt"
+        duplicate.write_bytes(b"must-not-be-silently-chosen")
+        with self.assertRaisesRegex(ValueError, "Ambiguous feature copies"):
+            resolve_feature_path(self.root / "features", self.names[0], dataset="ego4d")
+
+    def test_shards_are_disjoint_and_cover_filtered_selection(self):
+        selected = []
+        for rank in range(8):
+            args = extract.parse_args(self.extraction_args(["--num-shards", "8", "--shard-index", str(rank)]))
+            shard = extract.select_samples(args)
+            self.assertEqual(shard, self.names[rank::8])
+            selected.extend(shard)
+        self.assertCountEqual(selected, self.names)
+        for rank in range(2):
+            args = extract.parse_args(self.extraction_args(["--split", "train", "--num-shards", "2",
+                                                           "--shard-index", str(rank)]))
+            self.assertEqual(extract.select_samples(args), [self.names[rank * 2]])
+
+    def test_split_input_directory_is_protected_from_output_overwrite(self):
+        with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
+            extract.parse_args(self.extraction_args(["--output-root", str(self.root / "features/train")]))
+
     def test_invalid_cli_combinations(self):
         for extra in (["--samples", self.names[0]], ["--dataset", "easycom"], ["--start-index", "0"],
                       ["--limit", "0"], ["--max-steps", "20", "--reuse-existing"],
@@ -94,6 +133,7 @@ class Ego4DLocalWorkflowTests(unittest.TestCase):
             extract.main(self.extraction_args(["--dry-run"]))
 
     def test_one_encoder_per_batch_and_validated_reuse(self):
+        self.split_feature_layout()
         fake_encoder = SimpleNamespace(transformers_version="4.57.6", visual_dim=6, audio_dim=5,
                                        model_config_sha256="same-checkpoint")
         _, side, _ = fixture(dataset="ego4d")
@@ -174,6 +214,7 @@ class Ego4DLocalWorkflowTests(unittest.TestCase):
         self.assertFalse((self.root / "checkpoints").exists())
 
     def test_full_cpu_family_round_trip_and_summary(self):
+        self.split_feature_layout()
         env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
         completed = subprocess.run([sys.executable, str(runner.REPO / "run_ego4d_local_baselines.py"),
                                     *self.runner_args()], env=env, capture_output=True, text=True, timeout=90)