207 lines
12 KiB
Python
207 lines
12 KiB
Python
|
"""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()
|