trainHydra: record the split and Lance versions a run read

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",

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

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:

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
+    ``<stem>.<UTC time>.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`` (``<run>/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

tests/unit/test_run_manifest.py

@@ -0,0 +1,72 @@
+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.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",
+        "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] == 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] == ["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) == []