227 lines
10 KiB
Python
227 lines
10 KiB
Python
|
"""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()
|