Merge pull request #168 from BitMind-AI/fix/parallel-reader-legacy-resume

Support audited checkpoint-reader upgrades for existing runs

Author: dylanuys

pyproject.toml

@@ -63,6 +63,7 @@ where = ["src"]
 "" = "src"
 
 [tool.setuptools.package-data]
+"gasbench" = ["checkpoint_compatibility.json"]
 "gasbench.dataset" = ["configs/*.yaml"]
 
 [tool.black]

src/gasbench/benchmarks/common.py

@@ -275,6 +275,20 @@ def fingerprint_files(paths, root):
     return digest.hexdigest()
 
 
+def compatible_evaluator_identity(evaluator_dir, current, previous):
+    """Allow only explicitly audited, exact-build checkpoint reader upgrades."""
+    if current == previous:
+        return True
+    registry = Path(evaluator_dir) / "checkpoint_compatibility.json"
+    if not registry.is_file():
+        return False
+    try:
+        pairs = json.loads(registry.read_text())
+    except (OSError, ValueError):
+        return False
+    return isinstance(pairs, dict) and isinstance(pairs.get(current), list) and previous in pairs[current]
+
+
 def runtime_versions():
     versions = {
         "python": platform.python_version(),
@@ -461,6 +475,12 @@ def create_tracker(
     context = json.loads(json.dumps(context))
     if saved is not None:
         previous = saved.get("benchmark", {})
+        if compatible_evaluator_identity(
+            evaluator_dir, context["evaluator_sha256"], previous.get("evaluator_sha256")
+        ):
+            # Preserve the original manifest identity for both restored and new batches.
+            # Every other context field is still compared below without exceptions.
+            context["evaluator_sha256"] = previous["evaluator_sha256"]
         if {k: v for k, v in previous.items() if k != "samples"} != context:
             raise CheckpointError(
                 "Checkpoint configuration or model differs from this run"

src/gasbench/checkpoint_compatibility.json

@@ -0,0 +1,5 @@
+{
+  "cefcbe4740db068e9e34d3f4184cbde9d65b441482bc2ecdbcfeb98584634288": [
+    "ec5643b53236c144c9ed296f653a1dd39b852ee37ce4feaefa2179650b1ad2cb"
+  ]
+}

tests/unit/test_benchmark_resume.py

@@ -366,3 +366,34 @@ def test_augmentation_cache_cannot_change_across_attempts(benchmark, tmp_path):
         path.write_bytes(b"changed")
     with pytest.raises(CheckpointError, match="augmentation cache"):
         b.run(Session(b.model_dir), n_aug_per_dataset=3, aug_cache_dir=str(aug_dir))
+
+
+@pytest.mark.parametrize("change", ["reader", "unapproved-evaluator", "model"])
+def test_audited_reader_upgrade_preserves_resume_guards(benchmark, monkeypatch, tmp_path, change):
+    b = benchmark
+    original_fingerprint = common.fingerprint_files
+    evaluator_root = Path(common.__file__).resolve().parents[1]
+    identity = ["old"]
+
+    def fingerprint(paths, root):
+        return identity[0] if root == evaluator_root else original_fingerprint(paths, root)
+
+    monkeypatch.setattr(common, "fingerprint_files", fingerprint)
+    (tmp_path / "checkpoint_compatibility.json").write_text(json.dumps({"new": ["old"]}))
+    compatible = common.compatible_evaluator_identity
+    monkeypatch.setattr(common, "compatible_evaluator_identity", lambda root, current, previous: compatible(tmp_path, current, previous))
+    with pytest.raises(Interrupted):
+        b.run(Session(b.model_dir, stop_after=2))
+    identity[0] = "unapproved" if change == "unapproved-evaluator" else "new"
+    if change == "model":
+        (b.model_dir / "model.py").write_text("changed inference implementation")
+    resumed = Session(b.model_dir)
+    if change != "reader":
+        with pytest.raises(CheckpointError, match="differs"):
+            b.run(resumed)
+        assert not resumed.calls
+    else:
+        b.run(resumed)
+        assert len(resumed.calls) == 3
+        assert RecorderCheckpoint.read_manifest(b.checkpoint)["benchmark"]["evaluator_sha256"] == "old"
+        assert not compatible(tmp_path, "new", "unrelated-old")