[feat] support keep_checkpoint_max with async checkpoint pruning

[tiankongdeguiji]

What

Adds a keep_checkpoint_max training option that caps the number of retained model.ckpt-<step> directories. On long runs checkpoints currently accumulate without bound and can fill the disk; this lets users keep only the most recent N.

• TrainConfig.keep_checkpoint_max (default 0 = keep all → fully backward compatible; pruning is opt-in).
• Pruning is fully asynchronous: the training thread only enqueues a coalesced request; a single daemon worker does all filesystem work (directory glob, best-checkpoint lookup, recursive delete), so the loop never blocks on model_dir I/O — important for slow OSS-backed dirs.
• When export_config.exporter_type == "best", the current best-by-metric checkpoint is always retained even if it falls outside the recent-N window, so best-exporter export never breaks (effective retained count can be N+1).

Design

Checkpoint IO is consolidated into a new CheckpointManager (in checkpoint_util.py) rooted at model_dir, used for both save and load across train/eval/export/predict:

• save() writes the checkpoint then requests an async prune; close() (called at the end of training) drains pending deletions so the on-disk state is settled before export reads model_dir.
• Only rank 0 deletes. The queue is bounded by a _prune_pending coalescing flag (at most one in-flight + one pending pass). Deletions are filesystem-agnostic (local shutil.rmtree / fsspec fs.rm(recursive=True)).
• glob of model_dir is the source of truth every pass (no in-memory registry), so it stays correct under --continue_train.
• Load/discovery methods (latest_checkpoint, best_checkpoint, restore, restore_dataloader_state) delegate to the existing free functions, which remain canonical. fine_tune_checkpoint discovery stays a free-function call since it resolves a dir outside model_dir.

A partially-written trailing line in the eval-result file (concurrently appended by eval) can make a prune pass raise; this is caught and logged, and the next pass self-heals — no crash, no spurious deletions.

Testing

checkpoint_util_test.py: parameterized prune cases (keep-all / recent-N / protect-best / keep≥count), non-rank-0 no-op, coalesced-idempotent prune, and discovery-delegation. All green; pre-commit run -a clean on changed files.

[github-actions]

close() only runs if the training loop completes normally. If any step between the first save() (which starts the daemon worker) and here raises, close() is skipped: the final coalesced prune pass is abandoned and the worker thread is leaked (harmless across process exit since it's a daemon, but it breaks the documented "on-disk state settled before export reads model_dir" contract if anything downstream runs in the same process). Consider wrapping the loop body so close() runs in a finally.

[github-actions]

All FS work here (glob, the eval-file read inside best_checkpoint, and fs.rm) runs on the prune worker thread while the main thread is concurrently doing the next save() — and both go through the same process-global cached fsspec filesystem (filesystem_util._CACHED_FSSPEC_FILESYSTEMS[protocol]). Concurrent calls into one fsspec backend instance from two threads aren't guaranteed safe for all backends (e.g. ossfs runs async ops on a single per-instance loop). The target paths are disjoint (worker only touches older, already-protected dirs), so this is about the shared client object, not the paths. Worth confirming the OSS backend is thread-safe under concurrent access, or giving the worker its own filesystem handle.

[github-actions]

This comment reads as a hard cap, but when export_config.exporter_type == "best" the best checkpoint is retained in addition to the N recent ones, so the effective count can be N+1. The class docstring and train.md both note this; the proto comment is the only place that omits it.

Suggested change

	optional uint32 keep_checkpoint_max = 20 [default = 0];
	// max number of recent checkpoints to keep; 0 keeps all. When
	// export_config.exporter_type is "best", the best checkpoint is also retained.
	optional uint32 keep_checkpoint_max = 20 [default = 0];

[github-actions]

Review summary

Solid, well-designed change. The async pruning design is the right one: save() does the synchronous collective write and only enqueues a coalesced prune request, while a single rank-0 daemon worker does all the slow model_dir I/O (glob, eval-file read, recursive delete) off the training thread. The queue is provably bounded by the _prune_pending flag, glob is kept as the source of truth (no stale registry under --continue_train), and close() uses a FIFO stop-sentinel so disk state settles before export reads model_dir. The rank-0-only prune introduces no collective mismatch — prune()/close() contain no collective ops and run after the per-step collectives complete on every rank. The self-healing-on-corrupt-eval-line behavior is genuinely fail-safe: best_checkpoint() is called before the deletion loop, so any parse error aborts the whole pass with zero deletions and the next pass recovers. Docstrings are accurate and the test parameterization covers the retention policy well.

A few items worth considering (inline comments on the specifics):

• close() not exception-safe (main.py:530) — skipped if training raises, leaking the worker / abandoning the final prune.
• Shared fsspec filesystem across threads (checkpoint_util.py:407) — worth confirming the OSS backend is safe under concurrent worker+main access.
• Proto comment (train.proto:74) omits the N+1 (recent-N + best) case noted everywhere else.

Test gaps (non-blocking, but the highest-value follow-ups):

• save() itself is never exercised end-to-end — tests call prune() on hand-made dirs, so the save→prune wiring and the dataloader_state branch are unverified. A subprocess test mirroring the existing _save_restore_worker pattern would close this.
• The _remove_checkpoint fsspec branch and the corrupt-eval-line self-heal (a PR-advertised behavior) have no coverage; both are cheap to add with mocks.

Minor, pre-existing (while you're here): best_checkpoint() line 229 uses metric.values()[0], which raises TypeError in Py3 (dict_values isn't subscriptable). Not introduced by this PR, but _run_prune now relies on best_checkpoint, so a best-exporter run without an explicit best_exporter_metric would silently abort every prune pass. Suggest list(metric.values())[0].