Introduce a storage seam in BaseCheckpointManager

tests/unit_tests/test_checkpoint.py

@@ -25,11 +25,16 @@
 from torch.utils.data import DataLoader
 from torchtitan.components.checkpointer.base import (
     BaseCheckpointManager,
+    CheckpointStorage,
     MODEL,
     ModelWrapper,
     purge_thread,
 )
-from torchtitan.components.checkpointer.dcp import AsyncMode, CheckpointManager
+from torchtitan.components.checkpointer.dcp import (
+    _FilesystemCheckpointStorage,
+    AsyncMode,
+    CheckpointManager,
+)
 
 
 class FakeOptimizersContainer:
@@ -1071,6 +1076,13 @@ def _write(self, url):
         with fsspec.open(url, "wb") as f:
             f.write(b"x")
 
+    def _manager(self):
+        # The real storage adapter, so this keeps covering the fsspec path that
+        # torchtitan.tools.filesystem provides and the backend Storage does not.
+        manager = CheckpointManager.__new__(CheckpointManager)
+        manager._storage = _FilesystemCheckpointStorage()
+        return manager
+
     def test_returns_max_valid_step(self):
         self._write(f"{self.root}/step-10/.metadata")
         self._write(f"{self.root}/step-20/.metadata")
@@ -1079,12 +1091,107 @@ def test_returns_max_valid_step(self):
         # A non-checkpoint directory must be ignored.
         self._write(f"{self.root}/logs/events")
 
-        manager = CheckpointManager.__new__(CheckpointManager)
-        self.assertEqual(manager._find_load_step(folder=self.root), 20)
+        self.assertEqual(self._manager()._find_load_step(folder=self.root), 20)
 
     def test_missing_folder_returns_negative_one(self):
+        self.assertEqual(self._manager()._find_load_step(folder=self.root), -1)
+
+
+class TestFilesystemCheckpointStorage(unittest.TestCase):
+    """The DCP adapter must answer the CheckpointStorage protocol using
+    torchtitan.tools.filesystem, including for remote fsspec URIs."""
+
+    def setUp(self):
+        self.storage = _FilesystemCheckpointStorage()
+        self.root = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.root, ignore_errors=True)
+
+    def test_satisfies_the_protocol(self):
+        # A Protocol is structural, so nothing checks conformance at runtime.
+        # Assert it here or a renamed method would only surface as an
+        # AttributeError deep inside a save.
+        self.assertIsInstance(self.storage, CheckpointStorage)
+
+    def test_distinguishes_directories_from_files(self):
+        directory = os.path.join(self.root, "step-1")
+        os.makedirs(directory)
+        file_path = os.path.join(directory, ".metadata")
+        with open(file_path, "wb"):
+            pass
+
+        self.assertTrue(self.storage.isdir(directory))
+        self.assertFalse(self.storage.isfile(directory))
+        self.assertTrue(self.storage.isfile(file_path))
+        self.assertFalse(self.storage.isdir(file_path))
+        self.assertFalse(self.storage.isdir(os.path.join(self.root, "absent")))
+        self.assertFalse(self.storage.isfile(os.path.join(self.root, "absent")))
+
+    def test_listdir_returns_entry_names(self):
+        os.makedirs(os.path.join(self.root, "step-1"))
+        os.makedirs(os.path.join(self.root, "step-2"))
+
+        self.assertEqual({"step-1", "step-2"}, set(self.storage.listdir(self.root)))
+
+    def test_remove_deletes_a_populated_directory(self):
+        directory = os.path.join(self.root, "step-1")
+        os.makedirs(directory)
+        with open(os.path.join(directory, ".metadata"), "wb"):
+            pass
+
+        self.storage.remove(directory)
+
+        self.assertFalse(os.path.exists(directory))
+
+    def test_reaches_remote_uris_through_fsspec(self):
+        name = f"test-{uuid.uuid4().hex}"
+        root = f"memory://{name}"
+        memory_fs = fsspec.filesystem("memory")
+        self.addCleanup(memory_fs.rm, f"/{name}", recursive=True)
+        with fsspec.open(f"{root}/step-1/.metadata", "wb") as f:
+            f.write(b"x")
+
+        self.assertTrue(self.storage.isdir(f"{root}/step-1"))
+        self.assertTrue(self.storage.isfile(f"{root}/step-1/.metadata"))
+        self.assertEqual(["step-1"], self.storage.listdir(root))
+
+
+class TestShouldPurge(unittest.TestCase):
+    """_should_purge lives on the base so every manager, and TorchFT's
+    narrowing override, share one definition of who deletes checkpoints."""
+
+    def _manager(self, *, keep_latest_k: int, folder_exists: bool = True):
         manager = CheckpointManager.__new__(CheckpointManager)
-        self.assertEqual(manager._find_load_step(folder=self.root), -1)
+        manager.keep_latest_k = keep_latest_k
+        manager.folder = "/checkpoint"
+        manager._storage = mock.Mock()
+        manager._storage.isdir.return_value = folder_exists
+        return manager
+
+    def test_defined_on_the_base_not_the_dcp_manager(self):
+        self.assertNotIn("_should_purge", vars(CheckpointManager))
+        self.assertIn("_should_purge", vars(BaseCheckpointManager))
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_rank_zero_with_retention_and_an_existing_folder_purges(self, _rank):
+        self.assertTrue(self._manager(keep_latest_k=2)._should_purge())
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_retention_disabled_does_not_purge(self, _rank):
+        self.assertFalse(self._manager(keep_latest_k=0)._should_purge())
+
+    @mock.patch("torch.distributed.get_rank", return_value=1)
+    def test_nonzero_rank_does_not_purge(self, _rank):
+        self.assertFalse(self._manager(keep_latest_k=2)._should_purge())
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_absent_folder_does_not_purge(self, _rank):
+        manager = self._manager(keep_latest_k=2, folder_exists=False)
+
+        self.assertFalse(manager._should_purge())
+
+        # Probed through the seam rather than the filesystem module, so a
+        # manager on non-local storage asks its own backend.
+        manager._storage.isdir.assert_called_once_with("/checkpoint")
 
 
 class TestPurgeThread(unittest.TestCase):

tests/unit_tests/test_torch_checkpointing.py

@@ -122,6 +122,28 @@ def test_default_backend_configuration_owns_schema_and_barrier(self) -> None:
             DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT,
         )
 
+    def test_remote_checkpoint_paths_are_rejected_at_construction(self) -> None:
+        # Path() would turn "gs://bucket/x" into "gs:/bucket/x". Rejecting the
+        # two roots up front is what lets every path derived from them be
+        # converted without a further check -- and a save with retention off
+        # reaches the backend having run no probe that could have caught it.
+        for field, kwargs in (
+            ("checkpoint.folder", {"folder": "gs://bucket/checkpoint"}),
+            (
+                "checkpoint.initial_load_path",
+                {"initial_load_path": "gs://bucket/pretrained"},
+            ),
+        ):
+            with self.subTest(field=field):
+                config = TorchCheckpointingManager.Config(
+                    enable=True,
+                    keep_latest_k=0,
+                    initial_load_model_only=True,
+                    **kwargs,
+                )
+                with self.assertRaisesRegex(ValueError, rf"{field}.*not yet supported"):
+                    self._build_manager(config)
+
     def test_legacy_config_has_no_backend_selector(self) -> None:
         config = CheckpointManager.Config()
 

torchtitan/components/checkpointer/__init__.py

@@ -6,6 +6,7 @@
 
 from .base import (
     BaseCheckpointManager,
+    CheckpointStorage,
     DATALOADER,
     LR_SCHEDULER,
     MODEL,
@@ -19,6 +20,7 @@
     "AsyncMode",
     "BaseCheckpointManager",
     "CheckpointManager",
+    "CheckpointStorage",
     "DATALOADER",
     "LR_SCHEDULER",
     "MODEL",

torchtitan/components/checkpointer/base.py

@@ -12,9 +12,10 @@
 from collections.abc import Callable
 from concurrent.futures import Future
 from dataclasses import dataclass, field
-from typing import Any, Literal
+from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.distributed.tensor import DTensor
@@ -143,6 +144,46 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         self.cached_state_dict = self._get_state_dict()
 
 
+@runtime_checkable
+class CheckpointStorage(Protocol):
+    """The path operations a checkpoint manager needs from its storage.
+
+    Managers differ in how they read and write checkpoint bytes, but they ask
+    the same handful of questions about paths: is this a checkpoint directory,
+    did this metadata file land, which steps are on disk, delete this one. This
+    protocol is the whole of that surface, so policies like retention and
+    latest-step discovery can live on ``BaseCheckpointManager`` without knowing
+    which backend answers them.
+
+    Paths are ``str`` rather than ``Path`` because a checkpoint id may be a
+    remote URI (``gs://...``) that ``Path`` would mangle -- it collapses the
+    double slash. Carrying ``str`` keeps the vocabulary lossless; whether a
+    given implementation can actually reach a remote URI is up to that
+    implementation, which should reject what it cannot address rather than
+    silently rewrite it.
+
+    ``runtime_checkable`` so implementations can assert conformance in their
+    tests. It only checks that the method names exist, which is enough to catch
+    a rename that would otherwise surface as an ``AttributeError`` mid-save.
+    """
+
+    def isdir(self, path: str) -> bool:
+        """Whether ``path`` is an existing directory."""
+        ...
+
+    def isfile(self, path: str) -> bool:
+        """Whether ``path`` is an existing entry that is not a directory."""
+        ...
+
+    def listdir(self, path: str) -> list[str]:
+        """The entry names directly under the directory ``path``."""
+        ...
+
+    def remove(self, path: str) -> None:
+        """Recursively delete the directory ``path``."""
+        ...
+
+
 class BaseCheckpointManager(Configurable, ABC):
     """Contract every TorchTitan checkpoint manager implements.
 
@@ -154,6 +195,9 @@ class BaseCheckpointManager(Configurable, ABC):
 
     enable: bool
     save_future: Future | None
+    folder: str
+    keep_latest_k: int
+    _storage: CheckpointStorage
 
     # A disabled manager returns early from ``__init__`` without setting up any
     # state, so none of its attributes exist. The public entry points below own
@@ -219,6 +263,14 @@ def _maybe_wait_for_staging(self) -> None:
     def _close(self) -> None:
         """Implement ``close``. Only called when checkpointing is enabled."""
 
+    def _should_purge(self) -> bool:
+        """Whether this rank should purge stale checkpoints."""
+        return (
+            self.keep_latest_k > 0
+            and dist.get_rank() == 0
+            and self._storage.isdir(self.folder)
+        )
+
     @dataclass(kw_only=True, slots=True)
     class Config(Configurable.Config):
         """Checkpoint policies shared by concrete TorchTitan checkpoint managers."""