Fix checkpoint rotation after successful saves

[hjinnkim]

Summary

Checkpoint retention runs before checkpoint_state_save(). Rotating first, then writing, is
the wrong order for every consequence below; this PR writes first and rotates after the save
succeeds, and protects the checkpoint that was just written.

Measured on main (2d6df1245) by driving the real Trainer._run_standard_checkpoint against a
minimal stub:

• Steady-state disk usage is limit + 1. CheckpointManager.cleanup_checkpoints
  (checkpoint_manager.py:205-207) trims down to exactly limit, and the save at
  trainer.py:5250-5252 then adds one more. With checkpoints_total_limit=3 over six cycles the
  on-disk count goes [1, 2, 3, 4, 4, 4]. After the change it is [1, 2, 3, 3, 3, 3].
• A stale higher-numbered checkpoint permanently occupies a retention slot. Resuming from a
  lower step and continuing leaves the directory settled at
  ['checkpoint-400', 'checkpoint-5000'] forever, because numeric rotation keeps the largest step
  rather than the one just written. If a later save then fails, the only complete checkpoint left
  is the stale step-5000 one and recovery silently jumps 4900 steps. After the change the same
  sequence settles at ['checkpoint-400'], and the failed-save case leaves
  ['checkpoint-200', 'checkpoint-300-tmp'].
• checkpoints_total_limit=0 deletes every checkpoint. The guard at trainer.py:5244 is
  checkpoints_total_limit is not None, so 0 passes through to a cleanup that reads it as
  keep-zero. Measured: three existing checkpoints plus one save leaves ['checkpoint-400']. The
  field registry documents the opposite —
  field_registry/sections/training.py:198-212 carries
  ValidationRule(MIN, 0, 'Must be non-negative (0 = unlimited)') and the tooltip
  "Set to 0 for unlimited". After the change all four are retained.
• The background Hub upload re-resolves latest instead of using the checkpoint it was
  scheduled for. huggingface.py:301 calls find_latest_checkpoint() inside a
  single-worker background future (trainer.py:2112, :5262); measured resolving to
  checkpoint-200 for an upload scheduled with global_step=100.
• Rolling -tmp directories are mis-parsed, and it destroys data. Suffix is read as
  parts[2] (checkpoint_manager.py:249) / cs[2] (trainer.py:5865), so
  checkpoint-200-rolling-tmp is classified as suffix rolling. Measured:
  cleanup_checkpoints(1, "rolling") deleted the real checkpoint-200-rolling and kept the
  orphaned checkpoint-200-rolling-tmp.

One thing this PR also fixes, which is not reachable in a shipped configuration: the fallback
rotation path in trainer.py (checkpoint_state_cleanup's non-manager branch,
trainer.py:5878-5902) returns early when len(checkpoints) < limit and otherwise removes
len - limit + 1, so at limit=1 it deletes the only checkpoint before the save and a failed
save leaves nothing. That branch is dead today — trainer.py:2206-2208 constructs a
CheckpointManager whenever config.output_dir is truthy and the field default is
simpletuner-results — so this is a latent correctness fix, not an observable bug. It is listed
here so the - limit + 1 → - limit change is not mistaken for an off-by-one repair: that +1
was the correct compensation for pre-save ordering, and it changes only because the ordering is
being inverted.

Changes

• Clean stale temporary checkpoints before saving, but rotate completed standard and rolling
  checkpoints only after the new save succeeds.
• Protect the newly written save_path during rotation so it survives even when its step number
  is lower than another checkpoint on disk.
• CheckpointManager.cleanup_checkpoints() and the fallback path take an optional
  protected_checkpoint and use post-save retention counts.
• Keep temporary-checkpoint cleanup active when checkpoints_total_limit=0 while leaving
  completed-checkpoint retention unlimited, matching the documented meaning of 0.
• Pass the exact saved checkpoint path to the Hub uploader instead of re-resolving latest in the
  background task.
• Recognise both standard and rolling *-tmp directories during cleanup (parts[-1] / cs[-1]).
• Extract _save_rolling_checkpoint() so the rolling path follows the same order as the standard
  path.

Two changes that are not consequences of the above

Both are deliberate and both change runtime behaviour, so they are called out separately rather
than buried in the list.

A synchronous drain of pending Hub uploads before rotation. Rotation may delete a directory
that a background upload is still reading, so the new code blocks on
_drain_hub_upload_futures(wait=True) when the post-save count exceeds the limit. At steady state
that count is always limit + 1, so this fires on every checkpoint cycle.

main already has a blocking drain, but only in teardown: _finish_hub_uploads()
(trainer.py:2194-2198) is called from the finally of the run wrapper (:1908) and once more
after the final model upload (:7131). Inside the training loop the only drain is
trainer.py:2189 with wait=False, immediately before a new upload is scheduled. This PR
introduces the first blocking drain on the per-checkpoint path, and the uploader is a
single-worker ThreadPoolExecutor (trainer.py:2108-2113), so a slow push serialises into the
step loop.

The cost is confined to runs with push_to_hub enabled: with background push off,
_hub_upload_futures is empty and _drain_hub_upload_futures returns at its first line
(trainer.py:2149-2150). Reviewers who consider a per-checkpoint stall unacceptable should say so
— the alternative is to accept that rotation can delete a directory mid-upload.

A pre-save barrier for all-rank temporary-directory writes. When DeepSpeed or FSDP has every
rank saving and checkpointing_use_tempdir is set, the ranks must agree on the temp directory
before any of them writes into it, so checkpoint_state_save now calls
accelerator.wait_for_everyone() in exactly that case. It is a no-op for single-process runs and
for main-process-only saving.

Ordering constraint: please merge #2988 first

Do not merge this before #2988 (Fix RamTorch prefetch-order save race).
_save_rolling_checkpoint widens the rolling-save gate to
is_main_process or use_deepspeed_optimizer or fsdp_enable; main (trainer.py:6816-6825) has no
fsdp_enable there. With this PR alone, every FSDP rank newly enters accelerator.save_state on
the rolling path — and therefore reaches save_hooks.py:1095, where all ranks still race on one
shared ramtorch_prefetch_orders.json.tmp. Rolling checkpoints normally fire on a much tighter
interval than standard ones, so the exposure is not marginal. Merging #2988 first, or both
together, removes the interaction; the combined tree was verified green.

The two PRs touch disjoint files, so this is purely a sequencing constraint — there is no conflict
and no rebase needed either way.

Notes

• checkpoints_total_limit=0 remains unlimited for completed checkpoints while temp cleanup keeps
  running.
• Checkpoint format and resume-selection behaviour are unchanged.
• Distributed behaviour is covered with CPU mocks. No live multi-rank or GPU run was performed for
  this branch.

Verification

CPU-only container, current main (2d6df1245), 8 new test methods / 18 subtest cases in
tests/test_trainer.py.

tests.test_trainer          RED  Ran 82   FAILED (failures=10, errors=8)
                            GREEN Ran 82   OK

18 related modules          RED  Ran 236  FAILED (failures=10, errors=8)
                            GREEN Ran 236  OK

full suite, 317 modules     RED  Ran 4222 FAILED (failures=19, errors=23, skipped=143)
                            GREEN Ran 4222 FAILED (failures=9,  errors=15, skipped=143)
                            GREEN-only failures: NONE

RED is main plus this branch's test changes with the three production files restored from main.
All 8 new methods are RED there — none is green-on-arrival. The 18-failure delta between RED and
GREEN on the full suite is exactly the new tests; the 24 that remain on both are pre-existing on
main.

Each separable sub-fix was additionally mutation-tested by reverting it alone inside the green
tree, to confirm nothing is unpinned:

parts[-1] -> parts[2]                      FAILED (failures=4)
drop protected_checkpoint filtering        FAILED (failures=4)
drop the new wait_for_everyone() barrier   FAILED (failures=4)
restore the `< limit` / `- limit + 1` path  FAILED (failures=1)

[Copilot]

Pull request overview

This PR corrects checkpoint retention behavior by saving checkpoints first and rotating only after a successful save, while protecting the newly-written checkpoint from deletion. It also improves temp-checkpoint handling, aligns checkpoints_total_limit=0 with documented “unlimited” semantics for completed checkpoints, and makes Hub uploads use the exact checkpoint path scheduled for upload.

Changes:

• Reorders standard and rolling checkpoint flows to: temp cleanup → save → post-save rotation (with optional protected checkpoint).
• Updates checkpoint parsing to correctly recognize *-tmp for both standard and rolling checkpoints (using the last - segment).
• Passes an explicit checkpoint path to Hub upload and adds tests covering rotation/retention edge cases and distributed tempdir synchronization.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
tests/test_trainer.py	Adds regression tests for temp cleanup, post-save rotation behavior, protected checkpoints, Hub upload path usage, rolling checkpoint behavior, and distributed tempdir barriers.
simpletuner/helpers/utils/checkpoint_manager.py	Adds protected_checkpoint support during cleanup and fixes suffix parsing to correctly classify *-tmp checkpoints.
simpletuner/helpers/training/trainer.py	Implements post-save cleanup/rotation ordering, temp-only cleanup helper, protected checkpoint retention, explicit Hub upload path, rolling checkpoint refactor, and a distributed pre-save barrier for tempdir writes.
simpletuner/helpers/publishing/huggingface.py	Allows upload_latest_checkpoint() to accept an explicit checkpoint_path to avoid re-resolving latest in background upload tasks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.