icassp2027-experiment/test_ego4d_local_workflow.py

"""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()