Support remote fsspec checkpoint paths (e.g. gs://) via filesystem helpers

tests/unit_tests/test_checkpoint.py

@@ -5,19 +5,28 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
+import queue as queue_lib
 import shutil
 import tempfile
 import time
 import unittest
+import uuid
 from concurrent.futures import Future
 from types import SimpleNamespace
 from unittest import mock
 
+import fsspec
 import torch
 import torch.nn as nn
 from torch.distributed.checkpoint.state_dict_saver import AsyncSaveResponse
 from torch.utils.data import DataLoader
-from torchtitan.components.checkpoint import CheckpointManager, MODEL, ModelWrapper
+from torchtitan.components.checkpoint import (
+    CheckpointManager,
+    MODEL,
+    ModelWrapper,
+    purge_thread,
+    Terminate,
+)
 
 
 class FakeOptimizersContainer:
@@ -790,13 +799,17 @@ def test_sanity_and_range_checks(self):
             CheckpointManager.Config(exclude_from_loading=[MODEL])
 
     def test_path_normalization(self):
-        """Test that paths are stripped and must be absolute."""
-        # Test leading/trailing whitespace stripping
+        """initial_load_path is stripped and must be absolute or a remote URI."""
+        # Leading/trailing whitespace is stripped.
         cfg = CheckpointManager.Config(initial_load_path="  /absolute/path/step-100  ")
         self.assertEqual(cfg.initial_load_path, "/absolute/path/step-100")
 
-        # Test relative path rejection
-        with self.assertRaisesRegex(ValueError, "must be absolute"):
+        # A remote URI is accepted (and stripped).
+        cfg = CheckpointManager.Config(initial_load_path="  gs://bucket/run/step-100  ")
+        self.assertEqual(cfg.initial_load_path, "gs://bucket/run/step-100")
+
+        # A bare relative path is rejected.
+        with self.assertRaisesRegex(ValueError, "absolute path or a remote"):
             CheckpointManager.Config(initial_load_path="relative/path/step-100")
 
     def test_dependency_assertions(self):
@@ -820,6 +833,26 @@ def test_dependency_assertions(self):
         with self.assertRaisesRegex(ValueError, "requires last_save_model_only=True"):
             CheckpointManager.Config(last_save_in_hf=True, last_save_model_only=False)
 
+    def test_remote_hf_rejected(self):
+        """HF safetensors format is not supported with remote (fsspec) paths."""
+        with self.assertRaisesRegex(
+            ValueError, "last_save_in_hf is not supported with a remote"
+        ):
+            CheckpointManager.Config(
+                folder="gs://bucket/run",
+                last_save_in_hf=True,
+                last_save_model_only=True,
+            )
+
+        with self.assertRaisesRegex(
+            ValueError, "initial_load_in_hf is not supported with a remote"
+        ):
+            CheckpointManager.Config(
+                initial_load_in_hf=True,
+                initial_load_model_only=True,
+                initial_load_path="gs://bucket/run/step-1",
+            )
+
     def test_mode_normalization(self):
         """Test that async_mode is case-normalized."""
         cfg = CheckpointManager.Config(async_mode="ASYNC")
@@ -845,6 +878,66 @@ def test_warnings(self, mock_logger):
         )
 
 
+class TestFindLoadStepRemote(unittest.TestCase):
+    """_find_load_step must discover checkpoints on a remote (fsspec) folder,
+    exercising listdir-basename reduction and per-dir metadata probing over the
+    in-memory backend (no gcsfs/network)."""
+
+    def setUp(self):
+        self.name = f"test-{uuid.uuid4().hex}"
+        self.root = f"memory://{self.name}"
+        self._fs = fsspec.filesystem("memory")
+
+    def tearDown(self):
+        if self._fs.exists(f"/{self.name}"):
+            self._fs.rm(f"/{self.name}", recursive=True)
+
+    def _write(self, url):
+        with fsspec.open(url, "wb") as f:
+            f.write(b"x")
+
+    def test_returns_max_valid_step(self):
+        self._write(f"{self.root}/step-10/.metadata")
+        self._write(f"{self.root}/step-20/.metadata")
+        # step-30 has no core metadata -> not a valid checkpoint.
+        self._write(f"{self.root}/step-30/some_shard")
+        # 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)
+
+    def test_missing_folder_returns_negative_one(self):
+        manager = CheckpointManager.__new__(CheckpointManager)
+        self.assertEqual(manager._find_load_step(folder=self.root), -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."""
+
+    def test_survives_failed_delete(self):
+        q = queue_lib.Queue()
+        q.put("fails")
+        q.put("succeeds")
+        q.put(Terminate())
+
+        calls = []
+
+        def rmtree(path):
+            calls.append(path)
+            if path == "fails":
+                raise RuntimeError("transient backend error")
+
+        with mock.patch(
+            "torchtitan.components.checkpoint.filesystem.rmtree", side_effect=rmtree
+        ):
+            # Returns once Terminate is dequeued; must not raise.
+            purge_thread(q)
+
+        self.assertEqual(calls, ["fails", "succeeds"])
+
+
 class TestModelWrapper(unittest.TestCase):
     """ModelWrapper.state_dict() must keep stable tensor storage across calls so
     the async pinned-memory stager (keyed by source storage) reuses its host

tests/unit_tests/test_filesystem.py

@@ -0,0 +1,122 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import tempfile
+import unittest
+import uuid
+from unittest import mock
+
+import fsspec
+
+from torchtitan.tools import filesystem
+
+
+class TestIsRemote(unittest.TestCase):
+    def test_remote_uris(self):
+        self.assertTrue(filesystem.is_remote("gs://bucket/run"))
+        self.assertTrue(filesystem.is_remote("s3://bucket/run"))
+
+    def test_local_paths(self):
+        self.assertFalse(filesystem.is_remote("/abs/path"))
+        self.assertFalse(filesystem.is_remote("./relative"))
+        self.assertFalse(filesystem.is_remote("plain"))
+
+
+class TestJoin(unittest.TestCase):
+    def test_local_folder_is_prefixed(self):
+        self.assertEqual(filesystem.join("/dump", "checkpoint"), "/dump/checkpoint")
+
+    def test_remote_folder_is_not_prefixed(self):
+        self.assertEqual(
+            filesystem.join("/dump", "gs://bucket/run/ckpt"),
+            "gs://bucket/run/ckpt",
+        )
+
+
+class TestLocalOps(unittest.TestCase):
+    def test_local_filesystem_operations(self):
+        with tempfile.TemporaryDirectory() as d:
+            step = os.path.join(d, "step-10")
+            os.makedirs(step)
+            meta = os.path.join(step, ".metadata")
+            with open(meta, "w") as f:
+                f.write("x")
+
+            self.assertTrue(filesystem.exists(step))
+            self.assertTrue(filesystem.isdir(step))
+            self.assertTrue(filesystem.isfile(meta))
+            self.assertFalse(filesystem.isfile(step))
+            self.assertEqual(filesystem.listdir(d), ["step-10"])
+
+            filesystem.rmtree(step)
+            self.assertFalse(filesystem.exists(step))
+
+            # rmtree on a missing path is a no-op (ignore_errors=True).
+            filesystem.rmtree(step)
+
+
+class TestRemoteOps(unittest.TestCase):
+    """Exercise the fsspec path with the in-memory backend (no gcsfs/network)."""
+
+    def setUp(self):
+        # Unique root so tests do not interfere via the process-shared memory fs.
+        self.name = f"test-{uuid.uuid4().hex}"
+        self.root = f"memory://{self.name}"
+        self._fs = fsspec.filesystem("memory")
+
+    def tearDown(self):
+        if self._fs.exists(f"/{self.name}"):
+            self._fs.rm(f"/{self.name}", recursive=True)
+
+    def _write(self, url):
+        with fsspec.open(url, "wb") as f:
+            f.write(b"x")
+
+    def test_exists_isdir_isfile(self):
+        self._write(f"{self.root}/step-10/.metadata")
+        self.assertTrue(filesystem.exists(f"{self.root}/step-10"))
+        self.assertTrue(filesystem.isdir(f"{self.root}/step-10"))
+        self.assertTrue(filesystem.isfile(f"{self.root}/step-10/.metadata"))
+        self.assertFalse(filesystem.exists(f"{self.root}/step-99"))
+
+    def test_listdir_returns_basenames(self):
+        self._write(f"{self.root}/step-10/.metadata")
+        self._write(f"{self.root}/step-20/.metadata")
+        self.assertEqual(sorted(filesystem.listdir(self.root)), ["step-10", "step-20"])
+
+    def test_rmtree(self):
+        self._write(f"{self.root}/step-10/.metadata")
+        filesystem.rmtree(self.root)
+        self.assertFalse(filesystem.exists(self.root))
+
+    def test_rmtree_missing_is_noop(self):
+        # A missing remote path must not raise (parity with local rmtree),
+        # otherwise a failed purge would kill the background purge thread.
+        filesystem.rmtree(f"{self.root}/does-not-exist")
+
+
+class TestListdirDirectoryMarker(unittest.TestCase):
+    """Object stores (GCS, S3) can list a directory-marker entry for the listed
+    directory itself; ``os.listdir`` never returns the directory, so ``listdir``
+    must drop it while keeping a child that shares the parent's basename."""
+
+    def test_self_marker_is_dropped_by_full_path(self):
+        fake_fs = mock.Mock()
+        fake_fs.ls.return_value = [
+            "bucket/ckpt",  # directory marker for the listed dir itself
+            "bucket/ckpt/",  # same self-entry with a trailing slash
+            "bucket/ckpt/step-10",
+            "bucket/ckpt/step-20",
+            "bucket/ckpt/ckpt",  # legitimate child sharing the parent basename
+        ]
+        with mock.patch.object(
+            filesystem, "_resolve", return_value=(fake_fs, "bucket/ckpt")
+        ):
+            result = filesystem.listdir("gs://bucket/ckpt")
+
+        self.assertEqual(sorted(result), ["ckpt", "step-10", "step-20"])
+        fake_fs.ls.assert_called_once_with("bucket/ckpt", detail=False)