From 5a5a6a04ef40b5a2a991e0521e06bd5d92c9e719 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 20 Oct 2026 10:27:29 -0400 Subject: [PATCH 1/2] trainHydra: record the split and Lance versions a run read Every run writes split.json (train/valid dataset keys per train dataset) and, with Lance datasets, lance_versions.json (table versions and their tags) next to git_sha.txt; a requeue or resume that differs keeps the first file and writes a timestamped one. Resolving an untagged Lance version warns that the run is reproducible only until cleanup_old_versions drops it. Resuming from a checkpoint whose run validated on episodes this run trains on warns. Claude-Session: https://claude.ai/code/session_01Kt1E8T2x8kdV9akAvbJyVd --- egomimic/rldb/lance/dataset.py | 23 ++++++++++ egomimic/rldb/multi.py | 7 ++- egomimic/trainHydra.py | 10 +++++ egomimic/utils/run_manifest.py | 79 +++++++++++++++++++++++++++++++++ tests/unit/test_run_manifest.py | 70 +++++++++++++++++++++++++++++ 5 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 egomimic/utils/run_manifest.py create mode 100644 tests/unit/test_run_manifest.py diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py index 9ea873bff..69ff0cff4 100644 --- a/egomimic/rldb/lance/dataset.py +++ b/egomimic/rldb/lance/dataset.py @@ -678,6 +678,19 @@ def _open(self, name: str): self.versions[name] = want return uri, want, _table(uri, want) + def manifest(self) -> dict: + """What this resolver read: {"root", "versions": {table: version}, "tags": + {table: [tags on that version]}}.""" + tags = { + name: sorted( + t + for t, info in _table(f"{self.root}/{name}", v).tags.list().items() + if info["version"] == v + ) + for name, v in self.versions.items() + } + return {"root": self.root, "versions": dict(self.versions), "tags": tags} + def recordings(self) -> pd.DataFrame: """One row per recording.""" _, _, rec = self._open("recordings.lance") @@ -886,6 +899,16 @@ def resolve( ), **ann_kw, ) + untagged = sorted(n for n, t in self.manifest()["tags"].items() if not t) + if untagged: + logger.warning( + "read untagged versions of %s under %s (%s): this run can be " + "reproduced only until cleanup_old_versions drops them; tag them " + "(commit.tag) to keep them", + untagged, + self.root, + {n: self.versions[n] for n in untagged}, + ) if no_anchor: logger.warning( "%d recordings have no sampleable anchors, e.g. %s", diff --git a/egomimic/rldb/multi.py b/egomimic/rldb/multi.py index 1dda62bed..9409e45db 100644 --- a/egomimic/rldb/multi.py +++ b/egomimic/rldb/multi.py @@ -269,6 +269,7 @@ def __init__( # Dedup bounds-check warnings: keyed by f"bounds:{episode}:{zarr_key}" self._warned_violations: set[str] = set() self.train_collections: set = set() + self.source_manifest: dict | None = None self.valid_collections: set = set() if state is not None: @@ -677,7 +678,11 @@ def _from_resolver(cls, resolver, dataset_name: str | None = None, **kwargs): "were skipped by the resolver (see the log above)." ) - return cls(datasets=resolved, **kwargs) + ds = cls(datasets=resolved, **kwargs) + manifest = getattr(resolver, "manifest", None) + # Which table versions the leaves came from; trainHydra writes it to the run dir. + ds.source_manifest = manifest() if manifest is not None else None + return ds # ===================================================================== # Stats / normalization diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 41454404e..9bcba82a1 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -23,6 +23,7 @@ from egomimic.rldb.filters import PinError from egomimic.rldb.multi import EvenStrideDataset, MultiDataset, pinned_episode_subset from egomimic.rldb.resolve_memo import resolve_once +from egomimic.utils import run_manifest from egomimic.utils.checkpoint_utils import load_checkpoint_weights from egomimic.utils.compile_cache import set_per_job_compile_cache_dir from egomimic.utils.env import load_env @@ -804,6 +805,15 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: log_hyperparameters(object_dict) os.makedirs(os.path.join(trainer.default_root_dir, "videos"), exist_ok=True) + if trainer.is_global_zero: + manifest = run_manifest.run_manifest( + datamodule.train_datasets, datamodule.valid_datasets + ) + if mode == "train": + run_manifest.warn_on_split_drift( + cfg.get("ckpt_path"), manifest[run_manifest.SPLIT_FILE] + ) + run_manifest.write(trainer.default_root_dir, manifest) if mode == "train": if cfg.get("evaluator") is not None: diff --git a/egomimic/utils/run_manifest.py b/egomimic/utils/run_manifest.py new file mode 100644 index 000000000..e041241da --- /dev/null +++ b/egomimic/utils/run_manifest.py @@ -0,0 +1,79 @@ +"""Record which data a run trained on, next to its checkpoints: ``split.json`` +(train / valid dataset keys of each train dataset) and ``lance_versions.json`` +(the Lance table versions each dataset read).""" + +from __future__ import annotations + +import json +import logging +from datetime import datetime, timezone +from pathlib import Path + +log = logging.getLogger(__name__) + +SPLIT_FILE = "split.json" +LANCE_FILE = "lance_versions.json" + + +def run_manifest(train_datasets: dict, valid_datasets: dict) -> dict[str, dict]: + """{file name: contents} for ``write``; no Lance file without a Lance dataset.""" + split = { + name: { + "train": sorted(ds.train_collections), + "valid": sorted(ds.valid_collections), + } + for name, ds in train_datasets.items() + if hasattr(ds, "train_collections") + } + lance = { + f"{head}/{name}": ds.source_manifest + for head, datasets in (("train", train_datasets), ("valid", valid_datasets)) + for name, ds in datasets.items() + if getattr(ds, "source_manifest", None) + } + return {SPLIT_FILE: split, **({LANCE_FILE: lance} if lance else {})} + + +def write(root: str | Path, files: dict[str, dict]) -> None: + """Write each file into ``root``. A file already there with other contents (a + requeue or a resume into the same run dir) is kept, and the new one goes to + ``..json``, so the launch that trained the earlier checkpoints + keeps its record.""" + root = Path(root) + root.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + for name, contents in files.items(): + path = root / name + if path.exists(): + if json.loads(path.read_text(encoding="utf-8")) == contents: + continue + path = root / f"{path.stem}.{stamp}{path.suffix}" + path.write_text(json.dumps(contents, indent=1) + "\n", encoding="utf-8") + log.info("run manifest: wrote %s", path) + + +def warn_on_split_drift(ckpt_path: str | None, split: dict) -> list[str]: + """Warn about, and return, the keys the checkpoint's run validated on that this + run trains on: dataset keys added since then reshuffle ``split_dataset_names``. + Reads ``split.json`` of the run that wrote ``ckpt_path`` (``/checkpoints/``).""" + if not ckpt_path or not Path(ckpt_path).is_file(): + return [] + prev_file = Path(ckpt_path).resolve().parent.parent / SPLIT_FILE + if not prev_file.is_file(): + return [] + prev = json.loads(prev_file.read_text(encoding="utf-8")) + leaked = [] + for name, cur in split.items(): + moved = sorted(set(prev.get(name, {}).get("valid", [])) & set(cur["train"])) + if moved: + log.warning( + "dataset %r: %d episode(s) the checkpoint's run validated on are in " + "this run's train split (dataset keys changed since %s), e.g. %s; " + "its val metrics no longer measure held-out data", + name, + len(moved), + prev_file, + moved[:5], + ) + leaked += moved + return leaked diff --git a/tests/unit/test_run_manifest.py b/tests/unit/test_run_manifest.py new file mode 100644 index 000000000..a4647b79f --- /dev/null +++ b/tests/unit/test_run_manifest.py @@ -0,0 +1,70 @@ +import json + +import pytest + +lance = pytest.importorskip("lance") + +from fixtures.lance_tables import write_catalog # noqa: E402 + +from egomimic.rldb.lance import commit # noqa: E402 +from egomimic.rldb.lance.dataset import LanceEpisodeResolver # noqa: E402 +from egomimic.rldb.multi import MultiDataset # noqa: E402 +from egomimic.utils import run_manifest # noqa: E402 + +KEY_MAP = { + "left.action_ee_pose": { + "key_type": "action_keys", + "zarr_key": "left.obs_ee_pose", + "horizon": 3, + }, +} + + +def _dataset(root, **kw): + resolver = LanceEpisodeResolver(root, key_map=KEY_MAP, cameras={}, **kw) + return MultiDataset._from_resolver(resolver, mode="total", valid_ratio=0.5) + + +def test_manifest_names_the_split_and_the_versions_read(tmp_path, caplog): + root = tmp_path / "lance" + recs = write_catalog(root) + with caplog.at_level("WARNING"): + ds = _dataset(root) + assert "untagged" in caplog.text + files = run_manifest.run_manifest({"human_bimanual": ds}, {"human_bimanual": ds}) + split = files[run_manifest.SPLIT_FILE]["human_bimanual"] + assert sorted(split["train"] + split["valid"]) == sorted(r for r, _, _ in recs) + read = files[run_manifest.LANCE_FILE]["train/human_bimanual"] + assert read["versions"]["steps/aria.lance"] == 1 + assert read["tags"]["recordings.lance"] == [] + + commit.tag(root, "t1") + caplog.clear() + with caplog.at_level("WARNING"): + tagged = _dataset(root, tag="t1") + assert "untagged" not in caplog.text + assert tagged.source_manifest["tags"]["steps/aria.lance"] == ["t1"] + + +def test_write_keeps_the_first_launchs_record(tmp_path): + run_manifest.write(tmp_path, {"split.json": {"a": 1}}) + run_manifest.write(tmp_path, {"split.json": {"a": 1}}) + assert [p.name for p in tmp_path.iterdir()] == ["split.json"] + run_manifest.write(tmp_path, {"split.json": {"a": 2}}) + assert json.loads((tmp_path / "split.json").read_text()) == {"a": 1} + assert len(list(tmp_path.glob("split.*.json"))) == 1 + + +def test_resume_warns_when_old_val_episodes_now_train(tmp_path, caplog): + run = tmp_path / "run" + (run / "checkpoints").mkdir(parents=True) + ckpt = run / "checkpoints" / "last.ckpt" + ckpt.write_bytes(b"") + run_manifest.write(run, {"split.json": {"h": {"train": ["a"], "valid": ["b"]}}}) + same = {"h": {"train": ["a", "c"], "valid": ["b"]}} + assert run_manifest.warn_on_split_drift(str(ckpt), same) == [] + moved = {"h": {"train": ["a", "b"], "valid": ["c"]}} + with caplog.at_level("WARNING"): + assert run_manifest.warn_on_split_drift(str(ckpt), moved) == ["b"] + assert "validated on" in caplog.text + assert run_manifest.warn_on_split_drift(None, moved) == [] From 1ef7b5ed3acd36ea54db8a1db22e8bdd75d7a63d Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Fri, 23 Oct 2026 11:43:28 -0400 Subject: [PATCH 2/2] tests: run manifest names the per-(lab, embodiment, end_effector) steps table Steps tables moved to steps/{lab}/{embodiment}/{end_effector}.lance in the Lance schema change below this PR; the manifest test keyed versions by the old steps/aria.lance path. Claude-Session: https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw --- tests/unit/test_run_manifest.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_run_manifest.py b/tests/unit/test_run_manifest.py index a4647b79f..1a0a81113 100644 --- a/tests/unit/test_run_manifest.py +++ b/tests/unit/test_run_manifest.py @@ -8,9 +8,11 @@ from egomimic.rldb.lance import commit # noqa: E402 from egomimic.rldb.lance.dataset import LanceEpisodeResolver # noqa: E402 +from egomimic.rldb.lance.schema import steps_path # noqa: E402 from egomimic.rldb.multi import MultiDataset # noqa: E402 from egomimic.utils import run_manifest # noqa: E402 +STEPS = steps_path("rl2", "human_bimanual", "hand") KEY_MAP = { "left.action_ee_pose": { "key_type": "action_keys", @@ -35,7 +37,7 @@ def test_manifest_names_the_split_and_the_versions_read(tmp_path, caplog): split = files[run_manifest.SPLIT_FILE]["human_bimanual"] assert sorted(split["train"] + split["valid"]) == sorted(r for r, _, _ in recs) read = files[run_manifest.LANCE_FILE]["train/human_bimanual"] - assert read["versions"]["steps/aria.lance"] == 1 + assert read["versions"][STEPS] == 1 assert read["tags"]["recordings.lance"] == [] commit.tag(root, "t1") @@ -43,7 +45,7 @@ def test_manifest_names_the_split_and_the_versions_read(tmp_path, caplog): with caplog.at_level("WARNING"): tagged = _dataset(root, tag="t1") assert "untagged" not in caplog.text - assert tagged.source_manifest["tags"]["steps/aria.lance"] == ["t1"] + assert tagged.source_manifest["tags"][STEPS] == ["t1"] def test_write_keeps_the_first_launchs_record(tmp_path):