234 lines
12 KiB
Python
234 lines
12 KiB
Python
"""CPU checks for encoder-only packing and the unchanged four-epoch head."""
|
|
|
|
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_ego4d_encoder_only as runner
|
|
from test_local_turn_features import fixture
|
|
|
|
|
|
class EncoderOnlyWorkflowTests(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")
|
|
self.source = self.root / "local"
|
|
self.targets = self.root / "ego4d_target_frame"
|
|
self.targets.mkdir()
|
|
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)))
|
|
self.paths, self.payloads = {}, {}
|
|
for name, split in zip(self.names, self.splits):
|
|
# No global feature directory, raw video, annotation or model exists.
|
|
_, side, labels = fixture(name=name, dataset="ego4d")
|
|
side["reference_features"] = "/must-not-be-read/thinker.pt"
|
|
side["states"] = "POISON: never use this field as the encoder input"
|
|
path = self.source / split / f"{name}.pt"
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save(side, path)
|
|
np.save(self.targets / f"{name}.npy", labels)
|
|
self.paths[name], self.payloads[name] = path, side
|
|
|
|
def output_root(self, modality="av"):
|
|
return self.root / f"encoder_{modality}"
|
|
|
|
def prepare_args(self, modality="av", extra=()):
|
|
return ["--data-root", str(self.root), "--local-feature-root", str(self.source),
|
|
"--output-root", str(self.output_root(modality)), "--modalities", 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_pack_all_modalities_uses_only_native_vectors_and_validity(self):
|
|
name = self.names[0]
|
|
for modality, parts in MODALITIES.items():
|
|
with self.subTest(modality=modality):
|
|
result, contract = prepare.pack_sample(
|
|
self.paths[name], self.targets / f"{name}.npy", name, modality)
|
|
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(result["states"], expected))
|
|
self.assertTrue(torch.equal(result["times"], original["times"]))
|
|
self.assertTrue(torch.equal(result["label_indices"], original["label_indices"]))
|
|
self.assertEqual(result["mode"], f"encoder-only-{modality}")
|
|
self.assertIs(result["config"]["uses_thinker_states"], False)
|
|
self.assertEqual(contract["modalities"], modality)
|
|
|
|
def test_export_real_files_full_timeline_and_refuse_overwrite(self):
|
|
self.export()
|
|
output = self.output_root()
|
|
self.assertEqual((output / "complete.txt").read_text(), "3 clips\n")
|
|
manifest = json.loads((output / "export_manifest.json").read_text())
|
|
self.assertEqual(manifest["num_states"], 60)
|
|
self.assertEqual(manifest["input_width"], 13)
|
|
self.assertEqual(manifest["clips"], dict.fromkeys(self.names, 20))
|
|
for name in self.names:
|
|
path = output / f"{name}.pt"
|
|
self.assertFalse(path.is_symlink())
|
|
data = load_session_data(output, self.targets, name, dataset="ego4d", expected_mode="encoder-only-av")
|
|
self.assertEqual(data.states.shape, (20, 13))
|
|
self.assertIsNone(data.local_features)
|
|
before = (output / f"{self.names[0]}.pt").read_bytes()
|
|
with self.assertRaises(FileExistsError):
|
|
prepare.main(self.prepare_args())
|
|
self.assertEqual((output / f"{self.names[0]}.pt").read_bytes(), before)
|
|
|
|
def test_dry_run_does_not_write(self):
|
|
with patch.object(prepare.torch, "save") as save, contextlib.redirect_stdout(io.StringIO()):
|
|
prepare.main(self.prepare_args(extra=["--dry-run"]))
|
|
save.assert_not_called()
|
|
self.assertFalse(self.output_root().exists())
|
|
|
|
def test_corrupt_test_sidecars_rejected_before_output_creation(self):
|
|
name = self.names[1]
|
|
for issue in ("partial", "identity", "times", "indices", "contract", "nan", "invalid_mask"):
|
|
side = copy.deepcopy(self.payloads[name])
|
|
if issue == "partial":
|
|
side["partial"] = True
|
|
elif issue == "identity":
|
|
side["sample_id"] = self.names[0]
|
|
elif issue == "times":
|
|
side["times"][0] = .4
|
|
elif issue == "indices":
|
|
side["label_indices"][0] = 1
|
|
elif issue == "contract":
|
|
side["config"]["encoder_id"] = "different-encoder"
|
|
elif issue == "nan":
|
|
side["audio"][0, 0] = float("nan")
|
|
else:
|
|
side["visual"][0, 0] = 1 # Invalid visual rows must remain zero.
|
|
torch.save(side, self.paths[name])
|
|
with self.subTest(issue=issue), contextlib.redirect_stdout(io.StringIO()), self.assertRaises(ValueError):
|
|
prepare.main(self.prepare_args())
|
|
self.assertFalse(self.output_root().exists())
|
|
|
|
def test_target_mismatch_and_missing_sidecar_fail_before_writes(self):
|
|
name = self.names[1]
|
|
np.save(self.targets / f"{name}.npy", np.zeros((19, 3)))
|
|
with contextlib.redirect_stdout(io.StringIO()), self.assertRaisesRegex(ValueError, "feature/target mismatch"):
|
|
prepare.main(self.prepare_args())
|
|
self.assertFalse(self.output_root().exists())
|
|
self.paths[name].unlink()
|
|
with self.assertRaises(FileNotFoundError):
|
|
prepare.main(self.prepare_args())
|
|
|
|
def test_interrupted_export_has_no_complete_marker_or_temporary_file(self):
|
|
with patch.object(prepare.torch, "save", side_effect=RuntimeError("test write failure")), \
|
|
contextlib.redirect_stdout(io.StringIO()), self.assertRaisesRegex(RuntimeError, "write failure"):
|
|
prepare.main(self.prepare_args())
|
|
self.assertFalse((self.output_root() / "complete.txt").exists())
|
|
self.assertEqual(list(self.output_root().glob("*.tmp")), [])
|
|
|
|
def test_runner_commands_and_dry_run_never_select_thinker_or_dev(self):
|
|
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("--seed") + 1], str(seed))
|
|
self.assertEqual(train[train.index("--warmup-epochs") + 1], "1.6")
|
|
for command in (train, infer):
|
|
for forbidden in ("--local-feature-root", "--layer-fusion", "--dev-sessions", "--sessions"):
|
|
self.assertNotIn(forbidden, command)
|
|
self.assertEqual(command[command.index("--expected-mode") + 1], "encoder-only-av")
|
|
with patch.object(runner.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_missing_test_features_marker_and_wrong_contract(self):
|
|
self.export()
|
|
args = runner.parse_args(self.runner_args())
|
|
marker = self.output_root() / "complete.txt"
|
|
marker.write_text("2 clips\n")
|
|
with self.assertRaisesRegex(ValueError, "incomplete"):
|
|
runner.preflight(args)
|
|
marker.write_text("3 clips\n")
|
|
path = self.output_root() / f"{self.names[1]}.pt"
|
|
payload = torch.load(path, weights_only=False)
|
|
payload["config"]["uses_thinker_states"] = True
|
|
torch.save(payload, path)
|
|
with self.assertRaisesRegex(ValueError, "encoder-only export"):
|
|
runner.preflight(args)
|
|
path.unlink()
|
|
with self.assertRaises(FileNotFoundError):
|
|
runner.preflight(args)
|
|
|
|
def test_invalid_cli_and_existing_outputs_are_rejected(self):
|
|
for extra in (["--seeds", "0", "0"], ["--seeds", "-1"], ["--modalities", "thinker"]):
|
|
with self.subTest(extra=extra), contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
|
runner.parse_args(self.runner_args(extra=extra))
|
|
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
|
prepare.parse_args(self.prepare_args(extra=["--output-root", str(self.source)]))
|
|
self.export()
|
|
args = runner.parse_args(self.runner_args(extra=["--seeds", "0", "1"]))
|
|
checkpoint, _ = runner.run_paths(args, 1)
|
|
checkpoint.parent.mkdir()
|
|
checkpoint.write_bytes(b"preserve")
|
|
with self.assertRaises(FileExistsError):
|
|
runner.preflight(args)
|
|
self.assertEqual(checkpoint.read_bytes(), b"preserve")
|
|
|
|
def test_cpu_round_trip_all_three_modalities_four_epochs_without_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)
|
|
command = [sys.executable, str(Path(runner.__file__).resolve()), *self.runner_args(modality)]
|
|
result = subprocess.run(command, 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["epoch"], 4)
|
|
self.assertEqual(saved["dev_sessions"], [])
|
|
self.assertEqual(set(saved["train_sessions"]), {self.names[0], self.names[2]})
|
|
self.assertIsNone(saved["local_contract"])
|
|
width = sum(self.payloads[self.names[0]]["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.names[1]])
|
|
self.assertEqual(metrics["checkpoint_epoch"], 4)
|
|
self.assertEqual(np.load(output / "pred" / f"{self.names[1]}.npy").shape, (20, 3, 10))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|