"""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_failure_tail_keeps_traceback_and_long_diagnostics_but_progress_is_short(self): path = self.root / "worker.log" message = ("ValueError: Ego4D annotation frame exceeds source clip: " + self.names[0] + "; source_frames=8998; " + "diagnostic " * 100 + "first_10=[8998, 8999]") traceback = "Traceback (most recent call last):\n File 'extract.py', line 200\n" + message path.write_text("old progress\n" * 10000 + traceback + "\n") self.assertEqual(launch.log_tail(path, failure=True), traceback) self.assertLessEqual(len(launch.log_tail(path)), 300) path.write_text(message + "\n") self.assertEqual(launch.log_tail(path, failure=True), message) self.assertEqual(launch.log_tail(self.root / "missing.log", failure=True), "starting") def test_failure_diagnostics_survive_both_startup_and_polling_errors(self): message = f"clip={self.names[0]}; source_frames=8998; " + "detail " * 100 + "outliers=[8998, 8999]" for startup in (True, False): with self.subTest(startup=startup): folder = self.root / str(startup) folder.mkdir() # Simulate a flushed failing worker at startup or in the poll # loop without scheduler-dependent sleep/timing assertions. process = SimpleNamespace(returncode=1) process.poll = lambda: 1 process.wait = lambda **kwargs: 1 process.pid = None # Stub only; must never be passed to a real signal call. def start(command, **kwargs): kwargs["stdout"].write(("Traceback (most recent call last):\nValueError: " + message + "\n").encode()) kwargs["stdout"].flush() if not startup: polls = iter([None, 1, 1]) process.poll = lambda: next(polls, 1) return process args = SimpleNamespace(cpu_threads=1, startup_delay=0., status_interval=.01) plans = [{"gpu": 7, "samples": [self.names[0]], "command": ["unused"]}] with patch("extract_ego4d_local_multigpu.subprocess.Popen", side_effect=start), \ contextlib.redirect_stdout(io.StringIO()), self.assertRaises(RuntimeError) as raised: launch.run_workers(plans, args, folder) self.assertIn(message, str(raised.exception)) self.assertIn("Traceback (most recent call last):", str(raised.exception)) 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()