diff --git a/README.md b/README.md index d820fe3..6cc105b 100644 --- a/README.md +++ b/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/.pt` and `test/.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 diff --git a/docs/ego4d_local_baselines.md b/docs/ego4d_local_baselines.md index 74a31cc..032a563 100644 --- a/docs/ego4d_local_baselines.md +++ b/docs/ego4d_local_baselines.md @@ -50,6 +50,9 @@ LOCAL_ARGS=( ``` 原始输入应存在：`v1/clips/.mp4`、`v1/annotations/av_train.json` 和 `av_val.json`。 +全局特征既支持 `/.pt`，也直接支持原批量提取器的 +`/train/.pt`、`/test/.pt`；无需复制或建立软链接。 +提取、训练、推理使用同一个路径解析器；同一 UID 在这几个位置存在多份文件会报错，避免误读旧特征。 这里 `av_val.json` 是原 AVD 标注的文件名，不意味着把它直接当作本实验的 test 划分； train/test 始终按 EgoSpeak manifest。非标准导出路径可加 `--ego4d-annotations `。 局部编码器支持现有 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/.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 步骤验证。 diff --git a/easycom_head.py b/easycom_head.py index 385fd65..7dd97d4 100644 --- a/easycom_head.py +++ b/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, diff --git a/extract_ego4d_local_multigpu.py b/extract_ego4d_local_multigpu.py new file mode 100644 index 0000000..f82bfac --- /dev/null +++ b/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}/.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() diff --git a/extract_turn_local_features.py b/extract_turn_local_features.py index 3be96c4..e0ee3ff 100644 --- a/extract_turn_local_features.py +++ b/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") diff --git a/run_ego4d_local_baselines.py b/run_ego4d_local_baselines.py index 97b98cb..c4d619e 100644 --- a/run_ego4d_local_baselines.py +++ b/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: diff --git a/test_ego4d_local_multigpu.py b/test_ego4d_local_multigpu.py new file mode 100644 index 0000000..6c2daab --- /dev/null +++ b/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() diff --git a/test_ego4d_local_workflow.py b/test_ego4d_local_workflow.py index 2957b67..a44b621 100644 --- a/test_ego4d_local_workflow.py +++ b/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)