Repository: huggingface/transformers
Pull request: [Trainer] Move sort and rotate checkpoints to standalone functions

# What does this PR do?

This PR do the following:

- Extract `sorted_checkpoints` and `rotate_checkpoints` from Trainer class to standalone functions in `trainer_utils.py`
- Rename `sorted_checkpoints` to `sort_checkpoints`
- Simplify implementations for better readability
- Add and update tests 

Part of PRs to unbloat Trainer

Changed files:
filename: src/transformers/trainer.py
patch: @@ -24,7 +24,6 @@
 import math
 import os
 import random
-import re
 import shutil
 import sys
 import tempfile
@@ -124,8 +123,10 @@
     load_sharded_checkpoint,
     neftune_post_forward_hook,
     number_of_arguments,
+    rotate_checkpoints,
     seed_worker,
     set_seed,
+    sort_checkpoints,
     speed_metrics,
 )
 from .training_args import OptimizerNames, ParallelMode, TrainingArguments
@@ -2681,7 +2682,9 @@ def _inner_training_loop(
         self.log(metrics)
 
         run_dir = self._get_output_dir(trial)
-        checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir)
+        checkpoints_sorted = sort_checkpoints(
+            output_dir=run_dir, best_model_checkpoint=self.state.best_model_checkpoint
+        )
 
         # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save.
         if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
@@ -3167,8 +3170,13 @@ def _save_checkpoint(self, model, trial):
 
         # Maybe delete some older checkpoints.
         if self.args.should_save:
-            # we use mtime as default, filesystems without mtime support will be detected in `_sorted_checkpoints`
-            self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)
+            # we use mtime as default, filesystems without mtime support will be detected in `sort_checkpoints`
+            rotate_checkpoints(
+                output_dir=run_dir,
+                save_total_limit=self.args.save_total_limit,
+                best_model_checkpoint=self.state.best_model_checkpoint,
+                use_mtime=True,
+            )
 
     def _save_rng_state(self, output_dir):
         # Save RNG state in non-distributed training
@@ -4158,68 +4166,6 @@ def store_flos(self):
             self.state.total_flos += self.current_flos
             self.current_flos = 0
 
-    def _sorted_checkpoints(
-        self, output_dir=None, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False
-    ) -> list[str]:
-        ordering_and_checkpoint_path = []
-
-        glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]
-
-        for path in glob_checkpoints:
-            if use_mtime:
-                ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
-            else:
-                regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
-                if regex_match is not None and regex_match.groups() is not None:
-                    ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
-
-        checkpoints_sorted = sorted(ordering_and_checkpoint_path)
-        # mtime is not reliable on all filesystems, especially on some fuse fs in cloud environments
-        # so we check if the mtime is fake and fallback to numerical ordering if needed
-        if use_mtime and len(ordering_and_checkpoint_path) > 1:
-            mtime_diff = checkpoints_sorted[-1][0] - checkpoints_sorted[0][0]
-            if mtime_diff < 1.0:  # less than 1 second, which is almost impossible when mtime works fine
-                warnings.warn("mtime may not be reliable on this filesystem, falling back to numerical ordering")
-                return self._sorted_checkpoints(
-                    use_mtime=False, output_dir=output_dir, checkpoint_prefix=checkpoint_prefix
-                )
-        checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
-
-        # Make sure we don't delete the best model.
-        if (
-            self.state.best_model_checkpoint is not None
-            and str(Path(self.state.best_model_checkpoint)) in checkpoints_sorted
-        ):
-            best_model_index = checkpoints_sorted.index(str(Path(self.state.best_model_checkpoint)))
-            for i in range(best_model_index, len(checkpoints_sorted) - 2):
-                checkpoints_sorted[i], checkpoints_sorted[i + 1] = checkpoints_sorted[i + 1], checkpoints_sorted[i]
-        return checkpoints_sorted
-
-    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
-        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
-            return
-
-        # Check if we should delete older checkpoint(s)
-        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
-        if len(checkpoints_sorted) <= self.args.save_total_limit:
-            return
-
-        # If save_total_limit=1 with load_best_model_at_end=True, we could end up deleting the last checkpoint, which
-        # we don't do to allow resuming.
-        save_total_limit = self.args.save_total_limit
-        if (
-            self.state.best_model_checkpoint is not None
-            and self.args.save_total_limit == 1
-            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
-        ):
-            save_total_limit = 2
-
-        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
-        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
-        for checkpoint in checkpoints_to_be_deleted:
-            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
-            shutil.rmtree(checkpoint, ignore_errors=True)
-
     def evaluate(
         self,
         eval_dataset: Dataset | dict[str, Dataset] | None = None,

filename: src/transformers/trainer_utils.py
patch: @@ -23,10 +23,12 @@
 import os
 import random
 import re
+import shutil
 import threading
 import time
 from collections.abc import Callable
 from functools import partial
+from pathlib import Path
 from typing import Any, NamedTuple
 
 import numpy as np
@@ -46,10 +48,14 @@
     is_torch_npu_available,
     is_torch_xla_available,
     is_torch_xpu_available,
+    logging,
     requires_backends,
 )
 
 
+logger = logging.get_logger(__name__)
+
+
 if is_torch_available():
     import torch
     from safetensors.torch import load_file as safe_load_file
@@ -218,6 +224,113 @@ def get_last_checkpoint(folder):
     return os.path.join(folder, max(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0])))
 
 
+def sort_checkpoints(
+    output_dir: str,
+    checkpoint_prefix: str = PREFIX_CHECKPOINT_DIR,
+    use_mtime: bool = False,
+    best_model_checkpoint: str | None = None,
+) -> list[str]:
+    """
+    Return checkpoint directories sorted by step number (oldest first).
+
+    Args:
+        output_dir (`str`):
+            The directory containing the checkpoints.
+        checkpoint_prefix (`str`, *optional*, defaults to `"checkpoint"`):
+            The prefix used for checkpoint directory names.
+        use_mtime (`bool`, *optional*, defaults to `False`):
+            Whether to sort by modification time instead of step number.
+        best_model_checkpoint (`str`, *optional*):
+            If provided, this checkpoint is moved to second-to-last position to protect
+            it from deletion while keeping the most recent checkpoint last for resuming.
+
+    Returns:
+        `list[str]`: Sorted list of checkpoint directory paths (oldest first).
+    """
+    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]
+
+    ordering_and_checkpoint_path = []
+    for path in glob_checkpoints:
+        if use_mtime:
+            ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
+        else:
+            regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
+            if regex_match is not None and regex_match.groups() is not None:
+                ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+
+    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
+
+    # mtime is not reliable on some filesystems (e.g., cloud fuse filesystems)
+    # so we check if the mtime is fake and fall back to numerical ordering
+    if use_mtime and len(checkpoints_sorted) > 1:
+        mtime_diff = checkpoints_sorted[-1][0] - checkpoints_sorted[0][0]
+        if mtime_diff < 1.0:
+            logger.warning("mtime may not be reliable on this filesystem, falling back to numerical ordering")
+            return sort_checkpoints(
+                output_dir, checkpoint_prefix, use_mtime=False, best_model_checkpoint=best_model_checkpoint
+            )
+
+    checkpoints_sorted = [path for _, path in checkpoints_sorted]
+
+    # Move best_model_checkpoint to second-to-last position to protect it from deletion
+    # while keeping the most recent checkpoint at the end for resuming training.
+    if best_model_checkpoint is not None:
+        best_model_checkpoint = str(Path(best_model_checkpoint))
+        if best_model_checkpoint in checkpoints_sorted and checkpoints_sorted[-1] != best_model_checkpoint:
+            most_recent = checkpoints_sorted[-1]
+            checkpoints_sorted = [c for c in checkpoints_sorted if c not in {best_model_checkpoint, most_recent}]
+            checkpoints_sorted += [best_model_checkpoint, most_recent]
+
+    return checkpoints_sorted
+
+
+def rotate_checkpoints(
+    output_dir: str,
+    save_total_limit: int | None = None,
+    best_model_checkpoint: str | None = None,
+    use_mtime: bool = False,
+    checkpoint_prefix: str = PREFIX_CHECKPOINT_DIR,
+) -> None:
+    """
+    Delete older checkpoints, keeping at most `save_total_limit`.
+
+    Always preserves the most recent checkpoint and the best model checkpoint (if provided).
+
+    Args:
+        output_dir (`str`):
+            The directory containing the checkpoints.
+        save_total_limit (`int`, *optional*):
+            Maximum number of checkpoints to keep. No deletion if `None` or <= 0.
+        best_model_checkpoint (`str`, *optional*):
+            Path to best checkpoint (will always be preserved).
+        use_mtime (`bool`, *optional*, defaults to `False`):
+            Whether to sort by modification time instead of step number.
+        checkpoint_prefix (`str`, *optional*, defaults to `"checkpoint"`):
+            The prefix used for checkpoint directory names.
+    """
+    if save_total_limit is None or save_total_limit <= 0:
+        return
+
+    checkpoints = sort_checkpoints(output_dir, checkpoint_prefix, use_mtime)
+    if len(checkpoints) <= save_total_limit:
+        return
+
+    # Checkpoints that must not be deleted
+    protected = {checkpoints[-1]}  # most recent, for resuming
+    if best_model_checkpoint is not None:
+        protected.add(str(Path(best_model_checkpoint)))
+
+    # Delete oldest non-protected checkpoints until we have save_total_limit left
+    num_to_keep = max(save_total_limit, len(protected))
+    remaining = len(checkpoints)
+    for checkpoint in checkpoints:
+        if remaining <= num_to_keep:
+            break
+        if checkpoint not in protected:
+            shutil.rmtree(checkpoint, ignore_errors=True)
+            remaining -= 1
+
+
 class IntervalStrategy(ExplicitEnum):
     NO = "no"
     STEPS = "steps"

filename: tests/trainer/test_trainer.py
patch: @@ -108,6 +108,8 @@
     HPSearchBackend,
     check_target_module_exists,
     get_last_checkpoint,
+    rotate_checkpoints,
+    sort_checkpoints,
 )
 from transformers.training_args import OptimizerNames
 from transformers.utils import (
@@ -3945,11 +3947,38 @@ def assert_flos_extraction(trainer, wrapped_model_to_check):
             trainer.train()
             self.assertTrue(isinstance(trainer.state.total_flos, float))
 
+    def test_checkpoint_sorting(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Create fake checkpoints in non-sorted order
+            for n in [20, 5, 15, 25, 10]:
+                os.makedirs(os.path.join(tmp_dir, f"{PREFIX_CHECKPOINT_DIR}-{n}"))
+
+            # Test sorting by step number (oldest first)
+            sorted_cps = sort_checkpoints(tmp_dir)
+            values = [int(re.match(f".*{PREFIX_CHECKPOINT_DIR}-([0-9]+)", d).groups()[0]) for d in sorted_cps]
+            self.assertEqual(values, [5, 10, 15, 20, 25])
+
+            # Test with best_model_checkpoint - moved to second-to-last to protect from deletion
+            best = os.path.join(tmp_dir, f"{PREFIX_CHECKPOINT_DIR}-5")
+            sorted_cps = sort_checkpoints(tmp_dir, best_model_checkpoint=best)
+            values = [int(re.match(f".*{PREFIX_CHECKPOINT_DIR}-([0-9]+)", d).groups()[0]) for d in sorted_cps]
+            self.assertEqual(values, [10, 15, 20, 5, 25])
+
+            # Test with best_model_checkpoint already at end (stays at end)
+            best = os.path.join(tmp_dir, f"{PREFIX_CHECKPOINT_DIR}-25")
+            sorted_cps = sort_checkpoints(tmp_dir, best_model_checkpoint=best)
+            values = [int(re.match(f".*{PREFIX_CHECKPOINT_DIR}-([0-9]+)", d).groups()[0]) for d in sorted_cps]
+            self.assertEqual(values, [5, 10, 15, 20, 25])
+
     def check_checkpoint_deletion(self, trainer, output_dir, expected):
         # Make fake checkpoints
         for n in [5, 10, 15, 20, 25]:
             os.makedirs(os.path.join(output_dir, f"{PREFIX_CHECKPOINT_DIR}-{n}"), exist_ok=True)
-        trainer._rotate_checkpoints(output_dir=output_dir)
+        rotate_checkpoints(
+            output_dir=output_dir,
+            save_total_limit=trainer.args.save_total_limit,
+            best_model_checkpoint=trainer.state.best_model_checkpoint,
+        )
         glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{PREFIX_CHECKPOINT_DIR}-*")]
         values = [int(re.match(f".*{PREFIX_CHECKPOINT_DIR}-([0-9]+)", d).groups()[0]) for d in glob_checkpoints]
         self.assertSetEqual(set(values), set(expected))