From 18939d31935e67cb801760fc172bd276aa6bd3a5 Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Fri, 21 Aug 2026 18:30:59 -0700 Subject: [PATCH] Teach retention about staged and abandoned checkpoints Retention counts every step-N directory by name. That is right for a manager that writes each checkpoint straight to its published name, but not for one whose backend stages a save under a temporary name and renames it on success. Two things go wrong there. A directory left behind by a save that never got renamed holds nothing anyone can resume from, yet it occupies one of the k slots and evicts a checkpoint that can be resumed from. And the slot arithmetic assumes the new checkpoint is already on disk, which is false for a manager that purges before issuing its save. Both follow from one fact -- whether a save is being written right now -- so that is the single argument the caller passes. It is a caller's property, not the method's: the base does not implement _save, so it has no business knowing when a subclass purges, and keeping the value at the call site means moving the call cannot silently change retention. _parse_step widens to report whether a name is a staging directory. The DCP manager answers False always and passes the default, so its behavior is unchanged in every respect. Abandoned directories are deleted synchronously rather than queued. The reason they are safe to remove is that no save is being written, and that holds only at the instant of the check -- a queued delete could land after a retry at the same step recreated the directory it names. Test Plan: python3 -m pytest tests/unit_tests/test_checkpoint.py \ tests/unit_tests/test_torch_checkpointing.py \ torchtitan/experiments/torchft/tests/test_torchft_checkpoint.py -q -> 53 + 26 + 2 passed Tests cover an incomplete directory failing to evict a valid checkpoint, an in-flight directory not being mistaken for an abandoned one, and abandoned directories bypassing the purge queue. ufmt + flake8 clean. --- tests/unit_tests/test_checkpoint.py | 48 ++++++++++- torchtitan/components/checkpointer/base.py | 81 +++++++++++++++---- torchtitan/components/checkpointer/dcp.py | 9 ++- .../checkpointer/torch_checkpointing.py | 6 +- 4 files changed, 119 insertions(+), 25 deletions(-) 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(