"""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, load_session_data, resolve_feature_path 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 split_feature_layout(self, folder="features"): for name, split in zip(self.names, self.splits): destination = self.root / folder / split destination.mkdir(exist_ok=True) (self.root / folder / f"{name}.pt").rename(destination / f"{name}.pt") def test_native_split_paths_and_ambiguous_copies(self): self.split_feature_layout() self.split_feature_layout("local") for name, split in zip(self.names, self.splits): path = resolve_feature_path(self.root / "features", name, dataset="ego4d") self.assertEqual(path, self.root / "features" / split / f"{name}.pt") self.assertFalse(path.is_symlink()) data = load_session_data(self.root / "features", self.root / "targets", name, dataset="ego4d", expected_mode="open-stream-replay", local_feature_root=self.root / "local") self.assertEqual(len(data.states), 20) duplicate = self.root / "features" / f"{self.names[0]}.pt" duplicate.write_bytes(b"must-not-be-silently-chosen") with self.assertRaisesRegex(ValueError, "Ambiguous feature copies"): resolve_feature_path(self.root / "features", self.names[0], dataset="ego4d") def test_shards_are_disjoint_and_cover_filtered_selection(self): selected = [] for rank in range(8): args = extract.parse_args(self.extraction_args(["--num-shards", "8", "--shard-index", str(rank)])) shard = extract.select_samples(args) self.assertEqual(shard, self.names[rank::8]) selected.extend(shard) self.assertCountEqual(selected, self.names) for rank in range(2): args = extract.parse_args(self.extraction_args(["--split", "train", "--num-shards", "2", "--shard-index", str(rank)])) self.assertEqual(extract.select_samples(args), [self.names[rank * 2]]) def test_split_input_directory_is_protected_from_output_overwrite(self): with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit): extract.parse_args(self.extraction_args(["--output-root", str(self.root / "features/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): self.split_feature_layout() 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: {}) def fake_build(args, name, end, annotations): annotations[name]["frame_audit"] = {"ignored_frame_indices": []} return media, [] 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", side_effect=fake_build) 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_ego4d_noncanonical_binary_rows_are_preserved(self): target_path = self.root / "targets" / f"{self.names[0]}.npy" targets = np.load(target_path) targets[3] = 0 targets[4] = [1, 1, 0] np.save(target_path, targets) data = load_session_data( self.root / "features", self.root / "targets", self.names[0], dataset="ego4d", expected_mode="open-stream-replay", ) self.assertTrue(torch.equal(data.targets[3], torch.zeros(3))) self.assertTrue( torch.equal(data.targets[4], torch.tensor([1.0, 1.0, 0.0])) ) self.assertEqual(data.feature_metadata["num_noncanonical_target_rows"], 2) output = io.StringIO() with contextlib.redirect_stdout(output): runner.preflight(runner.parse_args(self.runner_args(["--variants", "baseline"]))) self.assertIn("constraint: 2/60 ([0,0,0]=1, [1,1,0]=1)", output.getvalue()) 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): # EgoSpeak's LSTR MCE consumes the released Ego4D rows as binary # weights. Preserve noncanonical rows, and do not derive transitions # when that auxiliary objective is disabled. target_path = self.root / "targets" / f"{self.names[0]}.npy" targets = np.load(target_path) targets[3] = 0 np.save(target_path, targets) self.split_feature_layout() 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()