diff --git a/tests/unit_tests/test_checkpoint.py b/tests/unit_tests/test_checkpoint.py index 064b2e625c..9dec02abcd 100644 --- a/tests/unit_tests/test_checkpoint.py +++ b/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.""" diff --git a/torchtitan/components/checkpointer/base.py b/torchtitan/components/checkpointer/base.py index 7a197ec2e0..e5ea6068bc 100644 --- a/torchtitan/components/checkpointer/base.py +++ b/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.""" diff --git a/torchtitan/components/checkpointer/dcp.py b/torchtitan/components/checkpointer/dcp.py index f1277397a6..678b4cbaf3 100644 --- a/torchtitan/components/checkpointer/dcp.py +++ b/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) diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py index 9701f98879..f2aa79e95c 100644 --- a/torchtitan/components/checkpointer/torch_checkpointing.py +++ b/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{re.escape(temp_dir_prefix)})?step-(?P\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."