Update landmarkdiff/checkpoint_manager.py to v0.3.2

landmarkdiff/checkpoint_manager.py

@@ -106,6 +106,19 @@ class CheckpointManager:
                 self._index = json.load(f)
             if "checkpoints" not in self._index:
                 self._index["checkpoints"] = {}
+            # Remove entries whose directories no longer exist on disk
+            # (can happen after a crash during pruning)
+            missing = [
+                name
+                for name in list(self._index["checkpoints"])
+                if not (self.output_dir / name).exists()
+            ]
+            if missing:
+                for name in missing:
+                    self._index["checkpoints"].pop(name, None)
+                self._update_best()
+                self._save_index()
+                self._update_symlinks()
 
     def _save_index(self) -> None:
         with open(self._index_path(), "w") as f:
@@ -166,7 +179,9 @@ class CheckpointManager:
         torch.save(state, ckpt_dir / "training_state.pt")
 
         # Compute checkpoint size
-        size_mb = sum(f.stat().st_size for f in ckpt_dir.rglob("*") if f.is_file()) / (1024 * 1024)
+        size_mb = sum(
+            f.stat().st_size for f in ckpt_dir.rglob("*") if f.is_file()
+        ) / (1024 * 1024)
 
         # Create metadata
         meta = CheckpointMetadata(
@@ -214,7 +229,7 @@ class CheckpointManager:
         entries.sort(key=lambda x: x[1], reverse=not self.lower_is_better)
 
         # Mark best
-        best_names = {e[0] for e in entries[: self.keep_best]}
+        best_names = {e[0] for e in entries[:self.keep_best]}
         for name, meta in self._index["checkpoints"].items():
             meta["is_best"] = name in best_names
 
@@ -243,11 +258,11 @@ class CheckpointManager:
             val = meta.get("metrics", {}).get(self.metric)
             if val is None:
                 continue
-            if (
-                best_val is None
-                or (self.lower_is_better and val < best_val)
-                or (not self.lower_is_better and val > best_val)
-            ):
+            if best_val is None:
+                best, best_val = name, val
+            elif self.lower_is_better and val < best_val:
+                best, best_val = name, val
+            elif not self.lower_is_better and val > best_val:
                 best, best_val = name, val
         return best
 
@@ -278,7 +293,7 @@ class CheckpointManager:
         keep = set()
 
         # Keep latest
-        for name in all_names[-self.keep_latest :]:
+        for name in all_names[-self.keep_latest:]:
             keep.add(name)
 
         # Keep best
@@ -292,7 +307,7 @@ class CheckpointManager:
                 ckpt_dir = self.output_dir / name
                 if ckpt_dir.exists():
                     shutil.rmtree(ckpt_dir)
-                del self._index["checkpoints"][name]
+                self._index["checkpoints"].pop(name, None)
 
         self._save_index()
 
@@ -321,7 +336,10 @@ class CheckpointManager:
 
     def total_size_mb(self) -> float:
         """Return total disk size of all tracked checkpoints."""
-        return sum(meta.get("size_mb", 0.0) for meta in self._index["checkpoints"].values())
+        return sum(
+            meta.get("size_mb", 0.0)
+            for meta in self._index["checkpoints"].values()
+        )
 
     def summary(self) -> str:
         """Return a human-readable summary of checkpoint state."""
@@ -346,7 +364,6 @@ class CheckpointManager:
 # Helpers
 # ------------------------------------------------------------------
 
-
 def _get_state_dict(module: torch.nn.Module) -> dict:
     """Extract state dict, handling DDP wrapper."""
     if hasattr(module, "module"):