vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: keep the training event pump alive so progress can't silently freeze (#6643)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-06-25 05:19:32 -07:00 • committed by GitHub
parent 4929c5f769
commit 1cb04be328
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 626 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

173
studio/backend/core/training/training.py
View file

@ -216,6 +216,9 @@ class TrainingBackend:
self._event_queue: Any = None
self._stop_queue: Any = None
self._pump_thread: Optional[threading.Thread] = None
# True while a pump thread should be running; cleared on intended exits.
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
self._lock = threading.Lock()
# Progress state (updated by pump thread from subprocess events)
@ -289,6 +292,9 @@ class TrainingBackend:
logger.warning("Previous pump thread did not exit within 5s — refusing to start")
return False
self._pump_thread = None
# Clear a stale crash flag from a prior died pump so the watchdog can't
# treat this fresh setup as a recoverable death.
self._pump_running = False
# Build config dict for the subprocess
config = {
@ -472,16 +478,21 @@ class TrainingBackend:
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Assign subprocess handles after state reset.
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
# Eagerly create DB run row so it appears in history during model loading.
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
return True
@ -606,6 +617,9 @@ class TrainingBackend:
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
@ -623,10 +637,44 @@ class TrainingBackend:
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
new_pump.start()
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
def _ensure_pump_alive(self) -> bool:
"""Restart the event pump if it crashed, even after the worker exited.
Defence in depth behind _pump_loop's guards. _pump_running stays True only
after an abnormal exit (the loop clears it on intended exits), so a True
flag plus a dead thread is an unambiguous crash. Restarts even after worker
exit so a fresh pump can drain the terminal events and finalize; otherwise
the run looks stuck "running" forever. Returns True if restarted.
"""
with self._lock:
if not self._pump_running:
return False
# A restarted pump needs the worker handle and queue to drain/finalize;
# their absence means nothing is left to recover.
if self._proc is None or self._event_queue is None:
return False
if self._pump_thread is not None and self._pump_thread.is_alive():
return False
logger.error(
"Training event pump thread died while the worker is still running; "
"restarting it so progress updates resume."
)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread = new_pump
# Start under the lock so a concurrent _ensure_pump_alive can't see
# this thread as not-yet-started and spawn yet another pump.
new_pump.start()
return True
def is_training_active(self) -> bool:
"""Check if training is currently active."""
# Self-heal a crashed pump first: a dead pump must never leave the worker
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
with self._lock:
if self._proc is not None and self._proc.is_alive():
return True
@ -727,51 +775,87 @@ class TrainingBackend:
# Event pump (background thread)
# ------------------------------------------------------------------
def _safe_handle_event(self, event: dict) -> None:
"""Apply one event, swallowing any handler error.
The pump is the only writer of the progress state every status surface
reads, so a malformed event must never propagate and kill it.
"""
try:
self._handle_event(event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Training event pump: failed to handle %s event; skipping", etype)
def _pump_loop(self) -> None:
"""Background thread: consume events from subprocess → update state."""
"""Background thread: consume subprocess events and update state.
Sole writer of the in-memory progress state that /progress, /status,
/metrics and DB history read. If it exited while the worker still ran, the
run would burn GPU with events piling up while every surface froze. So no
single bad event or transient queue/DB error may end it; it returns only
through intended exits (worker gone, respawn handed off, finalized).
"""
self._pump_running = True
while True:
if self._proc is None or self._event_queue is None:
self._pump_running = False
return
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
try:
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
except Exception:
# If a read keeps raising after the worker died, fall through to
# finalize instead of spinning; only retry while the worker lives.
logger.exception("Training event pump: queue read failed; continuing")
if self._proc is not None and self._proc.is_alive():
time.sleep(0.1)
continue
event = None
if event is not None:
self._handle_event(event)
self._safe_handle_event(event)
continue
if self._proc.is_alive():
continue
# Process exited — drain remaining events.
for e in self._drain_queue(self._event_queue):
self._handle_event(e)
# Worker exited. Drain the backlog and finalize, guarded so a slow or
# failing DB write can't strand the thread; we return either way.
try:
for e in self._drain_queue(self._event_queue):
self._safe_handle_event(e)
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Runs on THIS exiting pump thread and starts a fresh pump (never joins
# the current thread); DB run-state is preserved.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Starts a fresh pump on this thread (no self-join); it takes over
# _pump_running, so this exit leaves the flag set.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
self._progress.is_training = False
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error or "Training process exited unexpectedly"
)
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
self._progress.is_training = False
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error or "Training process exited unexpectedly"
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
self._pump_running = False
return
def _handle_event(self, event: dict) -> None:
@ -1094,6 +1178,8 @@ class TrainingBackend:
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
# A closed/broken queue reads as "no event"; any other error is left to
# _pump_loop's guarded block, which logs and backs off.
return None
@staticmethod
@ -1104,7 +1190,12 @@ class TrainingBackend:
events.append(q.get_nowait())
except queue.Empty:
return events
except (EOFError, OSError, ValueError):
except Exception:
# A drain error must not abort finalization: return what we have so
# the run finalizes rather than wedging "active" behind a dead worker.
logger.exception(
"Training event pump: queue drain failed; finalizing with drained events"
)
return events
# ------------------------------------------------------------------

494
studio/backend/tests/test_training_pump_resilience.py Normal file
View file

@ -0,0 +1,494 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Parent-side training event-pump resilience.
The pump is the only writer of the progress state /progress, /status, /metrics
and DB history read. If it died while the worker ran, the run would continue while
the UI froze -- the "training runs but no progress shows" symptom. These tests pin
two guards: a bad event/queue error can't kill the pump, and a dead pump is
detected and restarted (even after worker exit) so terminal events still finalize.
Fakes only; no GPU, network, or subprocess.
"""
from __future__ import annotations
import contextlib
import logging
import queue
import sys
import threading
import time
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub the heavy module-level imports of core/training/training.py so it imports
# under CPU-only/no-network, then restore them (see the restore loop below).
_SAVED: dict = {}
def _stub(name, mod):
_SAVED[name] = sys.modules.get(name)
sys.modules[name] = mod
_lg = _types.ModuleType("loggers")
_lg.get_logger = lambda name: logging.getLogger(name)
_stub("loggers", _lg)
_stub("structlog", _types.ModuleType("structlog"))
_mpl = _types.ModuleType("matplotlib")
_plt = _types.ModuleType("matplotlib.pyplot")
_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation
_mpl.pyplot = _plt
_stub("matplotlib", _mpl)
_stub("matplotlib.pyplot", _plt)
_hw = _types.ModuleType("utils.hardware")
_hw.prepare_gpu_selection = lambda *a, **k: (None, None)
_stub("utils.hardware", _hw)
_npl = _types.ModuleType("utils.native_path_leases")
_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext()
_npl.run_without_native_path_secret = lambda fn: fn
_stub("utils.native_path_leases", _npl)
_pth = _types.ModuleType("utils.paths")
_pth.outputs_root = lambda *a, **k: "/tmp/outputs"
_stub("utils.paths", _pth)
# Whether core.training.training was already imported before this file ran; only
# evict it below if we were the one to create the (stub-bound) module instance.
_TRAINING_PRE_IMPORTED = "core.training.training" in sys.modules
from core.training.training import TrainingBackend
# Restore every stubbed module so this file never pollutes the shared session.
for _name in (
"loggers",
"structlog",
"matplotlib",
"matplotlib.pyplot",
"utils.hardware",
"utils.native_path_leases",
"utils.paths",
):
_prev = _SAVED.get(_name)
if _prev is None:
sys.modules.pop(_name, None)
else:
sys.modules[_name] = _prev
# training imported its helpers while the stubs were active, binding them to stubs.
# If we created the cached module, evict it (and its parent) so a later test
# re-imports the real one.
if not _TRAINING_PRE_IMPORTED:
sys.modules.pop("core.training.training", None)
sys.modules.pop("core.training", None)
class _FakeProc:
"""A subprocess handle whose liveness the test drives directly."""
def __init__(self, alive: bool = True):
self._alive = alive
self.pid = 4321
def is_alive(self):
return self._alive
def join(self, timeout = None):
self._alive = False
class _IdleQueue:
"""get()/get_nowait() always signal "no event" so the pump idles."""
def put(self, *a, **k):
pass
def get(self, *a, **k):
raise queue.Empty
def get_nowait(self, *a, **k):
raise queue.Empty
class _ScriptedQueue:
"""Yields queued events once, then signals empty forever."""
def __init__(self, events):
self._events = list(events)
def put(self, *a, **k):
pass
def get(self, *a, **k):
if self._events:
return self._events.pop(0)
raise queue.Empty
def get_nowait(self, *a, **k):
if self._events:
return self._events.pop(0)
raise queue.Empty
def _dead_thread() -> threading.Thread:
t = threading.Thread(target = lambda: None)
t.start()
t.join()
return t
def _silence_db(monkeypatch, b):
"""Neutralize DB finalization so a started pump exits cleanly off-box."""
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_finalize_run_in_db", lambda **k: None)
def _wait_until(predicate, timeout = 5.0):
deadline = time.time() + timeout
while time.time() < deadline:
if predicate():
return True
time.sleep(0.01)
return predicate()
# ----------------------------------------------------------------------------
# Guarantee 1: a single bad event/queue error cannot kill the pump.
# ----------------------------------------------------------------------------
def test_pump_survives_handler_exception_and_keeps_processing(monkeypatch):
b = TrainingBackend()
_silence_db(monkeypatch, b)
handled: list = []
def fake_handle(ev):
if ev.get("type") == "boom":
raise RuntimeError("handler blew up")
handled.append(ev.get("type"))
monkeypatch.setattr(b, "_handle_event", fake_handle)
proc = _FakeProc(alive = True)
b._proc = proc
b._event_queue = _ScriptedQueue(
[{"type": "boom"}, {"type": "progress"}, {"type": "boom"}, {"type": "progress"}]
)
pump = threading.Thread(target = b._pump_loop, daemon = True)
pump.start()
try:
assert _wait_until(
lambda: handled.count("progress") == 2
), "pump must keep processing good events after handler exceptions"
assert pump.is_alive(), "pump thread must survive handler exceptions"
assert b._pump_running is True
finally:
proc._alive = False # let the loop reach its clean exit
pump.join(timeout = 5)
assert not pump.is_alive()
assert b._pump_running is False, "clean exit must clear the running flag"
def test_read_queue_narrow_contract():
class _Q:
def __init__(self, exc):
self.exc = exc
def get(self, *a, **k):
raise self.exc
# Expected closed/broken-queue signals read as "no event".
for exc in (queue.Empty(), EOFError(), OSError(), ValueError()):
assert TrainingBackend._read_queue(_Q(exc), 0.01) is None
# Anything unexpected propagates on purpose to _pump_loop's guarded block,
# which logs and backs off instead of swallowing it into a hot loop.
with pytest.raises(RuntimeError):
TrainingBackend._read_queue(_Q(RuntimeError("boom")), 0.01)
def test_pump_survives_queue_read_exception_and_recovers(monkeypatch):
# _read_queue raising an unexpected error must be caught by the pump's outer
# guard (log + backoff), not kill the pump; once reads recover it processes.
b = TrainingBackend()
_silence_db(monkeypatch, b)
handled: list = []
monkeypatch.setattr(b, "_handle_event", lambda ev: handled.append(ev.get("type")))
class _FlakyQueue:
def __init__(self):
self.calls = 0
def get(self, *a, **k):
self.calls += 1
if self.calls <= 3:
raise RuntimeError("transient queue read error")
if self.calls == 4:
return {"type": "progress", "step": 1}
raise queue.Empty
def get_nowait(self, *a, **k):
raise queue.Empty
proc = _FakeProc(alive = True)
b._proc = proc
b._event_queue = _FlakyQueue()
pump = threading.Thread(target = b._pump_loop, daemon = True)
pump.start()
try:
assert _wait_until(
lambda: handled == ["progress"]
), "pump must recover after read errors and process the next event"
assert pump.is_alive()
finally:
proc._alive = False
pump.join(timeout = 5)
def test_pump_finalizes_when_drain_queue_raises_unexpected_error(monkeypatch):
# Worker has exited; the final drain hits an unexpected error. The run must
# still be finalized (not wedged "active" with a dead worker).
b = TrainingBackend()
finalized: dict = {}
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
class _BadDrainQueue:
def get(self, *a, **k):
raise queue.Empty
def get_nowait(self, *a, **k):
raise RuntimeError("corrupt drain payload")
b._proc = _FakeProc(alive = False)
b._event_queue = _BadDrainQueue()
b._progress.is_training = True
b._pump_loop() # returns once it sees the dead worker
assert b._progress.is_training is False
assert b._progress.error == "Training process exited unexpectedly"
assert finalized.get("status") == "error"
assert b._pump_running is False
assert b.is_training_active() is False
def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch):
# An unexpected error escapes _read_queue to the pump's outer guard; if it
# keeps raising after worker exit, the loop must still finalize, not spin.
b = TrainingBackend()
finalized: dict = {}
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
class _BrokenReadQueue:
def get(self, *a, **k):
raise RuntimeError("broken queue pipe")
def get_nowait(self, *a, **k):
raise queue.Empty
b._proc = _FakeProc(alive = False)
b._event_queue = _BrokenReadQueue()
b._progress.is_training = True
pump = threading.Thread(target = b._pump_loop, daemon = True)
pump.start()
pump.join(timeout = 5)
assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising"
assert b._progress.is_training is False
assert finalized.get("status") == "error"
assert b._pump_running is False
def test_start_training_clears_stale_pump_running_flag():
# A prior pump that died abnormally leaves _pump_running True. The next
# start_training must clear it during reset so the start-time watchdog can't
# treat the fresh setup as a recoverable crash and spawn a duplicate pump.
b = TrainingBackend()
b._pump_running = True
b._pump_thread = None
b._proc = None
# No model_name -> start_training bails at kwargs["model_name"] (KeyError),
# but only AFTER the reset block that clears the stale flag.
with pytest.raises(KeyError):
b.start_training("job_stale_flag_test")
assert b._pump_running is False
# ----------------------------------------------------------------------------
# Guarantee 2: a pump that dies while the worker runs is detected + restarted.
# ----------------------------------------------------------------------------
def test_ensure_pump_alive_restarts_crashed_pump(monkeypatch):
b = TrainingBackend()
_silence_db(monkeypatch, b)
b._proc = _FakeProc(alive = True)
b._event_queue = _IdleQueue()
b._pump_running = True # a pump started, then died abnormally
dead = _dead_thread()
b._pump_thread = dead
assert b._ensure_pump_alive() is True
try:
assert b._pump_thread is not dead
assert b._pump_thread.is_alive(), "a fresh pump must be running"
finally:
b._proc._alive = False
b._pump_thread.join(timeout = 5)
def test_ensure_pump_alive_noop_when_pump_alive():
b = TrainingBackend()
b._proc = _FakeProc(alive = True)
b._event_queue = _IdleQueue()
b._pump_running = True
release = threading.Event()
alive = threading.Thread(target = release.wait, daemon = True)
alive.start()
b._pump_thread = alive
try:
assert b._ensure_pump_alive() is False
assert b._pump_thread is alive
finally:
release.set()
alive.join(timeout = 5)
def test_ensure_pump_alive_revives_crashed_pump_after_worker_exit(monkeypatch):
# True _pump_running + dead thread = a crash (the loop clears the flag on
# intended exits). The queue may still hold terminal events, so the pump must
# restart to drain and finalize, else the run is stuck "running" forever.
b = TrainingBackend()
_silence_db(monkeypatch, b)
b._proc = _FakeProc(alive = False)
b._event_queue = _IdleQueue()
b._progress.is_training = True
b._pump_running = True
b._pump_thread = _dead_thread()
assert b._ensure_pump_alive() is True
assert _wait_until(
lambda: b._progress.is_training is False
), "the restarted pump must drain + finalize the stranded run"
b._pump_thread.join(timeout = 5)
assert b._pump_running is False
assert b.is_training_active() is False
def test_ensure_pump_alive_noop_during_setup():
# _pump_running is False between state-reset and the first pump actually
# running; the watchdog must not race in and spawn a rogue pump.
b = TrainingBackend()
b._proc = _FakeProc(alive = True)
b._event_queue = _IdleQueue()
b._pump_running = False
b._pump_thread = None
assert b._ensure_pump_alive() is False
assert b._pump_thread is None
def test_is_training_active_revives_dead_pump(monkeypatch):
b = TrainingBackend()
_silence_db(monkeypatch, b)
b._proc = _FakeProc(alive = True)
b._event_queue = _IdleQueue()
b._pump_running = True
dead = _dead_thread()
b._pump_thread = dead
# The status poll the SSE stream makes every second both reports activity
# and heals the dead pump as a side effect.
assert b.is_training_active() is True
try:
assert b._pump_thread is not dead
assert b._pump_thread.is_alive()
finally:
b._proc._alive = False
b._pump_thread.join(timeout = 5)
# ----------------------------------------------------------------------------
# Guarantee 3: the DB run row exists before the pump consumes any event.
# ----------------------------------------------------------------------------
def _stub_spawn(monkeypatch):
"""Stub start_training's spawn surface (GPU pick, mp context, worker)."""
g = TrainingBackend.start_training.__globals__
class _SpawnProc:
pid = 4321
def start(self):
pass
def is_alive(self):
return True
class _Ctx:
def Queue(self):
return _IdleQueue()
def Process(self, **k):
return _SpawnProc()
# _CTX / prepare_gpu_selection resolve from the module globals; patch the
# function's own globals so the eviction of core.training.training (done at
# this test module's import for isolation) can't hand us a different copy.
monkeypatch.setitem(g, "_CTX", _Ctx())
monkeypatch.setitem(g, "prepare_gpu_selection", lambda *a, **k: (None, None))
hw = _types.ModuleType("utils.hardware")
hw.prepare_gpu_selection = lambda *a, **k: (None, None)
hw.hardware = type("HW", (), {"DEVICE": "cuda", "DeviceType": type("D", (), {"MLX": "mlx"})})()
monkeypatch.setitem(sys.modules, "utils.hardware", hw)
pl = _types.ModuleType("utils.process_lifetime")
pl.adopt_pid = lambda pid: None
monkeypatch.setitem(sys.modules, "utils.process_lifetime", pl)
worker = _types.ModuleType("core.training.worker")
worker.run_training_process = lambda **k: None
monkeypatch.setitem(sys.modules, "core.training.worker", worker)
def test_db_run_created_before_pump_consumes_events(monkeypatch):
# A fast terminal worker must not race the pump into creating the DB row: by
# the time the pump runs, start_training has already created it. The create
# sleep widens the window so the ordering is observed, not luck.
b = TrainingBackend()
_stub_spawn(monkeypatch)
def slow_create():
time.sleep(0.05)
b._db_run_created = True
seen = {}
def fake_pump():
seen["db_created"] = b._db_run_created
b._pump_running = False
monkeypatch.setattr(b, "_ensure_db_run_created", slow_create)
monkeypatch.setattr(b, "_pump_loop", fake_pump)
assert b.start_training("job_db_order", model_name = "m") is True
if b._pump_thread is not None:
b._pump_thread.join(timeout = 2.0)
# The pump observed an already-created run; it would be False if the pump
# were started before the eager create.
assert seen["db_created"] is True