Fix resume training crash recovery and MLX checkpoints - #6796
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables resuming training runs that ended in an error state, persists the output directory to the database, and adds support for writing MLX stop checkpoints when training is stopped. It also updates the frontend to allow resuming training runs directly from the historical view. The feedback suggests defensively creating the checkpoint directory before saving stop checkpoints, disabling the resume button in the UI when another training run is already active, and including the 'running' status in database and frontend checks to prevent older runs from being marked as resumable when a newer run is active.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 839365f836
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
1c4615e to
2765a85
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 99ac1c2717
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
… 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.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 095482aafc
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…nalize
- 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.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e3f58281d5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 52c80757e9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e8d9603328
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 427378cdb2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 927cf0f598
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fe06efccaa
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5b5f4e4933
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0fa0d0b692
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
|
Codex Review: Didn't find any major issues. Hooray! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3ef3b4648f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self.eval_step_history.clear() | ||
| self.eval_enabled = False | ||
| self._output_dir = None | ||
| self._output_dir = config.get("output_dir") if resume_source_run_id else None |
There was a problem hiding this comment.
When a user resumes from history/API, the route only sends resume_from_checkpoint (rewritten to the selected checkpoint-* path) and does not send output_dir, so with resume_source_run_id this initializes _output_dir to None. _ensure_db_run_created() then calls create_run(..., output_dir=None, resumed_from_run_id=...), whose source-claim query compares output_dir = ?; SQL NULL never matches the original run's saved output directory, so normal resume attempts terminate the worker and fail as “checkpoint is no longer available.” Derive the output directory from resume_from_checkpoint here or pass the validated resume output dir into start_training.
Useful? React with 👍 / 👎.
| UPDATE training_runs SET resume_blocked = 1 | ||
| WHERE id = ? AND status IN ('stopped', 'error') | ||
| AND output_dir = ? AND resume_blocked = 0 |
There was a problem hiding this comment.
When a user starts a resumed run and then chooses Stop without saving, the new run clears its own output_dir, but the original source run remains permanently resume_blocked=1 from this initial claim. Since the canceled continuation no longer has the same output_dir, the dynamic resumed_later check will not hide the source, yet can_resume_run() still rejects it via resume_blocked, so the unchanged original checkpoint can no longer be resumed; unblock the source when the continuation finalizes as a no-save cancel, or defer this permanent claim until the resumed run preserves a replacement output directory/checkpoint.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. The source is kept resume_blocked here by design so a partially consumed output_dir is never offered for resume twice; the on-disk checkpoint itself is untouched. Releasing it correctly needs a resumed_from_run_id link to know which source to unblock, so tracking that as a small follow-up rather than in this PR.
|
Focused CI passes on current PR #6796 head: https://github.com/Imagineer99/unsloth/actions/runs/29747712420 |
- 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
|
Pushed a small follow-up (
No behavior change for valid HF, sharded, legacy |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 296f1f7f66
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| UPDATE training_runs SET resume_blocked = 1 | ||
| WHERE id = ? AND status IN ('stopped', 'error') | ||
| AND output_dir = ? AND resume_blocked = 0 |
There was a problem hiding this comment.
When a continuation created from history is later deleted (for example, a resume attempt fails before useful progress and the user removes that row), this permanent update leaves the source run with resume_blocked=1. Since delete_run() only removes the continuation row and can_resume_run() rejects resume_blocked before checking the still-valid output directory, the original stopped/errored run loses its Resume button even though no newer row remains to supersede it; the block needs to be released when the claiming continuation is deleted or otherwise produces no replacement.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. The source is kept resume_blocked here by design so a partially consumed output_dir is never offered for resume twice; the on-disk checkpoint itself is untouched. Releasing it correctly needs a resumed_from_run_id link to know which source to unblock, so tracking that as a small follow-up rather than in this PR.
|
@codex review |
|
Codex Review: Didn't find any major issues. 🎉 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
- 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
|
@codex review |
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
4 participants
Fixes several resume-training issues in Studio: