"""EasyCom encoder-only split, packing and real CPU train/infer checks.""" from __future__ import annotations import contextlib import copy import io import json import os from pathlib import Path import subprocess import sys import tempfile import unittest from unittest.mock import patch import numpy as np import torch from easycom_head import load_session_data from local_turn_features import MODALITIES import prepare_encoder_only_features as prepare import run_easycom_encoder_only as runner from test_local_turn_features import fixture class EasyComEncoderOnlyTests(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"Session_{i}" for i in range(1, 13)] self.train = self.names[3:] self.test = ["Session_3", "Session_2", "Session_1"] self.source = self.root / "local" self.targets = self.root / "easycom_target_frame" self.source.mkdir() self.targets.mkdir() self.payloads = {} for name in self.names: # No Thinker features, manifest, media, annotations or model exists. _, side, target = fixture(name=name, dataset="easycom") side["reference_features"] = "/must-not-be-read/thinker.pt" side["states"] = "POISON: not an encoder input" torch.save(side, self.source / f"{name}.pt") np.save(self.targets / f"{name}.npy", target) self.payloads[name] = side def output_root(self, modality="av"): return self.root / f"encoder_{modality}" def prepare_args(self, modality="av", extra=()): return ["--dataset", "easycom", "--data-root", str(self.root), "--local-feature-root", str(self.source), "--modalities", modality, "--output-root", str(self.output_root(modality)), "--no-progress", *extra] def runner_args(self, modality="av", extra=()): return ["--data-root", str(self.root), "--feature-root", str(self.output_root(modality)), "--modalities", modality, "--checkpoint-root", str(self.root / "checkpoints"), "--result-root", str(self.root / "results"), "--seeds", "0", "--device", "cpu", "--no-progress", *extra] def export(self, modality="av"): with contextlib.redirect_stdout(io.StringIO()): prepare.main(self.prepare_args(modality)) def test_dataset_defaults_and_ascending_b1_training_order(self): args = prepare.parse_args(["--dataset", "easycom"]) root = Path("/data1/yfzx/wangx/Datasets/easycom") self.assertEqual(args.data_root, root) self.assertEqual(args.target_root, root / "easycom_target_frame") self.assertIsNone(args.manifest) run = runner.parse_args(["--feature-root", str(self.output_root())]) self.assertEqual(run.dataset, "easycom") self.assertEqual(run.data_root, root) self.assertEqual(run.target_root, root / "easycom_target_frame") self.assertIsNone(run.manifest) self.assertEqual(run.result_root.name, "easycom_encoder_only_1s") self.assertEqual(run.checkpoint_root.name, "easycom_encoder_only_1s") self.assertEqual(prepare.encoder_sample_split("easycom"), (self.names, self.train, self.test)) # Preserve the original export command's Ego4D default. self.assertEqual(prepare.parse_args([]).dataset, "ego4d") def test_all_modalities_export_twelve_real_files_and_exact_native_inputs(self): for modality, parts in MODALITIES.items(): with self.subTest(modality=modality): self.export(modality) output = self.output_root(modality) self.assertEqual((output / "complete.txt").read_text(), "12 sessions\n") manifest = json.loads((output / "export_manifest.json").read_text()) self.assertEqual(manifest["dataset"], "easycom") self.assertIsNone(manifest["manifest"]) self.assertEqual(manifest["train_samples"], self.train) self.assertEqual(manifest["test_samples"], self.test) self.assertEqual(manifest["sessions"], dict.fromkeys(self.names, 20)) self.assertEqual(manifest["num_states"], 240) for name in self.names: path = output / f"{name}.pt" self.assertTrue(path.is_file()) self.assertFalse(path.is_symlink()) packed = torch.load(path, weights_only=False) original = self.payloads[name] expected = torch.cat([tensor for part in parts for tensor in ( original[part].float(), original[f"{part}_valid"][:, None].float(), )], dim=-1).to(torch.bfloat16) self.assertTrue(torch.equal(packed["states"], expected)) self.assertTrue(torch.equal(packed["times"], original["times"])) self.assertTrue(torch.equal(packed["label_indices"], original["label_indices"])) self.assertEqual(packed["dataset"], "easycom") self.assertIs(packed["config"]["uses_thinker_states"], False) data = load_session_data(output, self.targets, name, expected_mode=f"encoder-only-{modality}") self.assertIsNone(data.local_features) self.assertEqual(tuple(data.states.shape), tuple(expected.shape)) with self.assertRaises(FileExistsError): prepare.main(self.prepare_args(modality)) def test_partial_wrong_dataset_and_missing_heldout_session_fail_before_write(self): name = "Session_3" path = self.source / f"{name}.pt" for key, value in (("partial", True), ("dataset", "ego4d")): side = copy.deepcopy(self.payloads[name]) side[key] = value torch.save(side, path) with self.subTest(key=key), self.assertRaises(ValueError): prepare.main(self.prepare_args()) self.assertFalse(self.output_root().exists()) path.unlink() with self.assertRaises(FileNotFoundError): prepare.main(self.prepare_args()) self.assertFalse(self.output_root().exists()) def test_easycom_rejects_ego4d_manifest(self): with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit): prepare.parse_args(self.prepare_args(extra=["--manifest", "/ego4d.tsv"])) with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit): runner.parse_args(self.runner_args(extra=["--manifest", "/ego4d.tsv"])) def test_dry_runs_write_nothing_and_use_fixed_sessions(self): with contextlib.redirect_stdout(io.StringIO()): prepare.main(self.prepare_args(extra=["--dry-run"])) self.assertFalse(self.output_root().exists()) self.export() args = runner.parse_args(self.runner_args(extra=["--seeds", "0", "1", "2"])) commands = list(runner.experiment_commands(args)) self.assertEqual(len(commands), 3) for seed, (train, infer) in enumerate(commands): self.assertEqual(train[train.index("--epochs") + 1], "4") self.assertEqual(train[train.index("--warmup-epochs") + 1], "1.6") self.assertEqual(train[train.index("--seed") + 1], str(seed)) self.assertEqual(train[train.index("--sessions") + 1:], self.train) self.assertEqual(infer[infer.index("--sessions") + 1:], self.test) for command in (train, infer): self.assertEqual(command[command.index("--dataset") + 1], "easycom") self.assertEqual(command[command.index("--expected-mode") + 1], "encoder-only-av") for forbidden in ("--manifest", "--local-feature-root", "--layer-fusion", "--dev-sessions"): self.assertNotIn(forbidden, command) with patch.object(runner.workflow.subprocess, "run") as run, contextlib.redirect_stdout(io.StringIO()): runner.main(self.runner_args(extra=["--dry-run"])) run.assert_not_called() self.assertFalse((self.root / "checkpoints").exists()) def test_runner_rejects_wrong_marker_missing_test_input_and_existing_results(self): self.export() args = runner.parse_args(self.runner_args()) marker = self.output_root() / "complete.txt" marker.write_text("12 clips\n") with self.assertRaisesRegex(ValueError, "incomplete"): runner.preflight(args) marker.write_text("12 sessions\n") _, output = runner.run_paths(args, 0) output.mkdir(parents=True) with self.assertRaises(FileExistsError): runner.preflight(args) output.rmdir() (self.output_root() / "Session_1.pt").unlink() with self.assertRaises(FileNotFoundError): runner.preflight(args) def test_cpu_round_trip_three_modalities_four_epochs_no_dev(self): env = os.environ.copy() env.update(OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1") for modality, parts in MODALITIES.items(): with self.subTest(modality=modality): self.export(modality) result = subprocess.run( [sys.executable, str(Path(runner.__file__).resolve()), *self.runner_args(modality)], env=env, capture_output=True, text=True, timeout=120, ) self.assertEqual(result.returncode, 0, result.stdout + result.stderr) args = runner.parse_args(self.runner_args(modality)) checkpoint, output = runner.run_paths(args, 0) saved = torch.load(checkpoint, map_location="cpu", weights_only=False) self.assertEqual(saved["dataset"], "easycom") self.assertEqual(saved["epoch"], 4) self.assertEqual(saved["train_sessions"], self.train) self.assertEqual(saved["dev_sessions"], []) self.assertIsNone(saved["local_contract"]) self.assertIsNone(saved["training_config"]["local_feature_root"]) width = sum(self.payloads["Session_1"]["config"][f"{part}_dim"] + 1 for part in parts) self.assertEqual(saved["head_config"]["hidden_size"], width) self.assertEqual(saved["head_config"]["local_visual_dim"], 0) self.assertEqual(saved["head_config"]["local_audio_dim"], 0) self.assertEqual(saved["head_config"]["layer_fusion"], "single") self.assertFalse(any("local_projections" in key for key in saved["head"])) metrics = json.loads((output / "metrics.json").read_text()) self.assertEqual(metrics["sessions"], self.test) self.assertEqual(metrics["checkpoint_epoch"], 4) for name in self.test: self.assertEqual(np.load(output / "pred" / f"{name}.npy").shape, (20, 3, 10)) self.assertEqual({p.stem for p in (output / "pred").glob("*.npy")}, set(self.test)) if __name__ == "__main__": unittest.main()