unsloth/studio
Daniel Han 1cb04be328
Studio: keep the training event pump alive so progress can't silently freeze (#6643)
* Studio: keep the training event pump alive so progress can't silently freeze

The parent-side event pump is the only writer of the in-memory progress state
that SSE /progress, /status, /metrics and the DB history all read. It ran in a
single unsupervised daemon thread with no guard around event handling, so one
malformed event or a transient queue/DB error would terminate it permanently.
The worker subprocess keeps training regardless (mp.Queue puts never block on an
unbounded queue), so a run kept burning GPU for hours while every progress
surface froze on the last step the pump saw.

- Guard each pump iteration: a bad event or queue-read error is logged and
  skipped instead of ending the loop. _read_queue now reads any error as
  "no event", not just Empty/EOFError/OSError/ValueError.
- Add a _pump_running flag and an _ensure_pump_alive watchdog wired into
  is_training_active, so a pump that dies while the worker is alive is restarted
  on the next status poll and the UI catches up from the still-open queue.
- Start respawned and restarted pumps under the lock so the watchdog can never
  spawn a duplicate during the brief start window.

Adds tests/test_training_pump_resilience.py covering both guarantees.

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

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

* Studio training pump: address review (drain guard, start race, read backoff, respawn flag)

Follow-up to the event-pump resilience change, closing four edge cases a
review surfaced in the same pump/queue surface:

- _drain_queue now tolerates any error during the worker-exit drain and
  finalizes with whatever it drained, instead of skipping finalization and
  leaving the run wedged "active" with a dead worker.
- start_training clears a stale _pump_running flag during reset and assigns
  the subprocess handles plus starts the pump under the lock, so a concurrent
  status/SSE poll can't spawn a duplicate pump during setup.
- _read_queue goes back to the narrow EOFError/OSError/ValueError catch;
  truly unexpected errors are left to _pump_loop's guarded read, which logs
  and backs off so a persistently raising queue can't spin a hot loop.
- The xet respawn-failure path clears _pump_running so a later run can't
  inherit a stale flag.

Adds regression tests for all four.

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

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

* Studio: revive a crashed pump after worker exit + stop test module pollution

Two review follow-ups on the training event pump:

- _ensure_pump_alive refused to restart once the worker had exited
  (not self._proc.is_alive()), so a pump that crashed just before the worker
  finished never drained the terminal complete/error events still sitting in
  the queue. progress.is_training stayed True and is_training_active() returned
  True forever, leaving the run stuck "running" behind a dead pump. A True
  _pump_running flag with a dead thread is an unambiguous crash regardless of
  worker state, so restart there too: the fresh pump drains the backlog and
  finalizes. Updated the watchdog test to assert the revive-and-finalize.

- The resilience test imports core.training.training while heavy module-level
  deps are stubbed, then restores the stubs -- but the cached training module
  kept the stubs bound in its globals, so a later test in the same session
  could exercise the fakes (e.g. prepare_gpu_selection) instead of the real
  code. Evict the training module (and its package) after import when this file
  created it, so subsequent tests re-import it cleanly.

* Studio: finalize training run when queue reads keep failing on a dead worker

reviewer.py follow-up. _read_queue only swallows EOFError/OSError/ValueError;
an unexpected error escapes to the pump's outer guard, which logged, slept and
`continue`d. If those reads keep raising after the worker has already exited
(e.g. a broken queue pipe), the loop never reaches the dead-worker finalize
block, so the pump spins on with _pump_running True and progress.is_training
stuck True -- the run looks like it is still training forever. On a read failure
now fall through to finalize when the worker is gone, only backing off and
retrying while it is still alive. Mirrors the data-recipe pump fix; added a
regression test.

* Tighten training pump resilience comments and docstrings

Condense the verbose explanatory comments and docstrings on the training event
pump and its tests to shorter, clearer forms. Comment/whitespace only; verified
no code changed via AST diff. No behaviour change.

* Studio: create the training DB run before starting the event pump

start_training started the event pump before the eager _ensure_db_run_created()
call, so for a worker that completes or fails immediately the pump could race the
main thread into creating and finalizing the same run row (duplicate INSERT, or a
finalize skipped while _db_run_created was still false). Create the run first; the
pump then only ever finalizes. Adds a regression test asserting the pump observes
an already-created run.

* [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>
2026-06-25 05:19:32 -07:00
..
backend Studio: keep the training event pump alive so progress can't silently freeze (#6643) 2026-06-25 05:19:32 -07:00
frontend Studio: UNSLOTH_NPM_REGISTRY opt-in for corporate npm mirrors (#6491) (#6663) 2026-06-25 04:01:43 -07:00
src-tauri Fix Linux AppImage packaging (#6657) 2026-06-24 19:40:00 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
install_llama_prebuilt.py Verify DiffusionGemma visual-server binary against approved checksums (#6635) 2026-06-24 06:37:41 -07:00
install_node_prebuilt.py Pin isolated Node.js installer to committed sha256 digests (#6625) 2026-06-24 05:47:58 -07:00
install_python_stack.py Make _uv_safe_path space-safe on macOS/Linux (#6503) (#6534) 2026-06-24 04:02:24 -07:00
LICENSE.AGPL-3.0 Add AGPL-3.0 license to studio folder 2026-03-09 19:36:25 +00:00
node_prebuilt_pins.json Pin isolated Node.js installer to committed sha256 digests (#6625) 2026-06-24 05:47:58 -07:00
package-lock.json ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
package.json ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
setup.bat Final cleanup 2026-03-12 18:28:04 +00:00
setup.ps1 Studio: UNSLOTH_NPM_REGISTRY opt-in for corporate npm mirrors (#6491) (#6663) 2026-06-25 04:01:43 -07:00
setup.sh Studio: UNSLOTH_NPM_REGISTRY opt-in for corporate npm mirrors (#6491) (#6663) 2026-06-25 04:01:43 -07:00
Unsloth_Studio_Colab.ipynb Fix/studio colab proxy and iframe - Unsloth Studio not loading in Colab (iframe "refused to connect" and wrong URL) (#5844) 2026-05-28 23:54:48 -07:00