diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index bbd9a895ab..17183484c5 100644 --- a/studio/backend/core/training/resume.py +++ b/studio/backend/core/training/resume.py @@ -4,6 +4,8 @@ """Helpers for validating resumable training outputs.""" import json +import pickletools +import zipfile from pathlib import Path from typing import Optional @@ -33,21 +35,158 @@ def _checkpoint_step(path: Path) -> int: return -1 -def get_resume_checkpoint_path(path_value: str) -> Optional[str]: +_MODEL_FILES = ( + "adapter_model.safetensors", + "adapter_model.bin", + "model.safetensors", + "pytorch_model.bin", +) +_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json") + + +def _valid_state_file(path: Path, require_tensor: bool = True) -> bool: + try: + if not path.is_file() or path.stat().st_size == 0: + return False + if path.suffix == ".safetensors": + try: + from safetensors import SafetensorError, safe_open + except ImportError: + return False + try: + with safe_open(str(path), framework = "np") as state: + return bool(state.keys()) + except SafetensorError: + return False + if path.suffix in {".bin", ".pt"}: + with zipfile.ZipFile(path) as state: + infos = state.infolist() + names = [info.filename for info in infos] + data_name = next( + (name for name in names if name == "data.pkl" or name.endswith("/data.pkl")), + None, + ) + if data_name is None: + return False + data_prefix = data_name.removesuffix("data.pkl") + "data/" + operations = list(pickletools.genops(state.read(data_name))) + if not operations or operations[-1][0].name != "STOP": + return False + if not require_tensor: + return True + # Require a non-empty tensor record; a zero-byte one fails torch.load. + return any( + info.filename.startswith(data_prefix) + and not info.is_dir() + and info.file_size > 0 + for info in infos + ) + # Unrecognized state-file formats are not usable resume state. + return False + except (OSError, ValueError, zipfile.BadZipFile): + return False + + +def _checkpoint_state(path: Path) -> Optional[int]: + try: + state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8")) + step = state.get("global_step") if isinstance(state, dict) else None + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + return None + if isinstance(step, bool) or not isinstance(step, int) or step < 0: + return None + directory_step = _checkpoint_step(path) + return step if directory_step < 0 or step == directory_step else None + + +_INDEX_SHARD_SUFFIX = { + "model.safetensors.index.json": ".safetensors", + "pytorch_model.bin.index.json": ".bin", +} + + +def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool: + # Shard must be a relative, in-format path contained in the checkpoint dir. + if not isinstance(shard, str) or not shard: + return False + if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix: + return False + try: + root = checkpoint.resolve(strict = True) + candidate = (checkpoint / shard).resolve(strict = True) + candidate.relative_to(root) + except (OSError, ValueError): + return False + return _valid_state_file(candidate) + + +def _has_model_state(path: Path) -> bool: + if any(_valid_state_file(path / name) for name in _MODEL_FILES): + return True + for name in _MODEL_INDEXES: + try: + index = json.loads((path / name).read_text(encoding = "utf-8")) + shards = set(index["weight_map"].values()) + except ( + AttributeError, + OSError, + KeyError, + TypeError, + UnicodeDecodeError, + json.JSONDecodeError, + ): + continue + expected_suffix = _INDEX_SHARD_SUFFIX[name] + if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards): + return True + return False + + +def is_resume_checkpoint_valid( + path: Path, + expected_step: Optional[int] = None, + backend: Optional[str] = None, +) -> bool: + step = _checkpoint_state(path) if path.is_dir() else None + step_valid = step is not None and (expected_step is None or step == expected_step) + if backend == "mlx": + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + else: + valid_bundle = ( + _has_model_state(path) + # optimizer/scheduler state can be validly tensor-free (e.g. SGD without + # momentum); _has_model_state still requires real model tensors. + and _valid_state_file(path / "optimizer.pt", require_tensor = False) + and _valid_state_file(path / "scheduler.pt", require_tensor = False) + ) + if backend is None and not valid_bundle: + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + return step_valid and valid_bundle + + +def get_resume_checkpoint_path( + path_value: str, expected_step: Optional[int] = None +) -> Optional[str]: path = resolve_output_dir(path_value) if not _is_under_outputs(path) or not path.is_dir(): return None - if (path / "trainer_state.json").is_file(): + if is_resume_checkpoint_valid(path, expected_step): return str(path) - checkpoints = [ - child - for child in path.glob("checkpoint-*") - if child.is_dir() and (child / "trainer_state.json").is_file() - ] - if not checkpoints: - return None - return str(max(checkpoints, key = _checkpoint_step)) + checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True) + return next( + ( + str(checkpoint) + for checkpoint in checkpoints + if _checkpoint_step(checkpoint) >= 0 + and is_resume_checkpoint_valid(checkpoint, expected_step) + ), + None, + ) def normalize_resume_output_dir(path_value: str) -> str: @@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool: def can_resume_run(run: dict) -> bool: if run.get("resumed_later"): return False + # Set when a stop-and-save failed to write a current-step checkpoint. + if run.get("resume_blocked"): + return False if _uses_s3_dataset(run): return False + status = run.get("status") + if status == "error": + # A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability. + return has_resume_state(run.get("output_dir")) + final_step = run.get("final_step") total_steps = run.get("total_steps") has_remaining_steps = ( @@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool: or total_steps <= 0 or final_step < total_steps ) - return ( - run.get("status") == "stopped" - and has_remaining_steps - and has_resume_state(run.get("output_dir")) - ) + return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir")) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index b407ba39a5..26d8c23b66 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -761,6 +761,7 @@ class TrainingBackend: # Left True after an abnormal death so _ensure_pump_alive spots a crash. self._pump_running: bool = False self._lock = threading.Lock() + self._run_intent_lock = threading.RLock() # 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 @@ -773,6 +774,7 @@ class TrainingBackend: self._progress = TrainingProgress() self._should_stop = False self._cancel_requested = False # True only for stop(save=False) + self._cancel_cleanup_output_dir: Optional[str] = None # Throttled training-status logging to the server log (not one line/step). self._last_progress_log_ts: float = 0.0 @@ -792,6 +794,8 @@ class TrainingBackend: # Job metadata self.current_job_id: Optional[str] = None self._output_dir: Optional[str] = None + self._resume_source_run_id: Optional[str] = None + self._terminal_finalize_payload: Optional[dict] = None # DB persistence self._metric_buffer: list[dict] = [] @@ -819,6 +823,7 @@ class TrainingBackend: job_id: str, *, before_spawn = None, + resume_source_run_id: Optional[str] = None, **kwargs, ) -> bool: """Spawn a subprocess to run the full training pipeline. @@ -956,6 +961,7 @@ class TrainingBackend: self.current_job_id = job_id self._should_stop = False self._cancel_requested = False + self._cancel_cleanup_output_dir = None self._complete_seen.clear() self._progress = TrainingProgress( is_training = True, status_message = "Initializing training..." @@ -972,7 +978,10 @@ class TrainingBackend: self.eval_loss_history.clear() self.eval_step_history.clear() self.eval_enabled = False - self._output_dir = None + self._output_dir = config.get("output_dir") if resume_source_run_id else None + self._progress.output_dir = self._output_dir + self._resume_source_run_id = resume_source_run_id + self._terminal_finalize_payload = None self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False @@ -990,6 +999,17 @@ class TrainingBackend: # 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() + if resume_source_run_id and not self._db_run_created: + if proc.is_alive(): + proc.terminate() + proc.join(timeout = 5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout = 2.0) + self._progress.is_training = False + self._progress.error = "Resume checkpoint is no longer available." + self._spawn_in_progress = False + return False # 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. @@ -1011,28 +1031,75 @@ class TrainingBackend: def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" - self._should_stop = True - if not save: - self._cancel_requested = True - with self._lock: - if self._stop_queue is not None: - try: - self._stop_queue.put({"type": "stop", "save": save}) - except (OSError, ValueError): - pass - # Update progress immediately for responsive UI. - 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) + with self._run_intent_lock: + with self._lock: + run_id = self.current_job_id + if not save and run_id: + persist_error: Optional[Exception] = None + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import mark_run_cancel_requested + + self._ensure_db_run_created() + with self._lock: + terminal_payload = self._terminal_finalize_payload + if ( + terminal_payload + and terminal_payload.get("expected_job_id") == run_id + ): + return False + if not mark_run_cancel_requested(run_id): + if self._db_run_created: + return False + raise RuntimeError( + "Training run disappeared before cancellation persisted" + ) + if self.current_job_id != run_id: + return False + self._should_stop = self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + persist_error = None + break + except Exception as exc: + persist_error = exc + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + if persist_error is not None: + raise RuntimeError("Failed to persist Stop-without-Save") from persist_error + with self._lock: + if self.current_job_id != run_id: + return False + if save or not run_id: + self._should_stop = True + if not save and not run_id: + self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + if self._stop_queue is not None: + try: + self._stop_queue.put({"type": "stop", "save": save}) + except (OSError, ValueError): + pass + self._progress.status_message = ( + "Stopping training and saving checkpoint..." + if save + else "Cancelling training..." + ) + self._start_stop_watchdog(cancel = not save, expected_job_id = run_id) return True - def _start_stop_watchdog(self, cancel: bool) -> None: + def _start_stop_watchdog( + self, + cancel: bool, + expected_job_id: Optional[str] = None, + ) -> 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 expected_job_id is not None and self.current_job_id != expected_job_id: + return proc = self._proc if proc is None or not proc.is_alive(): return @@ -1113,8 +1180,9 @@ class TrainingBackend: 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 + even if the worker is wedged in driver teardown; preserves output_dir on a save so + the checkpoint is kept, and clears it on a cancel (Stop without saving must not + offer resume/export). 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 @@ -1134,7 +1202,18 @@ class TrainingBackend: 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." + terminal_payload = self._terminal_finalize_kwargs() + status = terminal_payload["status"] + error_message = terminal_payload.get("error_message") + output_dir = terminal_payload["output_dir"] + clear_output_dir = terminal_payload["clear_output_dir"] + resume_blocked = bool(terminal_payload.get("resume_blocked")) + with self._lock: + if self.current_job_id != run_id: + return + self._progress.status_message = error_message or "Training stopped." + if error_message: + self._progress.error = error_message # 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() @@ -1148,7 +1227,8 @@ class TrainingBackend: batch: list = [] final_step = final_loss = duration = None loss_history: list = [] - output_dir = self._output_dir + if clear_output_dir: + self._output_dir = self._progress.output_dir = None if claim: self._run_finalized = True # claim this run's finalize batch = list(self._metric_buffer) @@ -1161,7 +1241,17 @@ class TrainingBackend: loss_history = list(self.loss_history) if claim: self._finish_stopped_run( - run_id, output_dir, batch, final_step, final_loss, duration, loss_history + run_id, + output_dir, + batch, + final_step, + final_loss, + duration, + loss_history, + status = status, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) with self._lock: if target_proc is None or self._proc is target_proc: @@ -1176,6 +1266,10 @@ class TrainingBackend: final_loss: Optional[float], duration: Optional[float], loss_history: list, + status: str = "stopped", + error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> 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, @@ -1194,14 +1288,16 @@ class TrainingBackend: sparkline = downsample(loss_history, 50) finish_run( id = run_id, - status = "stopped", + status = status, 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, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) return except Exception: @@ -1231,7 +1327,7 @@ class TrainingBackend: logger.info("Force-terminating training subprocess (pid=%s)", proc.pid) proc.terminate() cancelled = self._cancel_requested - output_dir = self._output_dir + output_dir = self._cancel_cleanup_output_dir or self._output_dir if proc is not None: proc.join(timeout = 5.0) @@ -1595,17 +1691,60 @@ class TrainingBackend: ) 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", - ) + terminal_payload = self._terminal_finalize_kwargs() + with self._lock: + if terminal_payload["clear_output_dir"]: + self._output_dir = self._progress.output_dir = None + if terminal_payload.get("error_message"): + self._progress.error = terminal_payload["error_message"] + self._progress.status_message = terminal_payload["error_message"] + self._finalize_run_in_db(**terminal_payload) except Exception: logger.exception("Training event pump: finalization after worker exit failed") self._pump_running = False return + def _has_current_resume_checkpoint(self, output_dir, step) -> bool: + # A valid checkpoint at the current step means the stop-and-save landed on + # disk even if the worker died before confirming it. + if not output_dir or not isinstance(step, int) or step <= 0: + return False + from core.training.resume import get_resume_checkpoint_path + return get_resume_checkpoint_path(output_dir, expected_step = step) is not None + + def _terminal_finalize_kwargs(self) -> dict: + with self._lock: + job_id = self.current_job_id + payload = self._terminal_finalize_payload + if payload and payload.get("expected_job_id") == job_id: + return dict(payload) + cancel, stopped = self._cancel_requested, self._should_stop + output_dir = None if cancel else self._output_dir + step = self._progress.step + existing_error = self._progress.error + status, error, blocked = ( + ("stopped", None, cancel) + if stopped + else ( + "error", + existing_error or "Training process terminated unexpectedly", + False, + ) + ) + # Block only when no valid current-step checkpoint actually landed. + if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step): + status = "error" + error = "Stop and Save ended before a valid current-step checkpoint was written." + blocked = True + return { + "status": status, + "error_message": error, + "output_dir": output_dir, + "clear_output_dir": cancel, + "resume_blocked": blocked, + "expected_job_id": job_id, + } + def _handle_event(self, event: dict) -> None: """Apply a subprocess event to local state. @@ -1764,6 +1903,15 @@ class TrainingBackend: elif etype == "eval_configured": self.eval_enabled = True + elif etype == "output_dir": + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = self._progress.output_dir = None + else: + self._output_dir = event_output_dir + db_action = "persist_output_dir" + elif etype == "status": self._progress.status_message = event.get("message", "") self._progress.is_training = True @@ -1778,7 +1926,12 @@ class TrainingBackend: self._complete_seen.set() self._progress.is_training = False self._progress.is_completed = not stopped - self._output_dir = event.get("output_dir") + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = None + else: + self._output_dir = event_output_dir self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: @@ -1788,11 +1941,16 @@ class TrainingBackend: db_action_kwargs = { "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) elif etype == "error": self._progress.is_training = False self._progress.error = event.get("error", "Unknown error") + if self._cancel_requested: + self._output_dir = self._progress.output_dir = None logger.error("Training error: %s", event.get("error")) stack = event.get("stack", "") if stack: @@ -1801,29 +1959,36 @@ class TrainingBackend: db_action = "create_and_finalize" else: db_action = "finalize" + stop_save_failed = ( + self._should_stop + and not self._cancel_requested + and not self._has_current_resume_checkpoint( + self._output_dir, self._progress.step + ) + ) db_action_kwargs = { - "status": "stopped" if self._should_stop else "error", + "status": "stopped" + if self._should_stop + and not stop_save_failed + and not event.get("keep_error_status") + else "error", "error_message": event.get("error", "Unknown error"), + "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "resume_blocked": stop_save_failed or bool(event.get("resume_blocked")), + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) # --- DB I/O outside the lock --- if db_action == "create_run": - try: - from storage.studio_db import create_run - - create_run( - id = db_action_kwargs["job_id"], - model_name = db_action_kwargs["model_name"], - dataset_name = db_action_kwargs["dataset_name"], - config_json = db_action_kwargs["config_json"], - started_at = db_action_kwargs["started_at"], - total_steps = db_action_kwargs["total_steps"], - ) - self._db_run_created = True + self._ensure_db_run_created() + if self._db_run_created: if db_action_kwargs["total_steps"]: self._db_total_steps_set = True - except Exception: - logger.warning("Failed to create DB run record", exc_info = True) + self._persist_output_dir() + elif db_action == "persist_output_dir": + self._persist_output_dir() elif db_action == "create_and_finalize": self._ensure_db_run_created() self._finalize_run_in_db(**db_action_kwargs) @@ -1842,6 +2007,22 @@ class TrainingBackend: if etype == "progress": self._log_training_progress() + def _persist_output_dir(self) -> None: + with self._lock: + if ( + not self._output_dir + or not self.current_job_id + or not self._db_run_created + or self._cancel_requested + ): + return + run_id, output_dir = self.current_job_id, self._output_dir + try: + from storage.studio_db import update_run_output_dir + update_run_output_dir(run_id, output_dir) + except Exception: + logger.warning("Failed to persist output_dir", exc_info = True) + def _log_training_progress(self) -> None: """One throttled training-status line to the server log (the per-step stream still goes to the UI via SSE): first step, then at most every 30s, plus the @@ -1875,6 +2056,7 @@ class TrainingBackend: 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).""" + self._run_intent_lock.acquire() with self._lock: if ( self._db_run_created @@ -1882,6 +2064,7 @@ class TrainingBackend: or not self.current_job_id or not self._db_config ): + self._run_intent_lock.release() return self._db_create_in_progress = True # only one caller creates job_id = self.current_job_id @@ -1898,6 +2081,12 @@ class TrainingBackend: or _s3_dataset_name(db_config.get("s3_dataset")) or "unknown" ) + with self._lock: + if self.current_job_id != job_id: + return + output_dir = self._output_dir + cancel_requested = self._cancel_requested + resumed_from_run_id = self._resume_source_run_id create_run( id = job_id, model_name = db_config["model_name"], @@ -1905,6 +2094,9 @@ class TrainingBackend: config_json = _json.dumps(db_config), started_at = started_at, total_steps = total_steps, + output_dir = output_dir, + cancel_requested = cancel_requested, + resumed_from_run_id = resumed_from_run_id, ) created = True except Exception: @@ -1919,12 +2111,15 @@ class TrainingBackend: if created: self._db_run_created = True # publish only after the insert commits self._db_create_in_progress = False + self._run_intent_lock.release() def _finalize_run_in_db( self, status: str, error_message: Optional[str] = None, output_dir: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, expected_job_id: Optional[str] = None, ) -> None: """Flush remaining metrics and mark a run finished in the DB. Claims the finalize @@ -1947,26 +2142,33 @@ class TrainingBackend: 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 + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import finish_run + from utils.downsample import downsample - sparkline = downsample(loss_history, 50) - finish_run( - id = run_id, - status = status, - 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 = error_message, - ) - 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) + finish_run( + id = run_id, + status = status, + ended_at = datetime.now(timezone.utc).isoformat(), + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, + loss_sparkline = _json.dumps(downsample(loss_history, 50)), + output_dir = output_dir, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, + ) + return + except Exception: + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + continue + with self._lock: + if self.current_job_id == run_id: + self._run_finalized = False + logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) 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, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 111f4fdd0f..2df4fa58c6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1840,8 +1840,15 @@ def _run_mlx_training(event_queue, stop_queue, config): # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it from utils.paths import ensure_dir - output_dir = _resolve_mlx_output_dir(config, model_name) + # Resume must land in the original run dir even when config lacks output_dir. + resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint( + resume_from_checkpoint + ) + output_dir = _resolve_mlx_output_dir( + {**config, "output_dir": resume_dir} if resume_dir else config, model_name + ) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 @@ -2067,6 +2074,17 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.add_eval_callback(_on_eval) + _opt_ref = [None] + _orig_build_optimizer = getattr(trainer, "_build_optimizer", None) + + if callable(_orig_build_optimizer): + + def _capture_optimizer(total_steps): + _opt_ref[0] = _orig_build_optimizer(total_steps) + return _opt_ref[0] + + trainer._build_optimizer = _capture_optimizer + # ── 11. Run training ── gc.collect() mx.synchronize() @@ -2082,31 +2100,58 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested: - if not _stop_save[0]: - # Cancel (save=False): skip saving. - _send("complete", output_dir = None, status_message = "Training cancelled") + def _finish_tracking() -> None: + # Runs on every save/finalize exit so TB/W&B never leak on early return. + if tb_writer is not None: + try: + tb_writer.close() + except Exception: + pass + if wandb_run is not None: + try: + wandb_run.finish() + except Exception: + pass + + def _stop_checkpoint_ok() -> bool: + if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir): + return True + _send( + "error", + error = ( + "Failed to save a resumable checkpoint after stop. " + "Model files were saved, but this run cannot be resumed." + ), + # A user stop finalizes as 'stopped'; keep this failure's error status so history explains it. + keep_error_status = True, + # Older checkpoints are stale; resuming would roll back past this stop. + resume_blocked = True, + ) + return False + + try: + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + # Stop-and-save promises a resumable checkpoint, not just model files. + if not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: - _send("status", status_message = "Saving stopped model...") + _send("status", status_message = "Saving model...") mx.synchronize() trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training stopped") - else: - _send("status", status_message = "Saving model...") - mx.synchronize() - trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training completed") - - if tb_writer is not None: - try: - tb_writer.close() - except Exception: - pass - if wandb_run is not None: - try: - wandb_run.finish() - except Exception: - pass + # A save-stop can race the natural final save; it made the same promise. + if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training completed") + finally: + _finish_tracking() def _is_current_process_apple_silicon() -> bool: @@ -3177,6 +3222,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) tensorboard_dir = config.get("tensorboard_dir") if config.get("enable_tensorboard", False): @@ -3296,6 +3342,61 @@ def _send_status(event_queue: Any, message: str) -> None: ) +def _emit_output_dir(event_queue: Any, output_dir: str) -> None: + try: + event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()}) + except Exception: + pass + + +def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool: + if step <= 0: + return False + from core.training.resume import is_resume_checkpoint_valid + return is_resume_checkpoint_valid( + Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx" + ) + + +def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool: + """Write a full resume checkpoint for a stopped MLX run. + + Returns True when a checkpoint for the current training step exists. + """ + step = int(getattr(trainer, "_global_step", 0) or 0) + # A periodic save or a resumed run may already cover the current step. + if _mlx_has_checkpoint_at_step(output_dir, step): + return True + if step <= 0 or optimizer is None: + return False + ckpt_dir = Path(output_dir) / f"checkpoint-{step}" + if ckpt_dir.is_symlink(): + # Refuse a symlinked dir: it could redirect writes outside output_dir. + logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir) + return False + try: + ckpt_dir.mkdir(parents = True, exist_ok = True) + from unsloth_zoo.mlx.utils import ( + save_optimizer_state, + save_trainable_adapters, + save_trainer_state, + ) + + save_trainable_adapters(trainer.model, str(ckpt_dir)) + save_optimizer_state(optimizer, str(ckpt_dir)) + save_trainer_state( + { + "global_step": step, + "train_loss_history": list(getattr(trainer, "_train_loss_history", [])), + }, + str(ckpt_dir), + ) + logger.info("Saved stop checkpoint to %s", ckpt_dir) + except Exception: + logger.exception("Failed to write stop checkpoint under %s", output_dir) + return _mlx_has_checkpoint_at_step(output_dir, step) + + def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None: """Self-contained embedding model training pipeline. @@ -3660,6 +3761,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> config.get("project_name"), ) output_dir = str(resolve_output_dir(output_dir)) + _emit_output_dir(event_queue, output_dir) num_epochs = config.get("num_epochs", 2) batch_size = config.get("batch_size", 256) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0b50f63b95..0f88b78f9f 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -505,6 +505,13 @@ class TrainingStartRequest(BaseModel): description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", ) + @field_validator("target_modules", mode = "before") + @classmethod + def _normalize_target_modules(cls, value: Any) -> Any: + # Sanitized non-LoRA history stores the unused value as null; treat it as a + # fresh request's omitted/default empty list on resume. + return [] if value is None else value + @model_validator(mode = "after") def _validate_streaming_splits(self) -> "TrainingStartRequest": # Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]" diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 53b1c4d991..a8a9874b1b 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -196,6 +196,7 @@ async def start_training( request.local_eval_datasets, "Local eval dataset" ) resume_output_dir: Optional[str] = None + resume_run: Optional[dict] = None if request.resume_from_checkpoint: try: resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint) @@ -208,7 +209,7 @@ async def start_training( if not resume_run or not can_resume_run(resume_run): raise HTTPException( status_code = 400, - detail = "Resume checkpoint must belong to a stopped run with saved trainer state.", + detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.", ) resume_checkpoint = get_resume_checkpoint_path(resume_output_dir) if not resume_checkpoint: @@ -458,7 +459,10 @@ async def start_training( try: success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + job_id = job_id, + before_spawn = _free_vram_for_training, + resume_source_run_id = resume_run["id"] if resume_run else None, + **training_kwargs, ) except SidecarSwapInProgress as exc: # Expected loss of the race against a sidecar install: a retryable @@ -521,7 +525,10 @@ async def stop_training( status = "idle", message = "No training job is currently running" ) - backend.stop_training(save = body.save) + if not backend.stop_training(save = body.save): + return TrainingStopResponse( + status = "idle", message = "No training job is currently running" + ) return TrainingStopResponse( status = "stopped", @@ -637,9 +644,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject "loss": getattr(progress, "loss", None), "learning_rate": getattr(progress, "learning_rate", None), } - output_dir = getattr(backend, "_output_dir", None) - if output_dir: - details["output_dir"] = output_dir + # Always present: an explicit null tells the client to drop a cached + # path (stop without save clears the run's output_dir). + details["output_dir"] = getattr(backend, "_output_dir", None) or None # Metric history for chart recovery after SSE reconnection. metric_history = None diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index d889894d04..6972e7b7ff 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -192,13 +192,18 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error_message TEXT, duration_seconds REAL, loss_sparkline TEXT, - display_name TEXT + display_name TEXT, + resume_blocked INTEGER NOT NULL DEFAULT 0 ) """ ) existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()} if "display_name" not in existing_cols: conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT") + if "resume_blocked" not in existing_cols: + conn.execute( + "ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS training_metrics ( @@ -734,16 +739,43 @@ def create_run( config_json: str, started_at: str, total_steps: Optional[int], + *, + output_dir: Optional[str] = None, + cancel_requested: bool = False, + resumed_from_run_id: Optional[str] = None, ) -> None: conn = get_connection() try: conn.execute( """ - INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO training_runs ( + id, model_name, dataset_name, config_json, started_at, total_steps, + output_dir, resume_blocked + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, model_name, dataset_name, config_json, started_at, total_steps), + ( + id, + model_name, + dataset_name, + config_json, + started_at, + total_steps, + None if cancel_requested else output_dir, + int(cancel_requested), + ), ) + if resumed_from_run_id: + claimed = conn.execute( + """ + UPDATE training_runs SET resume_blocked = 1 + WHERE id = ? AND status IN ('stopped', 'error') + AND output_dir = ? AND resume_blocked = 0 + """, + (resumed_from_run_id, output_dir), + ) + if claimed.rowcount != 1: + raise RuntimeError("Resume source is no longer available") conn.commit() finally: conn.close() @@ -786,6 +818,8 @@ def finish_run( loss_sparkline: Optional[str] = None, output_dir: Optional[str] = None, error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: conn = get_connection() try: @@ -793,9 +827,16 @@ def finish_run( """ UPDATE training_runs SET status = ?, ended_at = ?, final_step = ?, final_loss = ?, - duration_seconds = ?, loss_sparkline = ?, output_dir = ?, - error_message = ? - WHERE id = ? + duration_seconds = ?, loss_sparkline = ?, + output_dir = CASE + WHEN resume_blocked = 1 OR ? = 1 THEN NULL + WHEN ? IS NOT NULL THEN ? + WHEN ? IN ('error', 'stopped') THEN output_dir + ELSE NULL + END, + error_message = ?, + resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END + WHERE id = ? AND status = 'running' """, ( status, @@ -804,8 +845,13 @@ def finish_run( final_loss, duration_seconds, loss_sparkline, + int(clear_output_dir), output_dir, + output_dir, + status, error_message, + int(clear_output_dir), + int(resume_blocked), id, ), ) @@ -865,6 +911,38 @@ def update_run_display_name(id: str, display_name: Optional[str]) -> None: conn.close() +def update_run_output_dir(id: str, output_dir: Optional[str]) -> None: + conn = get_connection() + try: + conn.execute( + """ + UPDATE training_runs SET output_dir = ? + WHERE id = ? AND status = 'running' AND resume_blocked = 0 + """, + (output_dir, id), + ) + conn.commit() + finally: + conn.close() + + +def mark_run_cancel_requested(id: str) -> bool: + """Clear resume/export state only while the exact run is still active.""" + conn = get_connection() + try: + cursor = conn.execute( + """ + UPDATE training_runs SET output_dir = NULL, resume_blocked = 1 + WHERE id = ? AND status = 'running' + """, + (id,), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + def list_runs(limit: int = 50, offset: int = 0) -> dict: conn = get_connection() try: @@ -874,15 +952,15 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at, r.ended_at, r.total_steps, r.final_step, r.final_loss, r.output_dir, r.duration_seconds, r.error_message, - r.loss_sparkline, r.display_name, r.config_json, + r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -917,13 +995,13 @@ def get_run(id: str) -> Optional[dict]: """ SELECT r.*, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -958,12 +1036,12 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: 0 AS resumed_later FROM training_runs r WHERE r.output_dir = ? - AND r.status = 'stopped' + AND r.status IN ('stopped', 'error') AND NOT EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) ORDER BY r.started_at DESC @@ -1066,8 +1144,12 @@ def cleanup_orphaned_runs() -> None: conn.execute( """ UPDATE training_runs - SET status = 'error', - error_message = 'Server restarted during training', + SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END, + error_message = CASE + WHEN resume_blocked = 1 THEN NULL + ELSE 'Server restarted during training' + END, + output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END, ended_at = ? WHERE status = 'running' """, diff --git a/studio/backend/tests/test_mlx_stop_checkpoint.py b/studio/backend/tests/test_mlx_stop_checkpoint.py new file mode 100644 index 0000000000..d4a00cc6c8 --- /dev/null +++ b/studio/backend/tests/test_mlx_stop_checkpoint.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for MLX stop-and-save checkpoint handling.""" + +import importlib.util +import json +import sys +import types +from pathlib import Path + +import numpy as np +from safetensors.numpy import save_file + + +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load_worker_module(): + spec = importlib.util.spec_from_file_location( + "training_worker_under_test", + _BACKEND / "core" / "training" / "worker.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +worker = _load_worker_module() + + +class _FakeTrainer: + def __init__(self, step: int): + self._global_step = step + self._train_loss_history = [] + self.model = object() + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors") + save_file( + {"state": np.ones(1, dtype = np.float32)}, + checkpoint / "optimizer_state.safetensors", + ) + return checkpoint + + +def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._mlx_has_checkpoint_at_step(out, 5) is True + + +def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True + + +def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists( + tmp_path, monkeypatch +): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + saved_steps: list[int] = [] + + def _save_state(_value, path, name): + save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name)) + + def _save_trainer_state(state, ckpt_dir, **_kwargs): + Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8") + saved_steps.append(int(state["global_step"])) + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = lambda model, path: _save_state( + model, path, "adapters.safetensors" + ), + save_optimizer_state = lambda optimizer, path: _save_state( + optimizer, path, "optimizer_state.safetensors" + ), + save_trainer_state = _save_trainer_state, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True + assert saved_steps == [10] + assert (out / "checkpoint-10" / "trainer_state.json").is_file() + + +def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path): + out = tmp_path / "outputs" / "run_x" + ckpt = out / "checkpoint-5" + ckpt.mkdir(parents = True) + (ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8") + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path): + # An older checkpoint does not cover the current step, so this still fails. + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False + + +def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + def _boom(*_args, **_kwargs): + raise RuntimeError("save failed") + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = _boom, + save_optimizer_state = lambda *_a, **_k: None, + save_trainer_state = lambda *_a, **_k: None, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py index d75b205f35..5d2a218482 100644 --- a/studio/backend/tests/test_training_pump_resilience.py +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -310,6 +310,80 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): assert b._pump_running is False +def test_interrupted_cancel_clears_in_memory_output_dir(monkeypatch): + # Stop-without-save interrupted before its complete event: /status must not + # keep serving the cleared run's output_dir. + 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)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert b._output_dir is None + assert finalized.get("status") == "stopped" + assert finalized.get("output_dir") is None + assert finalized.get("clear_output_dir") is True + + +def test_worker_exit_reuses_terminal_stop_save_error(monkeypatch): + 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)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = False + b._output_dir = "/out/x" + b.current_job_id = "job-x" + b._terminal_finalize_payload = { + "status": "error", + "error_message": "checkpoint failed", + "output_dir": "/out/x", + "clear_output_dir": False, + "resume_blocked": True, + "expected_job_id": "job-x", + } + + b._pump_loop() + + assert b._output_dir == "/out/x" + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + assert finalized.get("resume_blocked") is True + + +def test_dead_worker_crash_preserves_output_dir(monkeypatch): + # A crash (no stop requested) after output_dir was emitted must keep the dir + # in the error finalize: checkpoints under it may still exist. + 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)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") 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 diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py index 91fdac9961..51425b0428 100644 --- a/studio/backend/tests/test_training_resume.py +++ b/studio/backend/tests/test_training_resume.py @@ -7,6 +7,9 @@ import importlib.util import json from pathlib import Path +import pytest +import torch + _BACKEND = Path(__file__).resolve().parents[1] @@ -25,6 +28,30 @@ def _load_resume_module(): resume = _load_resume_module() +def test_resume_request_accepts_sanitized_null_target_modules(): + from models.training import TrainingStartRequest + request = TrainingStartRequest( + model_name = "unsloth/Qwen3-0.6B", + training_type = "Full Finetuning", + format_type = "alpaca", + target_modules = None, + ) + + assert request.target_modules == [] + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + torch.save({"weight": torch.ones(1)}, checkpoint / "adapter_model.bin") + torch.save({"state": {0: torch.ones(1)}}, checkpoint / "optimizer.pt") + torch.save({"last_epoch": step}, checkpoint / "scheduler.pt") + return checkpoint + + def _stopped_run(**overrides): run = { "status": "stopped", @@ -44,6 +71,36 @@ def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch): assert resume.can_resume_run(_stopped_run()) is True +def test_can_resume_run_allows_errored_run_with_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error")) is True + + +def test_can_resume_run_rejects_errored_run_without_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: False) + + assert resume.can_resume_run(_stopped_run(status = "error")) is False + + +def test_can_resume_run_allows_errored_run_at_final_step(monkeypatch): + # A save-time crash records final_step == total_steps; resuming re-runs the + # final-save path from the checkpoint. + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(status = "error", final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is True + + +def test_can_resume_run_rejects_stopped_run_at_final_step(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is False + + def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) @@ -91,3 +148,444 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path) result = studio_db.list_runs() assert result["runs"][0]["config_json"] == config_json + + +def test_crashed_run_with_persisted_output_dir_is_resumable(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-crash", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-crash", str(out)) + conn = studio_db.get_connection() + conn.execute("UPDATE training_runs SET status = 'error' WHERE id = 'run-crash'") + conn.commit() + conn.close() + + run = studio_db.get_run("run-crash") + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is True + + +def test_checkpoint_discovery_skips_malformed_newest(monkeypatch, tmp_path): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_x" + valid = _write_checkpoint(out, 5) + (_write_checkpoint(out, 8) / "scheduler.pt").unlink() + malformed = out / "checkpoint-10" + malformed.mkdir() + (malformed / "trainer_state.json").write_text(json.dumps({"global_step": 10}), encoding = "utf-8") + (malformed / "adapter_model.bin").write_bytes(b"not a torch archive") + (malformed / "optimizer.pt").write_bytes(b"not a torch archive") + + assert resume.get_resume_checkpoint_path(str(out)) == str(valid) + + +def test_completed_run_keeps_output_dir_and_rejects_stale_cancel(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "completed", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.mark_run_cancel_requested("r") is False + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.get_run("r")["resume_blocked"] == 0 + + +def test_finish_run_clears_output_dir_for_stop_without_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + conn = studio_db.get_connection() + conn.execute( + "UPDATE training_runs SET status = 'running', output_dir = '/out/x', resume_blocked = 0 WHERE id = 'r'" + ) + conn.commit() + conn.close() + studio_db.mark_run_cancel_requested("r") + studio_db.cleanup_orphaned_runs() + assert studio_db.get_run("r")["status"] == "stopped" + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_clears_output_dir_on_cancel_error_finalize(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = "worker failed during cancel", + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_preserves_output_dir_for_interrupted_stop_and_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + + +def test_resumed_errored_run_is_not_offered_again(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + with pytest.raises(RuntimeError, match = "no longer available"): + studio_db.create_run( + id = "run-duplicate", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:01Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + assert studio_db.get_run("run-duplicate") is None + studio_db.finish_run( + id = "run-new", + status = "error", + ended_at = "2026-01-02T00:05:00Z", + final_step = 15, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed again", + ) + + old_run = studio_db.get_run("run-old") + new_run = studio_db.get_run("run-new") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert new_run["resumed_later"] == 0 + assert resume.can_resume_run(new_run) is True + assert studio_db.get_resumable_run_by_output_dir(str(out))["id"] == "run-new" + + +def test_running_continuation_blocks_older_resume(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + + old_run = studio_db.get_run("run-old") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert studio_db.get_resumable_run_by_output_dir(str(out)) is None + + +def test_stop_save_checkpoint_failure_keeps_error_status(monkeypatch, tmp_path): + # A stop-and-save whose checkpoint write failed must finalize as an error so + # history explains the missing resume state (keep_error_status flag). + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-failed-save", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-failed-save" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + } + ) + + run = studio_db.get_run("run-failed-save") + assert run["status"] == "error" + assert "resumable checkpoint" in run["error_message"] + + +def test_can_resume_run_rejects_resume_blocked_run(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error", resume_blocked = 1)) is False + + +def test_stop_save_checkpoint_failure_with_stale_checkpoint_is_not_resumable(monkeypatch, tmp_path): + # A failed stop-and-save must not offer Resume from an older periodic checkpoint; + # that would roll back past the recorded final step. + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-stale-ckpt", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-stale-ckpt", str(out)) + backend = TrainingBackend() + backend.current_job_id = "run-stale-ckpt" + backend._db_run_created = True + backend._should_stop = True + backend._output_dir = str(out) + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + "resume_blocked": True, + } + ) + + run = studio_db.get_run("run-stale-ckpt") + assert run["status"] == "error" + assert run["resume_blocked"] == 1 + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is False + + +def test_user_stop_error_without_checkpoint_ack_is_blocked(monkeypatch, tmp_path): + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-user-stop", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-user-stop" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event({"type": "error", "error": "interrupted"}) + + run = studio_db.get_run("run-user-stop") + assert run["status"] == "error" and run["resume_blocked"] == 1 + + +def test_terminal_fallback_keeps_resumable_when_current_checkpoint_landed(monkeypatch, tmp_path): + # Worker died before its terminal event, but a valid current-step checkpoint + # is on disk: the fallback must keep the run resumable, not block it. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_ok" + _write_checkpoint(out, 7) + + backend = TrainingBackend() + backend.current_job_id = "run-ok" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "stopped" + assert kwargs["resume_blocked"] is False + + +def test_terminal_fallback_blocks_when_no_current_checkpoint(monkeypatch, tmp_path): + # Same path, but only a stale (older-step) checkpoint exists: must block. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_stale" + _write_checkpoint(out, 5) + + backend = TrainingBackend() + backend.current_job_id = "run-stale" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "error" + assert kwargs["resume_blocked"] is True diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 0cd702bce2..cbe2082e82 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -353,7 +353,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): # 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)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) # wedged: still reports alive b._should_stop = True @@ -365,7 +365,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): 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 "valid current-step checkpoint" in b._progress.status_message assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" assert b.is_training_active() is False @@ -375,7 +375,7 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): # 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)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) b._should_stop = True @@ -390,6 +390,28 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): assert finstop[0][1] == "/tmp/outputs/run-123" +def test_finalize_after_escalation_clears_output_dir_on_cancel(monkeypatch): + # Stop-without-saving promises no resume: a cancel that escalates through the + # watchdog clears the persisted output_dir, not a checkpoint path. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append((a, k))) + + b._proc = _FakeProc(alive = True) + b._should_stop = True + b._cancel_requested = 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") + + assert finstop and finstop[0][0][0] == "job_c" + assert finstop[0][0][1] is None, "a cancelled run must not record a checkpoint path" + assert finstop[0][1].get("clear_output_dir") is True + assert b._output_dir is None, "/status must stop exposing the cancelled run's dir" + + def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch): # No worker -> nothing to escalate; the watchdog must not spawn. b = TrainingBackend() @@ -409,7 +431,7 @@ def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch): # 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)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: 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 @@ -430,7 +452,7 @@ def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch): # finalizes the captured run by id. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) proc = _FakeProc(alive = False) b._proc = proc @@ -451,7 +473,7 @@ def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypat # 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)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: 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 @@ -509,6 +531,7 @@ def _install_fake_db(monkeypatch): recs["insert_ids"].append(job_id), ) fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id")) + fake_db.mark_run_cancel_requested = lambda _run_id: True fake_storage.studio_db = fake_db monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) @@ -518,9 +541,48 @@ def _install_fake_db(monkeypatch): return recs +def test_stop_without_save_creates_missing_row_before_signal(monkeypatch): + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id, b._db_config = "job_missing", {"model_name": "m"} + b._stop_queue = queue.Queue() + assert b.stop_training(save = False) is True + assert [run["id"] for run in recs["created"]] == ["job_missing"] + assert b._stop_queue.get_nowait() == {"type": "stop", "save": False} + + b._cancel_requested = b._should_stop = False + sys.modules["storage.studio_db"].mark_run_cancel_requested = lambda _run_id: False + assert b.stop_training(save = False) is False + assert not b._cancel_requested and b._stop_queue.empty() + + new_queue = queue.Queue() + b.current_job_id, b._db_run_created = "job_old", True + b._cancel_requested = b._should_stop = False + + def _supersede(_run_id): + b.current_job_id = "job_new" + b._stop_queue = new_queue + return True + + sys.modules["storage.studio_db"].mark_run_cancel_requested = _supersede + assert b.stop_training(save = False) is False + assert not b._cancel_requested and new_queue.empty() + + 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) + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + attempts = 0 + + def flaky_finish(**kw): + nonlocal attempts + attempts += 1 + if attempts < 3: + raise RuntimeError("database is locked") + recs["finished"].append(kw) + + sys.modules["storage.studio_db"].finish_run = flaky_finish b = TrainingBackend() b.current_job_id = "job_x" b._db_run_created = True @@ -539,6 +601,7 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): t.join(timeout = 5) assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert attempts == 3 assert b._run_finalized is True @@ -646,7 +709,13 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) - b._ensure_db_run_created() + b._run_intent_lock.acquire() + creator = threading.Thread(target = b._ensure_db_run_created) + creator.start() + time.sleep(0.02) + assert b._db_create_in_progress is False + b._run_intent_lock.release() + creator.join(timeout = 5) assert observed["flag_during_create"] is False, "flag must not be published before insert" assert observed["in_progress_during_create"] is True @@ -718,6 +787,7 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b = TrainingBackend() b.current_job_id = "job_old" b._db_run_created = True + b._should_stop = True b._proc = _FakeProc(alive = False) b._progress.is_training = True b._progress.step = 42 @@ -726,7 +796,8 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): 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["finished"][0]["status"] == "error" + assert recs["finished"][0]["resume_blocked"] is True assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" assert b._metric_buffer == [], "the captured batch must be drained" @@ -737,7 +808,7 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): # 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)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: called.append(a)) b._proc = _FakeProc(alive = False) b.current_job_id = "job_q" @@ -785,7 +856,7 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): new_proc = _FakeProc(alive = True) b._proc = old_proc - def hijack(*a): + def hijack(*a, **k): b._proc = new_proc # a new run takes over during the finalize monkeypatch.setattr(b, "_finish_stopped_run", hijack) diff --git a/studio/frontend/src/features/studio/historical-training-view.tsx b/studio/frontend/src/features/studio/historical-training-view.tsx index b6ec06b06a..4a123875a8 100644 --- a/studio/frontend/src/features/studio/historical-training-view.tsx +++ b/studio/frontend/src/features/studio/historical-training-view.tsx @@ -2,19 +2,29 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { TrainingViewData } from "@/features/training"; -import { getTrainingRun, onTrainingRunUpdated } from "@/features/training"; +import { + getTrainingRun, + onTrainingRunUpdated, + useTrainingActions, + useTrainingRuntimeStore, +} from "@/features/training"; import type { TrainingRunDetailResponse } from "@/features/training"; import { parseBackendTrainingMethod } from "@/features/training/lib/training-methods"; import { type ReactElement, useEffect, useState } from "react"; import { ChartsSection } from "./sections/charts-section"; import { ProgressSection } from "./sections/progress-section"; import { mapRunConfigToOverride } from "./sections/run-config-override"; +import { Button } from "@/components/ui/button"; +import { Spinner } from "@/components/ui/spinner"; +import { PlayIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { translate, useT } from "@/i18n"; type StudioT = ReturnType; interface HistoricalTrainingViewProps { runId: string; + onResumeStarted?: () => void; } function mapToViewData( @@ -93,10 +103,27 @@ function mapToViewData( export function HistoricalTrainingView({ runId, + onResumeStarted, }: HistoricalTrainingViewProps): ReactElement { const t = useT(); const [detail, setDetail] = useState(null); const [error, setError] = useState(null); + const [resuming, setResuming] = useState(false); + const { resumeTrainingRunFromHistory } = useTrainingActions(); + const isStarting = useTrainingRuntimeStore((state) => state.isStarting); + const isTrainingRunning = useTrainingRuntimeStore( + (state) => state.isTrainingRunning, + ); + + const handleResume = async () => { + setResuming(true); + try { + const ok = await resumeTrainingRunFromHistory(runId); + if (ok) onResumeStarted?.(); + } finally { + setResuming(false); + } + }; // Derive loading from detail/error; no separate state. const loading = detail === null && error === null; @@ -152,6 +179,27 @@ export function HistoricalTrainingView({ return (
+ {detail.run.can_resume && ( +
+ +
+ )} other.id !== run.id && other.output_dir === run.output_dir && - (other.status === "stopped" || other.status === "completed") && + (other.status === "stopped" || + other.status === "completed" || + other.status === "error" || + other.status === "running") && new Date(other.started_at).getTime() > startedAt, ); } diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 4d3a32fab3..e575fbacd7 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -217,7 +217,13 @@ export function StudioPage(): ReactElement { {selectedHistoryRunId ? ( - + { + setSelectedHistoryRunId(null); + handleTabChange("current-run"); + }} + /> ) : ( { if (runId === currentJobId && isTrainingRunning) { diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index 0d32e066d5..9bb8206f28 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -217,7 +217,7 @@ export function useTrainingActions() { const detail = await getTrainingRun(runId); const outputDir = detail.run.output_dir; if (!detail.run.can_resume || !outputDir) { - throw new Error("Only stopped runs with a saved checkpoint can be resumed."); + throw new Error("Only stopped or errored runs with a saved checkpoint can be resumed."); } primeNativeNotificationPermission().catch(() => undefined); diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index acc80a3be2..131ea30d30 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -223,7 +223,10 @@ export const useTrainingRuntimeStore = create()((set) => ( typeof detailLr === "number" ? detailLr : state.currentLearningRate, currentEpoch: typeof detailEpoch === "number" ? detailEpoch : state.currentEpoch, - outputDir: payload.details?.output_dir ?? state.outputDir, + outputDir: + payload.details?.output_dir !== undefined + ? payload.details.output_dir + : state.outputDir, lossHistory: metricHistory.lossHistory ?? state.lossHistory, lrHistory: metricHistory.lrHistory ?? state.lrHistory, gradNormHistory: metricHistory.gradNormHistory ?? state.gradNormHistory, diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index 8ed0ce0037..373e876653 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -26,7 +26,8 @@ export interface TrainingStatusResponse { total_steps?: number; loss?: number; learning_rate?: number; - output_dir?: string; + // null = explicit clear (run stopped without saving); absent = unchanged. + output_dir?: string | null; } | null; metric_history?: { steps?: number[]; diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 3e1c1e5a76..4f859aa91c 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -1053,7 +1053,8 @@ export const en = { continueAction: "Continue Training", cancelAction: "Cancel Training", stopTitle: "Stop Training", - stopDescription: "Choose how you want to stop the current training run.", + stopDescription: + "Choose how you want to stop the current training run. Stop and Save writes a checkpoint you can resume from later; Stop cannot be resumed.", stopAction: "Stop", stopping: "Stopping...", stopAndSave: "Stop and Save", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 37c73086cf..9984fc324f 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -901,7 +901,8 @@ export const zhCN = { continueAction: "继续训练", cancelAction: "取消训练", stopTitle: "停止训练", - stopDescription: "选择如何停止当前训练运行。", + stopDescription: + "选择如何停止当前训练运行。“停止并保存”会写入检查点，之后可从该处恢复；“停止”则无法恢复。", stopAction: "停止", stopping: "停止中...", stopAndSave: "停止并保存",