unsloth/studio/backend/storage
Nilay f3c085ad9e
Fix resume training crash recovery and MLX checkpoints (#6796)
* Fix resume training crash recovery and MLX checkpoints

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: preserve interrupted stop-and-save output_dir, verify MLX checkpoint

- finish_run: add clear_output_dir flag; preserve output_dir for stopped/error
  unless cancel explicitly clears it (fixes pump finalization wiping persisted path).
- training pump: pass interrupted stop-and-save context into finalize_run_in_db.
- MLX stop-and-save: verify resumable checkpoint exists before sending complete;
  return bool from _write_mlx_stop_checkpoint and add regression tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address Codex review: MLX current-step checkpoint and cancel error finalize

- Only skip MLX stop checkpoint write when checkpoint-{current_step} exists;
  stale periodic checkpoints no longer mask missing stop saves.
- Pass clear_output_dir through error-event finalization so Stop-without-save
  cannot leave a persisted output_dir that still offers Resume.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address  review

* Address more reviews

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* more reviews

* clear in-memory output_dir on interrupted cancel

* allow resuming errored runs at the final step

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Clear persisted output_dir in cancel watchdog path

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Write MLX stop checkpoint in stop path, keep output_dir on crash finalize

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): harden resumable run finalization

* fix(studio): defer safetensors checkpoint import

* fix(studio): reject stale training cancellation

* fix(studio): replay null resume targets

* fix(studio): serialize terminal cancellation

* Harden resume checkpoint validation and fix stop-save cleanup

- Reject unrecognized shard formats and keep indexed shard paths inside the checkpoint dir
- Require a non-empty tensor record when validating .pt/.bin optimizer and model state
- Always finalize TensorBoard and W&B on stop-save-failure exits
- Refuse writing an MLX stop checkpoint through a symlinked directory
- Clarify the resume rejection message to cover errored runs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten resume/checkpoint comments

* Recover resumability when a valid stop checkpoint landed

- Re-validate the current-step checkpoint in the dead-worker and error finalization paths so a stop-and-save that actually wrote a valid checkpoint is not wrongly marked error/resume_blocked
- Accept a valid tensor-free optimizer state (e.g. SGD without momentum); the model-state check still requires real tensors
- Include errored runs in the frontend resume rejection message

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lyxot <longyixing331@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
2026-07-21 02:34:58 -07:00
..
__init__.py feat(studio): training history persistence and past runs viewer (#4501) 2026-03-25 00:58:55 -07:00
mcp_servers_db.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
providers_db.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
rag_db.py Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
studio_db.py Fix resume training crash recovery and MLX checkpoints (#6796) 2026-07-21 02:34:58 -07:00