* Studio: force-terminate a stuck training stop after a grace period
A Stop-with-save only signals the worker and waits for it to save and exit;
force_terminate() was reachable only from the /reset cancel path. On Windows +
ROCm the worker saves the adapter fine but then wedges in post-save GPU/HIP
teardown and never exits, so the run stays in "Stopping..." forever, is_training
stays true, and /reset returns 409.
Add a stop watchdog: when a stop is requested, a daemon escalates to
force_terminate() a short grace after the worker's "complete" (save done), or
after an absolute cap covering a hang during save. After escalation the parent
state is finalized (is_training=False, "Training stopped.") even if the OS never
reaps the wedged worker, so the UI leaves "Stopping..." and a new run can start.
No behavior change on a clean quick exit. Grace and timeout are configurable via
UNSLOTH_STUDIO_TRAINING_STOP_GRACE_S (15) and
UNSLOTH_STUDIO_TRAINING_STOP_TIMEOUT_S (120).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: harden the training stop watchdog per review
Address review feedback so a stop can never corrupt a checkpoint or leave the
run stuck:
- Never force-kill an in-progress save. The absolute cap is now a last-resort
backstop: raise the save default to 600s and only kill past that long window;
a not-yet-complete save is not treated as a hang. Cancels have nothing to save,
so they keep a shorter 120s cap via UNSLOTH_STUDIO_TRAINING_CANCEL_TIMEOUT_S.
The save vs cancel path is now explicit and the backstop logs a clear warning.
- Always finalize even if force_terminate raises on a wedged child (try/finally),
so the watchdog never dies leaving the run in "Stopping...".
- Preserve output_dir when the watchdog finalizes so a saved checkpoint is still
recorded in run history.
- Track the watched process per watchdog: a new run always gets its own watcher,
and a stale watchdog on an old proc no longer suppresses it.
- Terminate only the captured proc; force_terminate revalidates under the lock
that it is still the current worker, so it can never kill a fresh run.
- Name the watchdog thread for debuggability.
* studio: tighten training-stop watchdog comments
Comment-only pass: collapse the watchdog docstrings and inline notes to fewer
lines while keeping the rationale. No behavior change.
* Studio: make the stop watchdog safe against concurrent runs and the pump
Target-scope the escalation finalize so a stale watchdog can never clobber a
run that replaced its worker: capture the watched proc and job id, and no-op
the finalize (handle, progress, and DB) when a new run has already taken over.
Honor a later cancel by tightening an in-flight save watchdog to the shorter
cancel cap. Serialize the DB helpers on the lock so the watchdog and pump can
no longer double-create, double-finalize, or corrupt the metric buffer when a
force-terminate hands off to a still-finalizing pump.
Add regression tests: finalize no-ops when superseded, finalize runs for its
own worker, a later cancel tightens the cap, finalize is single-winner under
concurrency, finalize honors expected_job_id, and concurrent flushes claim
each metric exactly once.
* Studio: close the remaining stop-watchdog vs start/pump races
Guard the escalation finalize by the watched job id in addition to the proc:
start_training sets current_job_id before it installs the new _proc, so a stale
watchdog entering during that startup window still sees the old dead handle and
was not caught by the proc-only guard. Capture the job id when the watchdog
starts and require it to still match before touching state.
Snapshot the run id and final progress under the finalize lock and thread them
through the flush and finish_run calls, so a new run that starts between the
finalize claim and the DB writes cannot be flushed or marked stopped under the
old run's finalizer.
Publish _db_run_created only after create_run commits, gated by a dedicated
in-progress flag, so a concurrent finalize can no longer run finish_run against
a not-yet-inserted row and leave the run stuck as running.
Add regression tests for the startup-window job-id guard, run-id pinned flush,
snapshot-based finalize across a new run, and create-not-published-before-insert.
* Studio: finalize a force-stopped run by its captured id
If a new run starts in the gap after the watchdog clears _proc and marks the
backend idle, current_job_id changes, so the previous expected_job_id guard made
the finalize skip and left the stopped run recorded as running. Capture the run
id, metrics, and final progress under the lock (where current_job_id is still the
watched run) and finalize by that captured id via _finish_stopped_run: finish_run
is an idempotent UPDATE and insert_metrics_batch upserts, so a concurrent pump
finalize of the same run is harmless and a newly started run is never touched.
Add a test that the watched run is finalized by id with its buffered metrics, and
update the escalation tests to assert finalize goes through _finish_stopped_run.
* Studio: keep force-stop finalization retryable and unclaimed until the row exists
Only claim _run_finalized in the escalation when the DB row already exists; if an
early create failed and the pump is retrying it, claiming would make the pump's
later finalize no-op and strand the row as running, so leave the finalize to that
create-then-finalize path.
On a DB error in _finish_stopped_run (e.g. a transient SQLite lock), unclaim the
finalize and requeue the drained metrics when the run is still current, so the
pump or a later retry can still record the run stopped instead of leaving history
with an active run and lost metrics. A superseded run's state is never touched.
Add tests: no claim before the row exists, requeue+unclaim on a DB error, and a
superseded run left untouched on error.
* Studio: tighten stop-watchdog comments
Reduce the wording of the docstrings and inline comments added by this PR without
dropping any of the concurrency invariants (dual proc/job-id supersession guard,
finalize-by-captured-id, publish-after-commit, snapshot-under-lock, unclaim and
requeue on error). Comments and docstrings only; no code change.
* Studio: record the stopped run's DB state before dropping _proc
A wedged worker still reports alive, so the pump never reaches its own finalize
and bails on its _proc-is-None guard once the escalation drops the handle. So the
watchdog is the sole finalizer: record the terminal DB state (create the row if a
start-time create failed, then finish by captured id) BEFORE dropping _proc. While
the handle is held is_training_active() stays true, so no new run can start and
current_job_id stays the watched run for the write; _proc is dropped last, guarded
on target_proc so a run that did replace the worker keeps its handle.
_finish_stopped_run retries a transient DB error a few times (the pump can no
longer retry once _proc is gone) and unclaims on final failure only when the run
is still current. Add tests for create-then-finalize, retry-then-unclaim, and not
dropping a new run's handle.
* Studio: job-guard the DB create flags against a racing new run
_ensure_db_run_created publishes backend-wide _db_run_created and
_db_create_in_progress flags. When the watchdog creates a missing row for an
escalated stop, the killed worker lets a new /start proceed mid-create, so the
stale create could publish those flags against the new current_job_id, making the
new run skip inserting its own row (metric/finalize then target a missing run).
Publish the flags only when the captured job id is still current; the row is still
created by id, and the new run owns/creates its own. Also reset
_db_create_in_progress in start_training so a stale claim can't block a new run.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
4beb0a3a5f