Move checkpoint discovery and retention into BaseCheckpointManager

tests/unit_tests/test_checkpoint.py

@@ -1194,6 +1194,68 @@ def test_absent_folder_does_not_purge(self, _rank):
         manager._storage.isdir.assert_called_once_with("/checkpoint")
 
 
+class TestSharedDiscoveryAndRetention(unittest.TestCase):
+    """_find_load_step and _purge_stale_checkpoints live on the base; each
+    manager supplies only _parse_step and _is_valid_checkpoint."""
+
+    def _manager(self, *, keep_latest_k: int, entries: list[str]):
+        manager = CheckpointManager.__new__(CheckpointManager)
+        manager.keep_latest_k = keep_latest_k
+        manager.folder = "/checkpoint"
+        manager.purge_queue = queue_lib.Queue()
+        manager.purge_thread = object()
+        manager._storage = mock.Mock(spec=CheckpointStorage)
+        manager._storage.isdir.return_value = True
+        manager._storage.listdir.return_value = entries
+        manager._storage.isfile.return_value = True
+        return manager
+
+    def _purged(self, manager) -> set[str]:
+        purged = set()
+        while not manager.purge_queue.empty():
+            purged.add(manager.purge_queue.get_nowait())
+        return purged
+
+    def test_bodies_are_defined_on_the_base(self):
+        for name in ("_find_load_step", "_purge_stale_checkpoints"):
+            with self.subTest(name=name):
+                self.assertNotIn(name, vars(CheckpointManager))
+                self.assertIn(name, vars(BaseCheckpointManager))
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_purge_keeps_k_because_dcp_purges_after_saving(self, _rank):
+        # This manager purges once its checkpoint is already on disk, so it
+        # reserves nothing and keeps the full k.
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2", "step-3"])
+
+        manager._purge_stale_checkpoints()
+
+        self.assertEqual({"/checkpoint/step-1"}, self._purged(manager))
+
+    def test_parse_step_reads_step_numbers_unanchored(self):
+        manager = CheckpointManager.__new__(CheckpointManager)
+
+        self.assertEqual(7, 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"))
+
+    def test_valid_checkpoint_accepts_dcp_or_hf_markers(self):
+        manager = CheckpointManager.__new__(CheckpointManager)
+        manager._storage = mock.Mock(spec=CheckpointStorage)
+
+        for marker in (".metadata", "model.safetensors.index.json"):
+            with self.subTest(marker=marker):
+                manager._storage.isfile.side_effect = (
+                    lambda path, marker=marker: path.endswith(marker)
+                )
+                self.assertTrue(manager._is_valid_checkpoint("/checkpoint/step-1"))
+
+        manager._storage.isfile.side_effect = None
+        manager._storage.isfile.return_value = False
+        self.assertFalse(manager._is_valid_checkpoint("/checkpoint/step-1"))
+
+
 class TestPurgeThread(unittest.TestCase):
     """A single failed deletion must not kill the daemon purge thread; otherwise
     keep_latest_k would silently stop purging for the rest of the run."""

torchtitan/components/checkpointer/base.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import queue
+import threading
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Callable
@@ -197,6 +198,8 @@ class BaseCheckpointManager(Configurable, ABC):
     save_future: Future | None
     folder: str
     keep_latest_k: int
+    purge_thread: threading.Thread | None
+    purge_queue: queue.Queue[str | None]
     _storage: CheckpointStorage
 
     # A disabled manager returns early from ``__init__`` without setting up any
@@ -271,6 +274,69 @@ def _should_purge(self) -> bool:
             and self._storage.isdir(self.folder)
         )
 
+    @abstractmethod
+    def _parse_step(self, filename: str) -> int | 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.
+        """
+
+    @abstractmethod
+    def _is_valid_checkpoint(self, checkpoint_id: str) -> bool:
+        """Whether ``checkpoint_id`` holds a checkpoint this manager can load.
+
+        A directory whose save was interrupted exists but has no metadata, so
+        resuming from it would fail; this is what keeps it out of
+        ``_find_load_step``.
+        """
+
+    def _find_load_step(self, folder: str = "") -> int:
+        """The highest step in ``folder`` that can actually be loaded.
+
+        Args:
+            folder: Directory to scan. Defaults to ``self.folder``.
+
+        Returns:
+            The step number, or -1 when the folder holds no loadable checkpoint.
+
+        Note:
+            This is not remote friendly: it issues one listdir plus a metadata
+            probe per step folder, each a network round trip on remote (fsspec)
+            storage instead of a single batched listing. Acceptable for now
+            since it only runs once at load time.
+        """
+        folder = folder or self.folder
+        if not self._storage.isdir(folder):
+            return -1
+
+        valid_steps = []
+        for filename in self._storage.listdir(folder):
+            step = self._parse_step(filename)
+            if step is None:
+                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."""
+        if not self._should_purge():
+            return
+
+        discovered: list[tuple[int, str]] = []
+        for filename in self._storage.listdir(self.folder):
+            step = self._parse_step(filename)
+            if step is None:
+                continue
+            discovered.append((step, filesystem.join(self.folder, filename)))
+
+        discovered.sort()
+        for _, path in discovered[: -self.keep_latest_k]:
+            assert self.purge_thread is not None
+            self.purge_queue.put(path)
+
     @dataclass(kw_only=True, slots=True)
     class Config(Configurable.Config):
         """Checkpoint policies shared by concrete TorchTitan checkpoint managers."""

torchtitan/components/checkpointer/dcp.py

@@ -662,51 +662,20 @@ def _wait_for_saving(self) -> None:
         self.save_future.result()
         self.save_future = None
 
-    def _find_load_step(self, folder: str = "") -> int:
-        """Identify the highest available checkpoint step in the specified directory.
-
-        This method scans the target folder for subdirectories matching the
-        'step-N' pattern. A folder is only considered a valid checkpoint if
-        it contains either a DCP metadata file or a HuggingFace safetensors
-        index.
-
-        Args:
-            folder (str, optional): The directory to scan. Defaults to `self.folder`.
-
-        Returns:
-            int: The maximum step number found among valid checkpoints,
-                or -1 if no valid checkpoints are detected.
-
-        Note:
-            This function is not remote friendly: it issues one listdir plus
-            up to two isfile probes per step folder, each a network round trip
-            on remote (fsspec) storage instead of a single batched listing.
-            Acceptable for now since it only runs once at load time.
-        """
-
-        folder = folder or self.folder
-        if not self._storage.isdir(folder):
-            return -1
-
-        pattern = r"step-(\d+)"
-        valid_steps = []
-
-        for filename in self._storage.listdir(folder):
-            match = re.search(pattern, filename)
-            if not match:
-                continue
-
-            # A checkpoint is valid only if it contains core metadata
-            checkpoint_path = filesystem.join(folder, filename)
-            is_dcp = self._storage.isfile(filesystem.join(checkpoint_path, ".metadata"))
-            is_hf = self._storage.isfile(
-                filesystem.join(checkpoint_path, "model.safetensors.index.json")
-            )
-
-            if is_dcp or is_hf:
-                valid_steps.append(int(match.group(1)))
-
-        return max(valid_steps) if valid_steps else -1
+    def _parse_step(self, filename: str) -> int | None:
+        # Deliberately unanchored, matching what this manager has always
+        # accepted.
+        match = re.search(r"step-(\d+)", filename)
+        return None if match is None else int(match.group(1))
+
+    def _is_valid_checkpoint(self, checkpoint_id: str) -> bool:
+        # Either format DCP can read: a native DCP checkpoint or a HuggingFace
+        # safetensors directory.
+        return self._storage.isfile(
+            filesystem.join(checkpoint_id, ".metadata")
+        ) or self._storage.isfile(
+            filesystem.join(checkpoint_id, "model.safetensors.index.json")
+        )
 
     def _create_checkpoint_id(self, step: int, folder: str = "") -> str:
         """Generate the standardized filesystem path for a checkpoint
@@ -833,21 +802,3 @@ def _should_save(self, curr_step: int, last_step: bool = False) -> bool:
             return True
 
         return False
-
-    def _purge_stale_checkpoints(self):
-        """Remove older checkpoint directories from storage to maintain
-        only the most recent 'k' copies."""
-        if self._should_purge():
-            discovered_checkpoints = []
-            for filename in self._storage.listdir(self.folder):
-                match = re.search(r"step-(\d+)", filename)
-                if match:
-                    path = filesystem.join(self.folder, filename)
-                    discovered_checkpoints.append((int(match.group(1)), path))
-
-            discovered_checkpoints.sort()
-            to_delete = discovered_checkpoints[: -1 * self.keep_latest_k]
-
-            for _, path in to_delete:
-                assert self.purge_thread is not None
-                self.purge_queue.put(path)

torchtitan/components/checkpointer/torch_checkpointing.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import os
+import re
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -19,6 +20,9 @@
 from torch_checkpointing.checkpoint_writer import CheckpointWriterConfig
 from torch_checkpointing.config import AsyncCheckpointSaverConfig
 from torch_checkpointing.default_resharder import DefaultResharder
+from torch_checkpointing.distributed_metadata import (
+    METADATA_FILE_NAME as TORCH_CHECKPOINTING_METADATA_FILE_NAME,
+)
 from torch_checkpointing.schema import ItemSpec
 from torch_checkpointing.staging import CheckpointStagerConfig
 from torch_checkpointing.storage.base_storage import Storage
@@ -43,6 +47,11 @@
 _DEFAULT_BARRIER_TIMEOUT_SEC = 600
 
 
+def _step_dir_pattern(temp_dir_prefix: str) -> re.Pattern[str]:
+    """Match published and in-flight checkpoint directory names."""
+    return re.compile(rf"(?P<tmp>{re.escape(temp_dir_prefix)})?step-(?P<step>\d+)")
+
+
 class _BackendCheckpointStorage:
     """``CheckpointStorage`` backed by a ``torch_checkpointing`` ``Storage``.
 
@@ -183,6 +192,9 @@ def __init__(
             )
 
         manager_config = _default_backend_config()
+        self._step_dir_pattern = _step_dir_pattern(
+            manager_config.save.writer_config.temp_dir_prefix
+        )
         storage_config = manager_config.storage_config or LocalFileSystemStorageConfig()
         self._storage = _BackendCheckpointStorage(storage_config.create_storage())
         self._manager = manager_config.build()
@@ -209,6 +221,17 @@ def _wait_for_saving(self) -> None:
             "TorchCheckpointingManager does not implement saving yet."
         )
 
+    def _parse_step(self, filename: str) -> int | None:
+        match = self._step_dir_pattern.fullmatch(filename)
+        if match is None or match.group("tmp"):
+            return None
+        return int(match.group("step"))
+
+    def _is_valid_checkpoint(self, checkpoint_id: str) -> bool:
+        return self._storage.isfile(
+            filesystem.join(checkpoint_id, TORCH_CHECKPOINTING_METADATA_FILE_NAME)
+        )
+
     def _maybe_wait_for_staging(self) -> None:
         raise NotImplementedError(
             "TorchCheckpointingManager does not implement maybe_wait_for_staging() yet."