diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 290b4090e..9d991c651 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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 # ------------------------------------------------------------------ diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py new file mode 100644 index 000000000..d75b205f3 --- /dev/null +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -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