"""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()