147 lines
7.5 KiB
Python
147 lines
7.5 KiB
Python
"""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()
|