unsloth/studio/backend/core/training
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
..
__init__.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
resume.py feat(studio): implement S3 dataset loading (completes #5951) (#6222) 2026-06-12 14:52:04 +02:00
s3_dataset.py feat(studio): implement S3 dataset loading (completes #5951) (#6222) 2026-06-12 14:52:04 +02:00
trainer.py Studio: fix misleading "increase max_seq_length" message for train-on-completions (#6664) 2026-06-25 03:30:12 -07:00
training.py Studio: keep the training event pump alive so progress can't silently freeze (#6643) 2026-06-25 05:19:32 -07:00
worker.py Generalize transformers tier selection by probing AutoConfig (#6550) 2026-06-22 08:20:06 -07:00