"""CPU checks for manifest extraction, safe reuse and the four-run Ego4D family.""" from __future__ import annotations import contextlib import copy import io import json import os from pathlib import Path import subprocess import sys import tempfile from types import SimpleNamespace import unittest from unittest.mock import patch import numpy as np import torch from easycom_head import load_ego4d_manifest_split import extract_turn_local_features as extract import run_ego4d_local_baselines as runner from test_local_turn_features import fixture class Ego4DLocalWorkflowTests(unittest.TestCase): def setUp(self): torch.set_num_threads(1) 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 (4, 1, 5)] self.splits = ["train", "test", "train"] for folder in ("features", "targets", "local", "v1/clips", "v1/annotations"): (self.root / folder).mkdir(parents=True) self.manifest = self.root / "ego4d_egospeak.tsv" self.manifest.write_text("split\tuid\n" + "".join( f"{split}\t{name}\n" for split, name in zip(self.splits, self.names))) for name in self.names: reference, side, labels = fixture(name=name, dataset="ego4d") torch.save(reference, self.root / "features" / f"{name}.pt") torch.save(side, self.root / "local" / f"{name}.pt") np.save(self.root / "targets" / f"{name}.npy", labels) (self.root / "v1/clips" / f"{name}.mp4").write_bytes(b"fake-media-not-decoded") for annotation_split, split in (("train", "train"), ("val", "test")): payload = {"videos": [{"clips": [{"clip_uid": name, "persons": []} for name, part in zip(self.names, self.splits) if part == split]}]} (self.root / "v1/annotations" / f"av_{annotation_split}.json").write_text(json.dumps(payload)) def extraction_args(self, extra=()): return ["--dataset", "ego4d", "--data-root", str(self.root), "--feature-root", str(self.root / "features"), "--target-root", str(self.root / "targets"), "--manifest", str(self.manifest), "--output-root", str(self.root / "extracted"), *extra] def runner_args(self, extra=()): return ["--data-root", str(self.root), "--feature-root", str(self.root / "features"), "--target-root", str(self.root / "targets"), "--local-feature-root", str(self.root / "local"), "--checkpoint-root", str(self.root / "checkpoints"), "--result-root", str(self.root / "results"), "--seeds", "0", "--device", "cpu", "--no-progress", *extra] def test_manifest_order_split_and_selection(self): self.assertEqual(load_ego4d_manifest_split(self.manifest, "all"), self.names) args = extract.parse_args(self.extraction_args(["--split", "train", "--start-index", "2", "--limit", "1"])) self.assertEqual(extract.select_samples(args), [self.names[2]]) args = extract.parse_args(self.extraction_args(["--start-index", "2", "--end-index", "2"])) self.assertEqual(extract.select_samples(args), [self.names[1]]) args = extract.parse_args(self.extraction_args(["--start-index", "100"])) with self.assertRaisesRegex(ValueError, "No samples"): extract.select_samples(args) with self.manifest.open("a") as handle: handle.write(f"test\t{self.names[0]}\n") with self.assertRaisesRegex(ValueError, "Duplicate"): load_ego4d_manifest_split(self.manifest, "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"], ["--overwrite", "--reuse-existing"]): with self.subTest(extra=extra), contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit): extract.parse_args(self.extraction_args(extra)) def test_dry_run_no_model_decode_tracks_or_writes(self): with patch("extract_turn_local_features.FrozenQwenLocalEncoder") as encoder, \ patch("extract_turn_local_features.build_media_and_boxes") as media, \ patch("extract_turn_local_features.load_ego4d_boxes") as boxes: extract.main(self.extraction_args(["--dry-run"])) encoder.assert_not_called() media.assert_not_called() boxes.assert_not_called() self.assertFalse((self.root / "extracted").exists()) np.save(self.root / "targets" / f"{self.names[1]}.npy", np.zeros((19, 3))) with self.assertRaisesRegex(ValueError, "Reference/target mismatch"): extract.main(self.extraction_args(["--dry-run"])) def test_one_encoder_per_batch_and_validated_reuse(self): fake_encoder = SimpleNamespace(transformers_version="4.57.6", visual_dim=6, audio_dim=5, model_config_sha256="same-checkpoint") _, side, _ = fixture(dataset="ego4d") media = SimpleNamespace(clips=[], metadata=lambda: {}) with patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="4.57.6")}), \ patch("extract_turn_local_features.FrozenQwenLocalEncoder", return_value=fake_encoder) as encoder, \ patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])) as build, \ patch("extract_turn_local_features.extract_local", side_effect=lambda *a: copy.deepcopy(side)): extract.main(self.extraction_args()) encoder.assert_called_once() self.assertEqual(build.call_count, 3) encoder.reset_mock() build.reset_mock() extract.main(self.extraction_args(["--reuse-existing"])) encoder.assert_not_called() build.assert_not_called() with self.assertRaisesRegex(ValueError, "audio_window_seconds"): extract.main(self.extraction_args(["--reuse-existing", "--audio-window-seconds", "2"])) first = self.root / "extracted" / f"{self.names[0]}.pt" payload = torch.load(first, weights_only=False) payload["partial"] = True torch.save(payload, first) with self.assertRaisesRegex(ValueError, "Partial smoke-test"): extract.main(self.extraction_args(["--reuse-existing"])) def test_four_commands_keep_splits_epochs_and_feature_ablation(self): args = runner.parse_args(self.runner_args()) fusion = runner.preflight(args) self.assertEqual(fusion, "concat-mlp") commands = list(runner.experiment_commands(args, fusion)) self.assertEqual([name for name, _, _ in commands], [v + "_seed0" for v in runner.VARIANTS]) for name, train, infer in commands: self.assertEqual(train[train.index("--epochs") + 1], "4") for command in (train, infer): self.assertEqual(command[command.index("--dataset") + 1], "ego4d") self.assertEqual(command[command.index("--manifest") + 1], str(self.manifest.resolve())) self.assertNotIn("--sessions", command) self.assertNotIn("--dev-sessions", command) self.assertEqual("--local-feature-root" in command, name != "baseline_seed0") if name != "baseline_seed0": self.assertEqual(train[train.index("--local-modalities") + 1], name.split("_")[0]) def test_family_preflight_catches_missing_test_sidecar_before_training(self): (self.root / "local" / f"{self.names[1]}.pt").unlink() with patch("run_ego4d_local_baselines.subprocess.run") as run: with self.assertRaisesRegex(FileNotFoundError, self.names[1]): runner.main(self.runner_args()) run.assert_not_called() self.assertEqual(runner.preflight(runner.parse_args(self.runner_args(["--variants", "baseline"]))), "concat-mlp") def test_preflight_rejects_mixed_contracts_and_existing_outputs(self): path = self.root / "local" / f"{self.names[1]}.pt" side = torch.load(path, weights_only=False) side["config"]["encoder_id"] = "different-test-encoder" torch.save(side, path) with self.assertRaisesRegex(ValueError, "contract differs"): runner.preflight(runner.parse_args(self.runner_args())) (self.root / "results" / "baseline_seed0").mkdir(parents=True) with self.assertRaisesRegex(FileExistsError, "Refusing"): runner.preflight(runner.parse_args(self.runner_args(["--variants", "baseline"]))) def test_single_layer_auto_and_dry_run(self): for name in self.names: path = self.root / "features" / f"{name}.pt" reference = torch.load(path, weights_only=False) reference["states"] = reference["states"].mean(1) torch.save(reference, path) args = runner.parse_args(self.runner_args()) self.assertEqual(runner.preflight(args), "single") for _, train, _ in runner.experiment_commands(args, "single"): self.assertNotIn("--layer-fusion", train) with self.assertRaisesRegex(ValueError, "requires stacked"): runner.preflight(runner.parse_args(self.runner_args(["--layer-fusion", "concat-mlp"]))) with patch("run_ego4d_local_baselines.subprocess.run") as run: runner.main(self.runner_args(["--dry-run"])) run.assert_not_called() self.assertFalse((self.root / "checkpoints").exists()) def test_full_cpu_family_round_trip_and_summary(self): 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) self.assertEqual(completed.returncode, 0, completed.stdout + completed.stderr) self.assertIn("av - audio (seeds=[0])", completed.stdout) scores = {} for variant in runner.VARIANTS: path = self.root / "results" / f"{variant}_seed0" / "metrics.json" metrics = json.loads(path.read_text()) self.assertEqual(metrics["checkpoint_epoch"], 4) self.assertEqual(metrics["dataset"], "ego4d") self.assertEqual(metrics["checkpoint_dev_sessions"], []) self.assertEqual(metrics["sessions"], [self.names[1]]) self.assertEqual(metrics["checkpoint_train_sessions"], [self.names[0], self.names[2]]) self.assertEqual(bool(metrics["local_contract"]), variant != "baseline") scores[variant] = metrics["causal_valid_only"]["horizons"][0]["map"] expected_delta = 100 * (scores["av"] - scores["audio"]) self.assertIn(f"av - audio (seeds=[0]): {expected_delta:+.5f} ± n/a", completed.stdout) # Do not silently pool a mismatched epoch as a paired comparison. metrics["checkpoint_epoch"] = 25 path.write_text(json.dumps(metrics)) completed = subprocess.run([sys.executable, str(runner.REPO / "summarize_local_baselines.py"), "--root", str(self.root / "results"), "--dataset", "ego4d", "--paired-deltas"], env=env, capture_output=True, text=True, timeout=20) self.assertNotEqual(completed.returncode, 0) self.assertIn("different splits/epochs/loss/fusion", completed.stderr) if __name__ == "__main__": unittest.main()