Teach retention about staged and abandoned checkpoints

tests/unit_tests/test_checkpoint.py

@@ -1232,13 +1232,55 @@ def test_purge_keeps_k_because_dcp_purges_after_saving(self, _rank):
 
         self.assertEqual({"/checkpoint/step-1"}, self._purged(manager))
 
-    def test_parse_step_reads_step_numbers_unanchored(self):
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_incomplete_directory_cannot_evict_a_valid_checkpoint(self, _rank):
+        # An interrupted save leaves a step-N directory with no metadata. If it
+        # occupied a slot it would push a checkpoint we can actually resume from
+        # out of the retained set.
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2", "step-3"])
+        manager._storage.isfile.side_effect = lambda path: "step-3" not in path
+
+        manager._purge_stale_checkpoints(is_save_in_flight=False)
+
+        # k-1 valid ones stay (step-2), the incomplete step-3 is deleted rather
+        # than counted, and step-1 falls out normally.
+        self.assertEqual({"/checkpoint/step-1"}, self._purged(manager))
+        manager._storage.remove.assert_called_once_with("/checkpoint/step-3")
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_in_flight_checkpoint_is_not_treated_as_abandoned(self, _rank):
+        # This manager purges after dispatching its save, so the newest
+        # directory is legitimately incomplete and must not be deleted.
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2", "step-3"])
+        manager._storage.isfile.side_effect = lambda path: "step-3" not in path
+
+        manager._purge_stale_checkpoints(is_save_in_flight=True)
+
+        self.assertEqual({"/checkpoint/step-1"}, self._purged(manager))
+        manager._storage.remove.assert_not_called()
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_abandoned_directories_are_deleted_synchronously(self, _rank):
+        # Queuing these would let the delete land after a retry at the same step
+        # recreated the directory it names.
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2"])
+        manager._storage.isfile.return_value = False
+
+        manager._purge_stale_checkpoints(is_save_in_flight=False)
+
+        self.assertEqual(set(), self._purged(manager))
+        self.assertEqual(
+            [mock.call("/checkpoint/step-1"), mock.call("/checkpoint/step-2")],
+            manager._storage.remove.call_args_list,
+        )
+
+    def test_parse_step_never_reports_a_staging_directory(self):
         manager = CheckpointManager.__new__(CheckpointManager)
 
-        self.assertEqual(7, manager._parse_step("step-7"))
+        self.assertEqual((7, False), manager._parse_step("step-7"))
         self.assertIsNone(manager._parse_step("logs"))
         # Unanchored, matching what this manager has always accepted.
-        self.assertEqual(9, manager._parse_step("step-9.partial"))
+        self.assertEqual((9, False), manager._parse_step("step-9.partial"))
 
     def test_valid_checkpoint_accepts_dcp_or_hf_markers(self):
         manager = CheckpointManager.__new__(CheckpointManager)

torchtitan/components/checkpointer/base.py

@@ -275,12 +275,13 @@ def _should_purge(self) -> bool:
         )
 
     @abstractmethod
-    def _parse_step(self, filename: str) -> int | None:
+    def _parse_step(self, filename: str) -> tuple[int, bool] | None:
         """Read ``filename`` as a checkpoint directory name.
 
-        Returns its step number, or ``None`` when the name is not one this
-        manager writes. Names a manager does not recognize are left alone
-        rather than deleted.
+        Returns the step number and whether the directory is a staging one that
+        a backend renames into place once a save completes, or ``None`` when the
+        name is not one this manager writes. Names a manager does not recognize
+        are left alone rather than deleted.
         """
 
     @abstractmethod
@@ -313,30 +314,80 @@ def _find_load_step(self, folder: str = "") -> int:
 
         valid_steps = []
         for filename in self._storage.listdir(folder):
-            step = self._parse_step(filename)
-            if step is None:
+            parsed = self._parse_step(filename)
+            if parsed is None:
+                continue
+            step, is_staging = parsed
+            # A staging directory may already hold its metadata and so look
+            # complete, but the id we would rebuild from its step -- the
+            # published name -- does not exist yet.
+            if is_staging:
                 continue
             if self._is_valid_checkpoint(filesystem.join(folder, filename)):
                 valid_steps.append(step)
         return max(valid_steps) if valid_steps else -1
 
-    def _purge_stale_checkpoints(self) -> None:
-        """Delete the checkpoints beyond the ``keep_latest_k`` most recent."""
+    def _purge_stale_checkpoints(self, *, is_save_in_flight: bool = True) -> None:
+        """Delete the checkpoints beyond the ``keep_latest_k`` most recent.
+
+        Args:
+            is_save_in_flight: Whether a save is being written right now. True
+                for a manager that purges after dispatching its save; False for
+                one that purges before issuing it.
+
+        That single fact settles both of the things this method has to get
+        right, because both follow from it.
+
+        Slot accounting. With a save in flight its directory is already on disk
+        and counted, so all k slots go to what is there. Without one, the
+        imminent checkpoint has no directory yet, so a slot is held for it and
+        only k-1 existing ones stay -- the accounting the config assumes where
+        it rejects ``keep_latest_k == 1`` because "the last one may be in the
+        process of being saved".
+
+        What occupies a slot. A directory with no metadata cannot be resumed
+        from, so letting it hold a slot would evict one that can. With a save in
+        flight the newest directory is legitimately incomplete, so completeness
+        is not required and nothing is deleted outright. Without one, every
+        incomplete directory is abandoned: skipped for accounting, and removed.
+        """
         if not self._should_purge():
             return
 
-        discovered: list[tuple[int, str]] = []
+        durable: list[tuple[int, str]] = []
+        abandoned: list[str] = []
         for filename in self._storage.listdir(self.folder):
-            step = self._parse_step(filename)
-            if step is None:
+            parsed = self._parse_step(filename)
+            if parsed is None:
                 continue
-            discovered.append((step, filesystem.join(self.folder, filename)))
-
-        discovered.sort()
-        for _, path in discovered[: -self.keep_latest_k]:
+            step, is_staging = parsed
+            path = filesystem.join(self.folder, filename)
+            if is_save_in_flight:
+                durable.append((step, path))
+            elif is_staging or not self._is_valid_checkpoint(path):
+                abandoned.append(path)
+            else:
+                durable.append((step, path))
+
+        durable.sort()
+        retain = self.keep_latest_k - (0 if is_save_in_flight else 1)
+        for _, path in durable[:-retain]:
             assert self.purge_thread is not None
             self.purge_queue.put(path)
 
+        # Deleted here rather than queued. The reason these are safe to remove
+        # is that no save is being written, and that is only true at this
+        # instant: a queued delete could land after the next attempt at the same
+        # step has recreated the very directory it names.
+        for path in abandoned:
+            logger.info("Checkpointer is deleting the abandoned %s.", path)
+            try:
+                self._storage.remove(path)
+            except Exception as error:
+                logger.warning(
+                    "Checkpointer failed to delete %s: %s. Skipping.", path, error
+                )
+
     @dataclass(kw_only=True, slots=True)
     class Config(Configurable.Config):
         """Checkpoint policies shared by concrete TorchTitan checkpoint managers."""

torchtitan/components/checkpointer/dcp.py

@@ -662,11 +662,12 @@ def _wait_for_saving(self) -> None:
         self.save_future.result()
         self.save_future = None
 
-    def _parse_step(self, filename: str) -> int | None:
-        # Deliberately unanchored, matching what this manager has always
-        # accepted.
+    def _parse_step(self, filename: str) -> tuple[int, bool] | None:
+        # This manager writes each checkpoint straight to its published name, so
+        # it has no staging directories and never reports one. The match is
+        # deliberately unanchored, matching what it has always accepted.
         match = re.search(r"step-(\d+)", filename)
-        return None if match is None else int(match.group(1))
+        return None if match is None else (int(match.group(1)), False)
 
     def _is_valid_checkpoint(self, checkpoint_id: str) -> bool:
         # Either format DCP can read: a native DCP checkpoint or a HuggingFace

torchtitan/components/checkpointer/torch_checkpointing.py

@@ -221,11 +221,11 @@ def _wait_for_saving(self) -> None:
             "TorchCheckpointingManager does not implement saving yet."
         )
 
-    def _parse_step(self, filename: str) -> int | None:
+    def _parse_step(self, filename: str) -> tuple[int, bool] | None:
         match = self._step_dir_pattern.fullmatch(filename)
-        if match is None or match.group("tmp"):
+        if match is None:
             return None
-        return int(match.group("step"))
+        return int(match.group("step")), bool(match.group("tmp"))
 
     def _is_valid_checkpoint(self, checkpoint_id: str) -> bool:
         return self._storage.isfile(