diff --git a/tests/unit_tests/test_checkpoint.py b/tests/unit_tests/test_checkpoint.py index 9dec02abcd..dc814700b2 100644 --- a/tests/unit_tests/test_checkpoint.py +++ b/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) diff --git a/torchtitan/components/checkpointer/base.py b/torchtitan/components/checkpointer/base.py index e5ea6068bc..61f1d3b86c 100644 --- a/torchtitan/components/checkpointer/base.py +++ b/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.""" diff --git a/torchtitan/components/checkpointer/dcp.py b/torchtitan/components/checkpointer/dcp.py index 678b4cbaf3..36ffb89ae3 100644 --- a/torchtitan/components/checkpointer/dcp.py +++ b/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 diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py index f2aa79e95c..8b9ebd38da 100644 --- a/torchtitan/components/checkpointer/torch_checkpointing.py +++ b/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(