diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 2ddda19951..6b32ec873d 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -39,6 +39,27 @@ from utils.paths import outputs_root logger = get_logger(__name__) + +def _env_int(name: str, default: int) -> int: + try: + raw = (os.environ.get(name) or "").strip() + return int(raw) if raw else default + except ValueError: + return default + + +# Stop-watchdog escalation timeouts. Primary trigger: a short grace once "complete" +# (save done). Absolute cap is a backstop: long for save=True so a slow save is never +# killed mid-write, shorter for a cancel that has nothing to save. +_STOP_GRACE_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_GRACE_S", 15) +_STOP_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_TIMEOUT_S", 600) +_CANCEL_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_CANCEL_TIMEOUT_S", 120) + +# Watchdog DB finalize: a few short retries so a transient SQLite lock doesn't lose the +# terminal state, since the watchdog is the sole finalizer once _proc is dropped. +_DB_FINALIZE_RETRIES = 3 +_DB_FINALIZE_RETRY_S = 0.5 + _pyplot = None _pyplot_failed = False @@ -741,6 +762,13 @@ class TrainingBackend: self._pump_running: bool = False self._lock = threading.Lock() + # Stop watchdog: after a stop is requested, escalates to force_terminate() + # if the worker does not exit on its own within a bounded time. The watched + # proc is tracked so a new run always gets its own watcher. + self._stop_watchdog: Optional[threading.Thread] = None + self._stop_watchdog_proc: Optional[mp.Process] = None + self._complete_seen = threading.Event() + # Progress state (updated by pump thread from subprocess events) self._progress = TrainingProgress() self._should_stop = False @@ -765,6 +793,7 @@ class TrainingBackend: self._metric_buffer: list[dict] = [] self._run_finalized: bool = False self._db_run_created: bool = False + self._db_create_in_progress: bool = False self._db_total_steps_set: bool = False self._db_config: Optional[dict] = None self._db_started_at: Optional[str] = None @@ -896,6 +925,7 @@ class TrainingBackend: self.current_job_id = job_id self._should_stop = False self._cancel_requested = False + self._complete_seen.clear() self._progress = TrainingProgress( is_training = True, status_message = "Initializing training..." ) @@ -911,6 +941,7 @@ class TrainingBackend: self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False + self._db_create_in_progress = False # a stale watchdog create can't block this run self._db_total_steps_set = False self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() @@ -953,15 +984,212 @@ class TrainingBackend: self._progress.status_message = ( "Stopping training and saving checkpoint..." if save else "Cancelling training..." ) + # Guarantee the run finalizes even if the worker wedges after saving. + self._start_stop_watchdog(cancel = not save) return True - def force_terminate(self) -> None: - """Force-kill the training subprocess so state can be reset immediately.""" + def _start_stop_watchdog(self, cancel: bool) -> None: + """Start a daemon that force-terminates the worker if a requested stop does not + exit on its own. No-op if no worker is alive or a live watchdog already watches + this proc (a stale watchdog on an old proc never blocks a new run's watcher).""" with self._lock: - if self._proc is not None and self._proc.is_alive(): - logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid) - self._proc.terminate() proc = self._proc + if proc is None or not proc.is_alive(): + return + if ( + self._stop_watchdog is not None + and self._stop_watchdog.is_alive() + and self._stop_watchdog_proc is proc + ): + return + watchdog = threading.Thread( + target = self._stop_watchdog_loop, + args = (proc, cancel, self.current_job_id), + name = f"stop-watchdog-{self.current_job_id or 'unknown'}", + daemon = True, + ) + self._stop_watchdog = watchdog + self._stop_watchdog_proc = proc + watchdog.start() + + def _stop_watchdog_loop( + self, + target_proc: "mp.Process", + cancel: bool, + watched_job_id: Optional[str] = None, + ) -> None: + """Escalate a stuck stop to force_terminate(): grace after "complete", else the + absolute backstop (see the module timeouts). No-ops on a clean exit; exits + silently if a new run replaces the worker.""" + started = time.monotonic() + complete_at: Optional[float] = None + reason = "" + while True: + with self._lock: + superseded = self._proc is not target_proc + # A later cancel has nothing to save, so tighten an in-flight save + # watchdog to the shorter cancel cap. + cancelling = cancel or self._cancel_requested + if superseded or not target_proc.is_alive(): + return + now = time.monotonic() + abs_timeout = _CANCEL_TIMEOUT_S if cancelling else _STOP_TIMEOUT_S + if complete_at is None and self._complete_seen.is_set(): + complete_at = now + if complete_at is not None and now - complete_at >= _STOP_GRACE_S: + reason = "worker still alive after save" + break + if now - started >= abs_timeout: + reason = "worker did not exit within the absolute timeout" + break + time.sleep(0.5) + + with self._lock: + superseded = self._proc is not target_proc + if superseded or not target_proc.is_alive(): + return + if complete_at is None: + # Backstop fired pre-completion: a save may still be in progress. + logger.warning( + "Stop watchdog: absolute timeout with no completion signal; " + "force-terminating a possibly-mid-save worker: %s", + reason, + ) + else: + logger.warning("Stop watchdog force-terminating stuck training worker: %s", reason) + # force_terminate can raise on a wedged child; finalize regardless. + try: + self.force_terminate(target_proc = target_proc) + except Exception: + logger.exception("Stop watchdog: force_terminate failed; finalizing anyway") + finally: + self._finalize_stopped_after_escalation( + target_proc = target_proc, watched_job_id = watched_job_id + ) + + def _finalize_stopped_after_escalation( + self, + target_proc: "Optional[mp.Process]" = None, + watched_job_id: Optional[str] = None, + ) -> None: + """Finalize parent state after a force-terminate so the UI leaves "Stopping..." + even if the worker is wedged in driver teardown; preserves output_dir so a saved + checkpoint is kept. No-ops if a new run already replaced the watched worker, so a + stale watchdog never marks a fresh run stopped or drops its handle. + + Supersession is checked on both the watched proc and job id: start_training sets + current_job_id before it installs the new _proc, so a stale watchdog entering that + startup window still sees the old (dead) handle and is caught by the job-id guard. + + The run's terminal DB state is recorded (create-if-needed + finish by captured id) + BEFORE _proc is dropped: a wedged worker still reports alive, so the pump never + reaches its own finalize and would bail on its _proc-is-None guard once the handle + is gone. 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, + re-guarded on target_proc so a run that did replace the worker keeps its handle.""" + with self._lock: + if target_proc is not None and self._proc is not target_proc: + return # a new run replaced the worker; never touch its state + if watched_job_id is not None and self.current_job_id != watched_job_id: + return # a new run is already starting up; leave its state alone + run_id = self.current_job_id # == watched_job_id + self._progress.is_training = False + self._progress.status_message = "Training stopped." + # Create the row if a start-time create failed (no-op otherwise; skips when the pump + # is mid-create, in which case its create-then-finalize records the run instead). + self._ensure_db_run_created() + with self._lock: + claim = ( + bool(run_id) + and self.current_job_id == run_id + and self._db_run_created + and not self._run_finalized + ) + batch: list = [] + final_step = final_loss = duration = None + loss_history: list = [] + output_dir = self._output_dir + if claim: + self._run_finalized = True # claim this run's finalize + batch = list(self._metric_buffer) + del self._metric_buffer[: len(batch)] + final_step = self._progress.step + final_loss = self._progress.loss + if final_loss is not None and not math.isfinite(final_loss): + final_loss = None + duration = self._progress.elapsed_seconds + loss_history = list(self.loss_history) + if claim: + self._finish_stopped_run( + run_id, output_dir, batch, final_step, final_loss, duration, loss_history + ) + with self._lock: + if target_proc is None or self._proc is target_proc: + self._proc = None # drop only our handle, never a run that replaced it + + def _finish_stopped_run( + self, + run_id: str, + output_dir: Optional[str], + batch: list, + final_step: Optional[int], + final_loss: Optional[float], + duration: Optional[float], + loss_history: list, + ) -> None: + """Record a force-stopped run finished by its captured id, from state snapshotted + under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE, + so a concurrent pump finalize of the same run is harmless and a different current run + is never touched. The watchdog is the sole finalizer once _proc is dropped, so a + transient DB error (e.g. a SQLite lock) is retried a few times; on final failure the + finalize is unclaimed (only if the run is still current) so the row is not left + claimed-but-unfinalized.""" + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import finish_run, insert_metrics_batch + from utils.downsample import downsample + + if batch: + insert_metrics_batch(run_id, batch) + sparkline = downsample(loss_history, 50) + finish_run( + id = run_id, + status = "stopped", + ended_at = datetime.now(timezone.utc).isoformat(), + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, + loss_sparkline = _json.dumps(sparkline), + output_dir = output_dir, + error_message = None, + ) + return + except Exception: + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + continue + logger.warning( + "Failed to finalize stopped run %s in DB after %d attempts", + run_id, + _DB_FINALIZE_RETRIES, + exc_info = True, + ) + with self._lock: + # Only if still current; a new run's finalize state is never touched. + if self.current_job_id == run_id: + self._run_finalized = False + + def force_terminate(self, target_proc: "Optional[mp.Process]" = None) -> None: + """Force-kill the training subprocess so state can be reset immediately. With + ``target_proc``, terminate only that handle and no-op if a new run has replaced + it, so the watchdog can never kill a fresh worker.""" + with self._lock: + proc = self._proc + if target_proc is not None and proc is not target_proc: + return # superseded by a new run; do not touch the new worker + if proc is not None and proc.is_alive(): + logger.info("Force-terminating training subprocess (pid=%s)", proc.pid) + proc.terminate() cancelled = self._cancel_requested output_dir = self._output_dir @@ -1468,6 +1696,8 @@ class TrainingBackend: "training cancelled", "training stopped", } + # Save is done by now; let the stop watchdog start its grace timer. + self._complete_seen.set() self._progress.is_training = False self._progress.is_completed = not stopped self._output_dir = event.get("output_dir") @@ -1532,90 +1762,135 @@ class TrainingBackend: self._finalize_run_in_db(**db_action_kwargs) def _ensure_db_run_created(self) -> None: - """Create the DB row if it doesn't exist yet. Called outside the lock.""" - if self._db_run_created or not self.current_job_id or not self._db_config: - return + """Create the DB row if it doesn't exist yet. An in-progress flag lets only one + caller create at a time, and ``_db_run_created`` is published only after + ``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a + not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running).""" + with self._lock: + if ( + self._db_run_created + or self._db_create_in_progress + or not self.current_job_id + or not self._db_config + ): + return + self._db_create_in_progress = True # only one caller creates + job_id = self.current_job_id + db_config = self._db_config + started_at = self._db_started_at or datetime.now(timezone.utc).isoformat() + total_steps = self._progress.total_steps or None + created = False try: from storage.studio_db import create_run dataset_name = ( - self._db_config.get("hf_dataset") - or next(iter(self._db_config.get("local_datasets") or []), None) - or _s3_dataset_name(self._db_config.get("s3_dataset")) + db_config.get("hf_dataset") + or next(iter(db_config.get("local_datasets") or []), None) + or _s3_dataset_name(db_config.get("s3_dataset")) or "unknown" ) create_run( - id = self.current_job_id, - model_name = self._db_config["model_name"], + id = job_id, + model_name = db_config["model_name"], dataset_name = dataset_name, - config_json = _json.dumps(self._db_config), - started_at = self._db_started_at or datetime.now(timezone.utc).isoformat(), - total_steps = self._progress.total_steps or None, + config_json = _json.dumps(db_config), + started_at = started_at, + total_steps = total_steps, ) - self._db_run_created = True + created = True except Exception: logger.warning("Failed to create DB run record for early failure", exc_info = True) + finally: + with self._lock: + # Publish the flags only if this is still the current run. A killed worker + # lets a new /start proceed mid-create, and these flags are backend-wide, so + # a stale create for the captured job must not satisfy the new run's DB state + # (the row was still created by id; the new run owns/creates its own row). + if self.current_job_id == job_id: + if created: + self._db_run_created = True # publish only after the insert commits + self._db_create_in_progress = False def _finalize_run_in_db( self, status: str, error_message: Optional[str] = None, output_dir: Optional[str] = None, + expected_job_id: Optional[str] = None, ) -> None: - """Flush remaining metrics and mark a run as finished in the DB.""" - if not self.current_job_id or not self._db_run_created or self._run_finalized: - return - self._flush_metrics_to_db() + """Flush remaining metrics and mark a run finished in the DB. Claims the finalize + under the lock so the watchdog and pump can't double-finalize, and no-ops when + ``expected_job_id`` no longer matches (a new run took over). The run id and final + progress are snapshotted under the lock and threaded through the flush/finish calls, + so a new run racing between this claim and the DB writes can't be flushed or marked + stopped under the old run's finalize.""" + with self._lock: + if expected_job_id is not None and self.current_job_id != expected_job_id: + return + if not self.current_job_id or not self._db_run_created or self._run_finalized: + return + self._run_finalized = True + run_id = self.current_job_id + final_step = self._progress.step + final_loss = self._progress.loss + if final_loss is not None and not math.isfinite(final_loss): + final_loss = None + duration = self._progress.elapsed_seconds + loss_history = list(self.loss_history) + self._flush_metrics_to_db(run_id = run_id) try: from storage.studio_db import finish_run from utils.downsample import downsample - sparkline = downsample(self.loss_history, 50) + sparkline = downsample(loss_history, 50) finish_run( - id = self.current_job_id, + id = run_id, status = status, ended_at = datetime.now(timezone.utc).isoformat(), - final_step = self._progress.step, - final_loss = self._progress.loss - if (self._progress.loss is not None and math.isfinite(self._progress.loss)) - else None, - duration_seconds = self._progress.elapsed_seconds, + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, loss_sparkline = _json.dumps(sparkline), output_dir = output_dir, error_message = error_message, ) - self._run_finalized = True except Exception: + with self._lock: + self._run_finalized = False # unclaim so a later flush can retry logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) - def _flush_metrics_to_db(self) -> None: - """Flush buffered metrics to the database and update live progress.""" - if not self._metric_buffer or not self.current_job_id or not self._db_run_created: - return - # Cap buffer to bound memory growth. - if len(self._metric_buffer) > 500: - logger.warning( - "Metric buffer exceeded 500 entries (%d) — trimming oldest", - len(self._metric_buffer), - ) - self._metric_buffer = self._metric_buffer[-500:] - # Snapshot before insert so metrics arriving during the write survive. - batch = list(self._metric_buffer) + def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None: + """Flush buffered metrics to the DB and update live progress. The target run id, + metric batch, and progress snapshot are all taken under the lock, so a concurrent + flush can't double-remove metrics and a racing new run can't redirect the write to + a different job. A finalizer passes ``run_id`` to pin the target to its captured run.""" + with self._lock: + target = run_id if run_id is not None else self.current_job_id + if not self._metric_buffer or not target or not self._db_run_created: + return + # Cap buffer to bound memory growth. + if len(self._metric_buffer) > 500: + logger.warning( + "Metric buffer exceeded 500 entries (%d) — trimming oldest", + len(self._metric_buffer), + ) + del self._metric_buffer[:-500] + # Claim the batch under the lock so a concurrent flush can't re-remove it. + batch = list(self._metric_buffer) + del self._metric_buffer[: len(batch)] + step = self._progress.step + loss = self._progress.loss + if loss is not None and not math.isfinite(loss): + loss = None + duration = self._progress.elapsed_seconds try: from storage.studio_db import insert_metrics_batch, update_run_progress - - insert_metrics_batch(self.current_job_id, batch) - del self._metric_buffer[: len(batch)] - update_run_progress( - id = self.current_job_id, - step = self._progress.step, - loss = self._progress.loss - if (self._progress.loss is not None and math.isfinite(self._progress.loss)) - else None, - duration_seconds = self._progress.elapsed_seconds, - ) + insert_metrics_batch(target, batch) + update_run_progress(id = target, step = step, loss = loss, duration_seconds = duration) except Exception: - # Leave buffer intact for retry on next flush + # Re-queue the claimed batch at the front so it retries on the next flush. + with self._lock: + self._metric_buffer[:0] = batch logger.warning("Failed to flush metrics to DB", exc_info = True) @staticmethod diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py new file mode 100644 index 0000000000..457dfc8ea2 --- /dev/null +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -0,0 +1,824 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Stop-watchdog escalation for a stuck training stop. + +A save-stop signals the worker and waits for it to save and exit. On some platforms the +worker saves but then wedges in post-save GPU/driver teardown and never exits, leaving the +run stuck in "Stopping..." forever. These tests pin the bounded recovery: the watchdog +escalates to force_terminate() a short grace after "complete" (save done) or after an +absolute timeout (hang during save), and never force-kills a worker that exits cleanly. +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 + +_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 + +if not _TRAINING_PRE_IMPORTED: + sys.modules.pop("core.training.training", None) + sys.modules.pop("core.training", None) + +# The module globals hold the escalation timeouts and are the watchdog's own +# namespace; patch them here so tests run in well under a second. +_G = TrainingBackend._stop_watchdog_loop.__globals__ + + +class _FakeProc: + """A subprocess handle whose liveness and kill calls the test observes.""" + + def __init__(self, alive: bool = True): + self._alive = alive + self.pid = 4321 + self.terminated = False + self.killed = False + + def is_alive(self): + return self._alive + + def terminate(self): + self.terminated = True + + def kill(self): + self.killed = True + + def join(self, timeout = None): + pass + + +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() + + +def _record_force_terminate(monkeypatch, b): + """Replace force_terminate + escalation finalize with recorders (no DB/OS).""" + calls: list = [] + monkeypatch.setattr(b, "force_terminate", lambda target_proc = None: calls.append("force")) + monkeypatch.setattr( + b, + "_finalize_stopped_after_escalation", + lambda target_proc = None, watched_job_id = None: calls.append("final"), + ) + return calls + + +# ---------------------------------------------------------------------------- +# (a) Escalate a short grace after "complete" (save done) if still alive. +# ---------------------------------------------------------------------------- + + +def test_watchdog_escalates_after_grace_once_complete_seen(monkeypatch): + monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # ensure grace, not timeout, fires + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + proc = _FakeProc(alive = True) + b._proc = proc + b._complete_seen.set() # worker reported "complete" -> save is done + + b._start_stop_watchdog(cancel = False) + assert _wait_until( + lambda: calls == ["force", "final"] + ), "watchdog must force_terminate a worker still alive after the post-save grace" + b._stop_watchdog.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# (b) The absolute cap is a last-resort backstop, not a save killer. +# ---------------------------------------------------------------------------- + + +def test_watchdog_does_not_kill_save_still_saving_within_window(monkeypatch): + # save=True, no "complete" yet: a slow save in progress must not be force-killed + # inside the (long) absolute window. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + proc = _FakeProc(alive = True) + b._proc = proc + b._start_stop_watchdog(cancel = False) + + time.sleep(0.3) + assert calls == [], "an in-progress save must not be killed within the absolute window" + assert b._stop_watchdog.is_alive() + + proc._alive = False + b._stop_watchdog.join(timeout = 5) + + +def test_watchdog_backstop_fires_for_save_after_absolute_timeout(monkeypatch): + # Past the long save=True cap with no completion: force-terminate as last resort. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) # never trips (no complete) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + b._proc = _FakeProc(alive = True) + b._start_stop_watchdog(cancel = False) + assert _wait_until( + lambda: calls == ["force", "final"] + ), "the absolute backstop must force_terminate a save that never completes" + b._stop_watchdog.join(timeout = 5) + + +def test_cancel_uses_shorter_absolute_timeout(monkeypatch): + # A cancel has nothing to save, so it escalates on the shorter cancel cap even before + # the long save cap elapses. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # save cap would not fire + monkeypatch.setitem(_G, "_CANCEL_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + b._proc = _FakeProc(alive = True) + b._start_stop_watchdog(cancel = True) + assert _wait_until( + lambda: calls == ["force", "final"] + ), "a cancel must escalate on the shorter cancel timeout" + b._stop_watchdog.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# (c) No force-kill when the worker exits cleanly and promptly. +# ---------------------------------------------------------------------------- + + +def test_watchdog_no_op_on_clean_quick_exit(monkeypatch): + monkeypatch.setitem(_G, "_STOP_GRACE_S", 5.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 10.0) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + proc = _FakeProc(alive = True) + b._proc = proc + b._complete_seen.set() # save done; worker is about to exit on its own + + b._start_stop_watchdog(cancel = False) + # Worker exits promptly, well before the grace period elapses. + time.sleep(0.1) + proc._alive = False + + b._stop_watchdog.join(timeout = 5) + assert not b._stop_watchdog.is_alive() + assert calls == [], "a clean quick exit must not trigger force_terminate" + + +def test_watchdog_no_op_when_worker_superseded(monkeypatch): + # A stale watchdog from a prior run must never kill a new run's worker: once + # self._proc is replaced, it exits silently. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + old_proc = _FakeProc(alive = True) + b._proc = old_proc + b._complete_seen.set() + b._start_stop_watchdog(cancel = False) + + # A new run takes over the handle before the grace elapses. + b._proc = _FakeProc(alive = True) + + b._stop_watchdog.join(timeout = 5) + assert calls == [], "watchdog must not force_terminate a superseded worker" + + +def test_new_run_gets_its_own_watchdog(monkeypatch): + # A stale watchdog sleeping on an old proc must not stop a new run's stop from + # creating its own watcher. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) + b = TrainingBackend() + _record_force_terminate(monkeypatch, b) + + old_proc = _FakeProc(alive = True) + b._proc = old_proc + b._start_stop_watchdog(cancel = False) + first_wd = b._stop_watchdog + + # New run: fresh worker replaces the handle; its stop must get a new watcher + # even though the old (superseded) watchdog is still alive. + new_proc = _FakeProc(alive = True) + b._proc = new_proc + b._start_stop_watchdog(cancel = False) + second_wd = b._stop_watchdog + + try: + assert first_wd.is_alive() + assert second_wd is not first_wd, "a new run must get its own watchdog" + assert b._stop_watchdog_proc is new_proc + finally: + old_proc._alive = False + new_proc._alive = False + first_wd.join(timeout = 5) + second_wd.join(timeout = 5) + + +def test_force_terminate_targets_only_captured_proc(): + # Superseded: force_terminate(target) must not touch a different current worker. + b = TrainingBackend() + old_proc = _FakeProc(alive = True) + new_proc = _FakeProc(alive = True) + b._proc = new_proc + b.force_terminate(target_proc = old_proc) + assert new_proc.terminated is False, "must not terminate the new run's worker" + assert old_proc.terminated is False, "must not terminate a handle that is not current" + + # Matching: the captured handle is the current worker, so it is terminated. + p = _FakeProc(alive = True) + b._proc = p + b.force_terminate(target_proc = p) + assert p.terminated is True + + +# ---------------------------------------------------------------------------- +# Post-escalation finalize leaves the parent ready for a new run. +# ---------------------------------------------------------------------------- + + +def test_finalize_runs_even_if_force_terminate_raises(monkeypatch): + # A wedged child can make force_terminate() raise; finalize must still run so the + # run does not stay stuck in "Stopping...". + monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) + b = TrainingBackend() + + def _boom(target_proc = None): + raise RuntimeError("kill() failed on wedged child") + + finalized: list = [] + monkeypatch.setattr(b, "force_terminate", _boom) + monkeypatch.setattr( + b, + "_finalize_stopped_after_escalation", + lambda target_proc = None, watched_job_id = None: finalized.append(True), + ) + + b._proc = _FakeProc(alive = True) + b._complete_seen.set() + b._start_stop_watchdog(cancel = False) + + assert _wait_until( + lambda: finalized == [True] + ), "finalize must run even when force_terminate raises" + b._stop_watchdog.join(timeout = 5) + + +def test_finalize_after_escalation_clears_state(monkeypatch): + # Even if the OS never reaps the wedged worker, the parent must report the run + # stopped so the UI leaves "Stopping..." and a new run can start. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + b._proc = _FakeProc(alive = True) # wedged: still reports alive + b._should_stop = True + b.current_job_id = "job_c" + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + assert b._proc is None, "the wedged handle must be dropped so is_training_active clears" + assert b._progress.is_training is False + assert b._progress.status_message == "Training stopped." + assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" + assert b.is_training_active() is False + + +def test_finalize_after_escalation_preserves_output_dir(monkeypatch): + # A save-stop that already emitted "complete" has the checkpoint dir; run history + # must record it even if the watchdog wins the finalize race against the pump. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + b._proc = _FakeProc(alive = True) + b._should_stop = True + b.current_job_id = "job_c" + b._db_run_created = True + b._output_dir = "/tmp/outputs/run-123" + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + # _finish_stopped_run(run_id, output_dir, batch, final_step, final_loss, duration, loss_history) + assert finstop and finstop[0][0] == "job_c" + assert finstop[0][1] == "/tmp/outputs/run-123" + + +def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch): + # No worker -> nothing to escalate; the watchdog must not spawn. + b = TrainingBackend() + b._proc = None + assert b.stop_training(save = True) is True + assert b._stop_watchdog is None + + +# ---------------------------------------------------------------------------- +# (d) A stale watchdog must never clobber a run that replaced its worker. +# ---------------------------------------------------------------------------- + + +def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch): + # A /start can slip in while the watchdog force-terminates the old worker + # (is_training_active() is False once _should_stop is set and the old proc is dead). + # The escalation finalize must then leave the NEW run untouched, not drop its handle. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + old_proc = _FakeProc(alive = False) # force-terminated worker we were watching + new_proc = _FakeProc(alive = True) # a new run already took over + b._proc = new_proc + b.current_job_id = "job_new" + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = old_proc) + + assert b._proc is new_proc, "must not drop the new run's handle" + assert b._progress.is_training is True, "must not mark the new run stopped" + assert finstop == [], "must not finalize the new run in the DB" + + +def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch): + # Common case: the watched worker is still current, so finalize proceeds and + # finalizes the captured run by id. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + proc = _FakeProc(alive = False) + b._proc = proc + b.current_job_id = "job_a" + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = proc, watched_job_id = "job_a") + + assert b._proc is None + assert b._progress.is_training is False + assert finstop and finstop[0][0] == "job_a", "must finalize the captured run by id" + + +def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypatch): + # start_training updates current_job_id BEFORE it installs the new _proc, so a stale + # watchdog can enter while _proc is still the old (dead) handle. The job-id guard must + # catch this even though the proc-only guard would not. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet + b._proc = old_proc # still the old handle (== target), so proc guard would pass + b.current_job_id = "job_new" # but the new run already claimed the job id + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old") + + assert b._proc is old_proc, "must not drop the handle during a new run's startup" + assert b._progress.is_training is True, "must not mark the starting run stopped" + assert finstop == [], "must not finalize while a new run is starting up" + + +# ---------------------------------------------------------------------------- +# (e) A later cancel (save=False) tightens an in-flight save watchdog. +# ---------------------------------------------------------------------------- + + +def test_later_cancel_tightens_watchdog_timeout(monkeypatch): + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) # never trips (no complete) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # save cap would not fire + monkeypatch.setitem(_G, "_CANCEL_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + b._proc = _FakeProc(alive = True) + b._start_stop_watchdog(cancel = False) # started as a save-stop with the long cap + time.sleep(0.15) + assert calls == [], "a save-stop must not escalate on the short cancel cap yet" + + # The user now cancels the in-flight stop: the watchdog must tighten its cap. + b._cancel_requested = True + assert _wait_until( + lambda: calls == ["force", "final"] + ), "a later cancel must tighten the watchdog to the shorter cancel cap" + b._stop_watchdog.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# (f) DB finalize/flush are safe when the watchdog and pump race (see Item 4). +# ---------------------------------------------------------------------------- + + +def _install_fake_db(monkeypatch): + """Stub storage.studio_db + utils.downsample so the real DB helpers run without + SQLite. Returns the recorder dict.""" + recs = {"created": [], "finished": [], "inserted": [], "insert_ids": [], "progress_ids": []} + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + fake_db.create_run = lambda **kw: recs["created"].append(kw) + fake_db.finish_run = lambda **kw: recs["finished"].append(kw) + fake_db.insert_metrics_batch = lambda job_id, batch: ( + recs["inserted"].extend(batch), + recs["insert_ids"].append(job_id), + ) + fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id")) + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + fake_ds = _types.ModuleType("utils.downsample") + fake_ds.downsample = lambda seq, n: list(seq)[:n] + monkeypatch.setitem(sys.modules, "utils.downsample", fake_ds) + return recs + + +def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): + # The watchdog and pump can both finalize; only one call may reach finish_run. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_x" + b._db_run_created = True + b._run_finalized = False + + start = threading.Barrier(8) + + def worker(): + start.wait() + b._finalize_run_in_db(status = "stopped") + + threads = [threading.Thread(target = worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + + assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert b._run_finalized is True + + +def test_finalize_run_in_db_no_ops_on_job_mismatch(monkeypatch): + # A finalize captured for an old job must not finalize the run that replaced it. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_new" + b._db_run_created = True + b._run_finalized = False + + b._finalize_run_in_db(status = "stopped", expected_job_id = "job_old") + + assert recs["finished"] == [], "a superseded job id must not finalize the current run" + assert b._run_finalized is False + + +def test_concurrent_flush_claims_each_metric_once(monkeypatch): + # Concurrent flushes (pump periodic flush vs watchdog finalize flush) must not + # double-remove or drop buffered metrics. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_y" + b._db_run_created = True + b._metric_buffer[:] = [{"step": i} for i in range(200)] + + start = threading.Barrier(6) + + def worker(): + start.wait() + for _ in range(50): + b._flush_metrics_to_db() + + threads = [threading.Thread(target = worker) for _ in range(6)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + b._flush_metrics_to_db() # drain any remainder + + steps = sorted(m["step"] for m in recs["inserted"]) + assert steps == list(range(200)), "each metric must be inserted exactly once" + assert b._metric_buffer == [], "the buffer must be fully drained" + + +def test_flush_pins_to_passed_run_id(monkeypatch): + # A finalizer flushes to the run it captured, even if a new /start has already + # changed current_job_id. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_new" # a new run is already live + b._db_run_created = True + b._metric_buffer[:] = [{"step": 1}, {"step": 2}] + + b._flush_metrics_to_db(run_id = "job_old") + + assert recs["insert_ids"] == ["job_old"], "metrics must go to the captured run, not the new one" + assert recs["progress_ids"] == ["job_old"] + + +def test_finalize_uses_snapshot_run_id_across_new_run(monkeypatch): + # If a new /start changes current_job_id after the finalize claim but before the DB + # writes, finish_run must still target the run captured under the lock. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_x" + b._db_run_created = True + b._run_finalized = False + + def hijack(run_id = None): + # Simulate a new run taking over during the flush (after the finalize claim). + b.current_job_id = "job_y" + + monkeypatch.setattr(b, "_flush_metrics_to_db", hijack) + + b._finalize_run_in_db(status = "stopped", expected_job_id = "job_x") + + assert [f["id"] for f in recs["finished"]] == [ + "job_x" + ], "finish_run must target the captured run, not the run that replaced it" + + +# ---------------------------------------------------------------------------- +# (g) DB row creation must not be published before the insert commits. +# ---------------------------------------------------------------------------- + + +def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): + # _db_run_created must stay False while create_run is in flight, so a concurrent + # finalize can't run finish_run (an UPDATE) against a not-yet-inserted row. + b = TrainingBackend() + b.current_job_id = "job_z" + b._db_config = {"model_name": "m"} + observed: dict = {} + + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + + def _create(**kw): + observed["flag_during_create"] = b._db_run_created + observed["in_progress_during_create"] = b._db_create_in_progress + + fake_db.create_run = _create + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + + b._ensure_db_run_created() + + assert observed["flag_during_create"] is False, "flag must not be published before insert" + assert observed["in_progress_during_create"] is True + assert b._db_run_created is True, "flag must be published after a successful insert" + assert b._db_create_in_progress is False + + +def test_ensure_db_run_created_stays_unpublished_on_failure(monkeypatch): + # If create_run raises, neither flag stays set, so a later caller can retry. + b = TrainingBackend() + b.current_job_id = "job_z" + b._db_config = {"model_name": "m"} + + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + + def _boom_create(**kw): + raise RuntimeError("insert failed") + + fake_db.create_run = _boom_create + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + + b._ensure_db_run_created() + + assert b._db_run_created is False, "a failed insert must not publish the row as created" + assert b._db_create_in_progress is False, "the in-progress flag must be cleared on failure" + + +def test_ensure_db_run_created_does_not_publish_for_a_new_run(monkeypatch): + # A killed worker lets a new /start proceed while the watchdog is still creating the old + # run's row. The stale create must not publish the backend-wide flags against the new + # current_job_id, or the new run would skip inserting its own row. + b = TrainingBackend() + b.current_job_id = "job_old" + b._db_config = {"model_name": "m"} + b._db_run_created = False + b._db_create_in_progress = False + + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + + def _create(**kw): + b.current_job_id = "job_new" # a new run takes over during the slow create + + fake_db.create_run = _create + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + + b._ensure_db_run_created() + + assert b._db_run_created is False, "must not publish the created flag against the new run" + # The stale claim is left for start_training to reset, not satisfied for the new run. + assert b._db_create_in_progress is True, "must not clear the claim once the run is not current" + + +# ---------------------------------------------------------------------------- +# (h) The escalation finalizes the watched run by id (so it is never left running). +# ---------------------------------------------------------------------------- + + +def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): + # Exercise the real _finish_stopped_run against a fake DB. The watched run is finalized + # by its captured id with its buffered metrics, so a new run that starts in the gap + # after the backend goes idle can never leave the stopped run recorded running. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_old" + b._db_run_created = True + b._proc = _FakeProc(alive = False) + b._progress.is_training = True + b._progress.step = 42 + b._metric_buffer[:] = [{"step": 41}, {"step": 42}] + + b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old") + + assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id" + assert recs["finished"][0]["status"] == "stopped" + assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" + assert b._metric_buffer == [], "the captured batch must be drained" + + +def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): + # If the row does not exist and cannot be created here (no db_config, or the pump is + # mid-create), the escalation must not claim _run_finalized or call _finish_stopped_run, + # so the pump's create-then-finalize records the run. Parent state still clears. + b = TrainingBackend() + called: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a)) + + b._proc = _FakeProc(alive = False) + b.current_job_id = "job_q" + b._db_run_created = False # row not created yet + b._db_config = None # ... and cannot be created here + b._run_finalized = False + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_q") + + assert called == [], "must not finalize when the row can't be established here" + assert b._run_finalized is False, "must not claim the finalize the pump still owes" + assert b._progress.is_training is False, "parent state must still clear so the UI unsticks" + assert b._proc is None + + +def test_escalation_creates_row_then_finalizes_when_start_create_failed(monkeypatch): + # A wedged worker's pump can never finalize and would bail once _proc is dropped, so if + # the row was never created (start-time create failed) the escalation creates it and + # finalizes by id itself, recording the terminal state before dropping the handle. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_s" + b._db_config = {"model_name": "m"} # so _ensure_db_run_created can create the row + b._db_run_created = False # start-time create failed + b._proc = _FakeProc(alive = True) # wedged: still reports alive + b._should_stop = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_s") + + assert [c["id"] for c in recs["created"]] == ["job_s"], "must create the missing row" + assert [f["id"] for f in recs["finished"]] == ["job_s"], "must finish the created row by id" + assert b._proc is None, "handle dropped only after the terminal state is recorded" + assert b._db_run_created is True + + +def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): + # If a run replaces the worker while the finalize DB write is in flight, the final _proc + # drop must leave the new run's handle intact (re-guarded on target_proc). + b = TrainingBackend() + b.current_job_id = "job_old" + b._db_run_created = True + old_proc = _FakeProc(alive = False) + new_proc = _FakeProc(alive = True) + b._proc = old_proc + + def hijack(*a): + b._proc = new_proc # a new run takes over during the finalize + + monkeypatch.setattr(b, "_finish_stopped_run", hijack) + + b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old") + + assert b._proc is new_proc, "must not drop the handle a new run installed during finalize" + + +def _make_finish_raise(monkeypatch, calls): + fn = sys.modules["storage.studio_db"] + + def _boom(**kw): + calls.append(kw) + raise RuntimeError("database is locked") + + fn.finish_run = _boom + + +def test_finish_stopped_run_retries_then_unclaims_on_db_error(monkeypatch): + # The watchdog is the sole finalizer once _proc is dropped, so a transient DB error is + # retried a few times; on final failure the finalize is unclaimed (run still current). + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + _install_fake_db(monkeypatch) + tries: list = [] + _make_finish_raise(monkeypatch, tries) + b = TrainingBackend() + b.current_job_id = "job_r" + b._run_finalized = True # the caller (escalation) already claimed + + b._finish_stopped_run("job_r", None, [{"step": 1}], 1, None, None, []) + + assert len(tries) == 3, "a transient DB error must be retried before giving up" + assert b._run_finalized is False, "a persistent DB error must unclaim the finalize" + + +def test_finish_stopped_run_error_leaves_new_run_untouched(monkeypatch): + # If the watched run was superseded, a DB error must not unclaim the new run's finalize. + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + _install_fake_db(monkeypatch) + _make_finish_raise(monkeypatch, []) + b = TrainingBackend() + b.current_job_id = "job_new" # a new run is live + b._run_finalized = True # the new run's flag + + b._finish_stopped_run("job_old", None, [{"step": 1}], 1, None, None, []) + + assert b._run_finalized is True, "must not unclaim the new run's finalize"