f372da407b* Expose MLX grad value clipping in Studio * update test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * dataset ordering + wd * fix mlx smoke step expectations * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cast norm activation output back to original input dtype * address mlx studio review feedback * Fix present-but-None seed override for PR #5656 studio/backend/core/training/worker.py `config.get("model_random_state", random_seed)` only fills the default when the key is absent. When a caller passes `config["model_random_state"] = None` explicitly (which happens any time a JSON payload sends an explicit `null`), the old code forwarded `None` to FastMLXModel and disabled deterministic init silently. Same for `lora_random_state`. Treat absent and explicit None the same way: fall back to random_seed. studio/backend/tests/test_training_raw_support.py Update the source-string assertions to match the new lines. * Guard optional MLXTrainingConfig fields and normalize random_seed for PR #5656 The MLX worker now passes `cast_norm_output_to_input_dtype` and `dataset_order` only when the linked unsloth-zoo dataclass actually declares them. Released zoo trees that predate the paired PR can still construct `MLXTrainingConfig` without raising `TypeError: unexpected keyword argument`. Once the dependency floor is bumped to a release that contains both fields, the feature-detect guards become no-ops. `random_seed = config.get("random_seed", 3407)` was unguarded against explicit `None` from raw / backend callers. The same value seeded the trainer and was the fallback target for `model_random_state` / `lora_random_state`. Normalize once at the top of the function and use the normalized value everywhere so an explicit `None` cannot reach FastMLXModel / get_peft_model / MLXTrainingConfig. Existing seed source-pattern test updated to match the new normalize helper. New test asserts the feature-detection guards exist and that the unconditional kwargs do not include the gated fields. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Normalize seed / cast / max_grad_value at TrainingBackend for PR #5656 Round-3 review consensus: the per-field guards that landed in the MLX worker only protect the MLX path. The same `TrainingBackend.start_training` config still reaches the CUDA/text trainer at `worker.py:2267`, the embedding LoRA init at `worker.py:2450`, and embedding TrainingArguments at `worker.py:2624` with raw `None` values, so an explicit `random_seed=None` from a raw / backend caller still breaks non-MLX training even after the previous fix. Move the normalization into `TrainingBackend.start_training` itself, where it runs once for every training mode: - `_coerce_seed(value)`: explicit `None`, non-int, or absent all become 3407. Every downstream worker now sees an int. - `_coerce_optional_bool(value, default)`: explicit `None` falls back to `default` instead of `bool(None) == False`. Also normalizes the common raw-config / YAML string aliases ("true" / "false" / "0" / "1"). Used for `cast_norm_output_to_input_dtype`. - `_coerce_optional_nonneg_float(name, value)`: rejects negative numerics from raw / backend callers, matching the Pydantic `ge=0` constraint the HTTP route already enforces. Used for `max_grad_value`. worker.py MLX path: the existing `bool(config.get(key, True))` for `cast_norm_output_to_input_dtype` was changed to also fall back on explicit `None`, so direct worker callers (bypassing `TrainingBackend.start_training`) are equally safe. `max_grad_value` also raises on negative values inside the worker for the same reason. TrainingStartRequest.random_seed default bumped from 42 to 3407 so direct REST callers that omit the field receive the same default as the Studio frontend and the MLX worker. New regression test exercises the three new helpers across explicit None, valid values, string aliases, and negative-value rejection. * Tighten feature-detect test paren tracking for PR #5656 The block-extraction used , which stops at the first inner closing paren (e.g. ) and would silently miss a future unconditional / added later in the same dict literal. Switched to proper paren-depth tracking so the unconditional block is checked end-to-end. * Shorten verbose comments in MLX Studio backend * Handle MLX Studio EOS appending by mode * Wire MLX leaf norm clipping through Studio * Respect VLM layer filters for explicit LoRA targets Rationale / guardrails for the local Studio/vision push: When callers provide explicit VLM LoRA target_modules together with layer filters, FastVisionModel still needs to route the explicit targets through get_peft_regex. Otherwise the layer filters are ignored and adapters can be attached outside the requested language/vision scope. Do not revert this to plain list(target_modules) for explicit module lists. The CUDA/Studio-facing contract is that explicit targets and layer filters compose: target_modules selects module names, while finetune_language_layers / finetune_vision_layers / finetune_attention_modules / finetune_mlp_modules constrain where those targets are allowed. The regression test covers the language-only explicit q_proj case and source-checks that explicit targets are wrapped through get_peft_regex when filters are active. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refresh MLX smoke clip-config note for leaf_norm default Trim the 11-line comment block to 5 lines and correct the stale claim that MLXTrainingConfig defaults to max_grad_value=1.0. The new default is max_grad_leaf_norm=1.0 (same memory profile as elementwise but direction-preserving). The smoke still pins max_grad_value=1.0 explicitly to keep the 13-seed pass-rate fixture stable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Forward max_grad_leaf_norm through the training route and warn when layer filters constrain explicit target_modules for PR #5656 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han-Chen <info@unsloth.ai> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
1054 lines
42 KiB
Python
1054 lines
42 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Training backend — subprocess orchestrator.
|
|
|
|
Each job runs in a fresh spawn subprocess (solving transformers version-switching);
|
|
the in-process UnslothTrainer singleton is only used inside the worker. This file
|
|
orchestrates the subprocess lifecycle, pumps events from the worker's mp.Queue, and
|
|
exposes the same API to routes/training.py. Pattern follows data_recipe/jobs/manager.py.
|
|
"""
|
|
|
|
import json as _json
|
|
import math
|
|
import multiprocessing as mp
|
|
import os
|
|
import queue
|
|
import re
|
|
import shutil
|
|
import threading
|
|
import time
|
|
import structlog
|
|
from datetime import datetime, timezone
|
|
from loggers import get_logger
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple, Any
|
|
|
|
import matplotlib.pyplot as plt
|
|
from utils.hardware import prepare_gpu_selection
|
|
from utils.native_path_leases import (
|
|
native_path_secret_removed_for_child_start,
|
|
run_without_native_path_secret,
|
|
)
|
|
from utils.paths import outputs_root
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def _coerce_seed(value, default = 3407) -> int:
|
|
"""Normalize None / non-int to `default` (transformers.set_seed(None) raises)."""
|
|
if value is None:
|
|
return int(default)
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return int(default)
|
|
|
|
|
|
def _coerce_optional_bool(value, default: bool) -> bool:
|
|
"""Treat explicit None as `default` instead of `bool(None) == False`."""
|
|
if value is None:
|
|
return bool(default)
|
|
if isinstance(value, str):
|
|
normalized = value.strip().lower()
|
|
if normalized in ("true", "1", "yes", "on"):
|
|
return True
|
|
if normalized in ("false", "0", "no", "off", ""):
|
|
return False
|
|
return bool(value)
|
|
|
|
|
|
def _coerce_optional_nonneg_float(name: str, value):
|
|
"""Reject negatives; HTTP `ge=0` doesn't cover raw `**kwargs` callers."""
|
|
if value is None:
|
|
return None
|
|
try:
|
|
coerced = float(value)
|
|
except (TypeError, ValueError):
|
|
raise ValueError(f"Unsloth: {name}={value!r} must be a non-negative float or None.")
|
|
if coerced < 0:
|
|
raise ValueError(f"Unsloth: {name}={coerced} must be >= 0 (use 0 or None to disable).")
|
|
return coerced
|
|
|
|
|
|
_HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$")
|
|
|
|
|
|
def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]:
|
|
db_config = {
|
|
k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"}
|
|
}
|
|
s3_config = config.get("s3_config")
|
|
if hasattr(s3_config, "model_dump"):
|
|
s3_config = s3_config.model_dump()
|
|
if isinstance(s3_config, dict) and s3_config:
|
|
db_config["dataset_source"] = "s3"
|
|
db_config["s3_dataset"] = {
|
|
"bucket": s3_config.get("bucket"),
|
|
"region": s3_config.get("region"),
|
|
"prefix": s3_config.get("prefix"),
|
|
"use_iam_role": bool(s3_config.get("use_iam_role")),
|
|
}
|
|
return db_config
|
|
|
|
|
|
def _s3_dataset_name(s3_dataset: Any) -> Optional[str]:
|
|
if not isinstance(s3_dataset, dict):
|
|
return None
|
|
bucket = s3_dataset.get("bucket")
|
|
if not bucket:
|
|
return None
|
|
prefix = s3_dataset.get("prefix")
|
|
return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}"
|
|
|
|
|
|
def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
|
|
"""Remove only HF Trainer ``tmp-checkpoint-<step>/`` partials after a cancel.
|
|
|
|
Completed ``checkpoint-<int>/`` dirs survive. Symlinked output_dir / children
|
|
are skipped so containment can't be bypassed.
|
|
"""
|
|
out = Path(output_dir)
|
|
if not out.exists() or not out.is_dir() or out.is_symlink():
|
|
return
|
|
try:
|
|
out_real = out.resolve()
|
|
out_root_real = Path(outputs_root()).resolve()
|
|
except OSError:
|
|
return
|
|
try:
|
|
out_real.relative_to(out_root_real)
|
|
except ValueError:
|
|
logger.warning(
|
|
"Skipping checkpoint cleanup - %s is not under outputs_root %s",
|
|
out_real,
|
|
out_root_real,
|
|
)
|
|
return
|
|
removed = 0
|
|
for entry in out.iterdir():
|
|
if not entry.is_dir() or entry.is_symlink():
|
|
continue
|
|
if not _HF_TMP_CHECKPOINT_RE.match(entry.name):
|
|
continue
|
|
try:
|
|
shutil.rmtree(entry, ignore_errors = False)
|
|
removed += 1
|
|
except OSError as exc:
|
|
logger.warning("Could not remove %s: %s", entry, exc)
|
|
logger.info(
|
|
"Cancelled-run cleanup removed %d in-flight tmp-checkpoint dir(s) under %s",
|
|
removed,
|
|
out,
|
|
)
|
|
|
|
|
|
_CTX = mp.get_context("spawn")
|
|
|
|
# Plot styling constants
|
|
PLOT_WIDTH = 8
|
|
PLOT_HEIGHT = 3.5
|
|
|
|
|
|
@dataclass
|
|
class TrainingProgress:
|
|
"""Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules."""
|
|
|
|
epoch: float = 0
|
|
step: int = 0
|
|
total_steps: int = 0
|
|
loss: Optional[float] = None
|
|
learning_rate: Optional[float] = None
|
|
is_training: bool = False
|
|
is_completed: bool = False
|
|
error: Optional[str] = None
|
|
status_message: str = "Ready to train"
|
|
elapsed_seconds: Optional[float] = None
|
|
eta_seconds: Optional[float] = None
|
|
grad_norm: Optional[float] = None
|
|
num_tokens: Optional[int] = None
|
|
eval_loss: Optional[float] = None
|
|
peak_memory_gb: Optional[float] = None
|
|
|
|
|
|
class TrainingBackend:
|
|
"""
|
|
Training orchestration backend — subprocess-based.
|
|
Launches a fresh subprocess per job, communicates via mp.Queue.
|
|
"""
|
|
|
|
FLUSH_THRESHOLD: int = 10
|
|
|
|
def __init__(self):
|
|
# Subprocess state
|
|
self._proc: Optional[mp.Process] = None
|
|
self._event_queue: Any = None
|
|
self._stop_queue: Any = None
|
|
self._pump_thread: Optional[threading.Thread] = None
|
|
self._lock = threading.Lock()
|
|
|
|
# Progress state (updated by pump thread from subprocess events)
|
|
self._progress = TrainingProgress()
|
|
self._should_stop = False
|
|
self._cancel_requested = False # True only for stop(save=False)
|
|
|
|
# Training metrics (consumed by routes for SSE and /metrics)
|
|
self.loss_history: list = []
|
|
self.lr_history: list = []
|
|
self.step_history: list = []
|
|
self.grad_norm_history: list = []
|
|
self.grad_norm_step_history: list = []
|
|
self.eval_loss_history: list = []
|
|
self.eval_step_history: list = []
|
|
self.eval_enabled: bool = False
|
|
self.current_theme: str = "light"
|
|
|
|
# Job metadata
|
|
self.current_job_id: Optional[str] = None
|
|
self._output_dir: Optional[str] = None
|
|
|
|
# DB persistence
|
|
self._metric_buffer: list[dict] = []
|
|
self._run_finalized: bool = False
|
|
self._db_run_created: bool = False
|
|
self._db_total_steps_set: bool = False
|
|
self._db_config: Optional[dict] = None
|
|
self._db_started_at: Optional[str] = None
|
|
|
|
logger.info("TrainingBackend initialized (subprocess mode)")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public API (called by routes/training.py)
|
|
# ------------------------------------------------------------------
|
|
|
|
def start_training(self, job_id: str, **kwargs) -> bool:
|
|
"""Spawn a subprocess to run the full training pipeline.
|
|
|
|
All kwargs are serialized into a config dict and sent to the worker.
|
|
Returns True if the subprocess started successfully.
|
|
"""
|
|
with self._lock:
|
|
if self._proc is not None and self._proc.is_alive():
|
|
logger.warning("Training subprocess already running")
|
|
return False
|
|
|
|
# Join prior pump thread — refuse to start if it won't die
|
|
if self._pump_thread is not None and self._pump_thread.is_alive():
|
|
self._pump_thread.join(timeout = 5.0)
|
|
if self._pump_thread.is_alive():
|
|
logger.warning("Previous pump thread did not exit within 5s — refusing to start")
|
|
return False
|
|
self._pump_thread = None
|
|
|
|
# Build config dict for the subprocess
|
|
config = {
|
|
"model_name": kwargs["model_name"],
|
|
"training_type": kwargs.get("training_type", "LoRA/QLoRA"),
|
|
"hf_token": kwargs.get("hf_token", ""),
|
|
"load_in_4bit": kwargs.get("load_in_4bit", True),
|
|
"max_seq_length": kwargs.get("max_seq_length", 2048),
|
|
"vision_image_size": kwargs.get("vision_image_size"),
|
|
"hf_dataset": kwargs.get("hf_dataset", ""),
|
|
"local_datasets": kwargs.get("local_datasets"),
|
|
"local_eval_datasets": kwargs.get("local_eval_datasets"),
|
|
"format_type": kwargs.get("format_type", ""),
|
|
"subset": kwargs.get("subset"),
|
|
"train_split": kwargs.get("train_split", "train"),
|
|
"eval_split": kwargs.get("eval_split"),
|
|
"eval_steps": kwargs.get("eval_steps", 0.00),
|
|
"dataset_slice_start": kwargs.get("dataset_slice_start"),
|
|
"dataset_slice_end": kwargs.get("dataset_slice_end"),
|
|
"custom_format_mapping": kwargs.get("custom_format_mapping"),
|
|
"is_dataset_image": kwargs.get("is_dataset_image", False),
|
|
"is_dataset_audio": kwargs.get("is_dataset_audio", False),
|
|
"is_embedding": kwargs.get("is_embedding", False),
|
|
"num_epochs": kwargs.get("num_epochs", 3),
|
|
"learning_rate": kwargs.get("learning_rate", "2e-4"),
|
|
"embedding_learning_rate": kwargs.get("embedding_learning_rate"),
|
|
"batch_size": kwargs.get("batch_size", 2),
|
|
"gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4),
|
|
"warmup_steps": kwargs.get("warmup_steps"),
|
|
"warmup_ratio": kwargs.get("warmup_ratio"),
|
|
"max_steps": kwargs.get("max_steps", 0),
|
|
"save_steps": kwargs.get("save_steps", 0),
|
|
"weight_decay": kwargs.get("weight_decay", 0.001),
|
|
"max_grad_norm": kwargs.get("max_grad_norm", 0.0),
|
|
"max_grad_value": _coerce_optional_nonneg_float(
|
|
"max_grad_value", kwargs.get("max_grad_value")
|
|
),
|
|
"max_grad_leaf_norm": _coerce_optional_nonneg_float(
|
|
"max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm")
|
|
),
|
|
"cast_norm_output_to_input_dtype": _coerce_optional_bool(
|
|
kwargs.get("cast_norm_output_to_input_dtype"), True
|
|
),
|
|
# MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises).
|
|
"random_seed": _coerce_seed(kwargs.get("random_seed")),
|
|
"packing": kwargs.get("packing", False),
|
|
"optim": kwargs.get("optim", "adamw_8bit"),
|
|
"lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"),
|
|
"use_lora": kwargs.get("use_lora", True),
|
|
"lora_r": kwargs.get("lora_r", 16),
|
|
"lora_alpha": kwargs.get("lora_alpha", 16),
|
|
"lora_dropout": kwargs.get("lora_dropout", 0.0),
|
|
"target_modules": kwargs.get("target_modules"),
|
|
"gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"),
|
|
"use_rslora": kwargs.get("use_rslora", False),
|
|
"use_loftq": kwargs.get("use_loftq", False),
|
|
"train_on_completions": kwargs.get("train_on_completions", False),
|
|
"finetune_vision_layers": kwargs.get("finetune_vision_layers", True),
|
|
"finetune_language_layers": kwargs.get("finetune_language_layers", True),
|
|
"finetune_attention_modules": kwargs.get("finetune_attention_modules", True),
|
|
"finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True),
|
|
"enable_wandb": kwargs.get("enable_wandb", False),
|
|
"wandb_token": kwargs.get("wandb_token"),
|
|
"wandb_project": kwargs.get("wandb_project", "unsloth-training"),
|
|
"enable_tensorboard": kwargs.get("enable_tensorboard", False),
|
|
"tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
|
|
"resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
|
|
"trust_remote_code": kwargs.get("trust_remote_code", False),
|
|
"gpu_ids": kwargs.get("gpu_ids"),
|
|
"s3_config": kwargs.get("s3_config"),
|
|
}
|
|
|
|
# Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request.
|
|
if config["training_type"] == "Full Finetuning":
|
|
config["load_in_4bit"] = False
|
|
|
|
# Spawn into locals so state is untouched on failure.
|
|
from utils.hardware import hardware as _hw
|
|
|
|
if _hw.DEVICE == _hw.DeviceType.MLX:
|
|
config["resolved_gpu_ids"] = None
|
|
config["gpu_selection"] = None
|
|
else:
|
|
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
|
|
kwargs.get("gpu_ids"),
|
|
model_name = config["model_name"],
|
|
hf_token = config["hf_token"] or None,
|
|
training_type = config["training_type"],
|
|
load_in_4bit = config["load_in_4bit"],
|
|
batch_size = config.get("batch_size", 4),
|
|
max_seq_length = config.get("max_seq_length", 2048),
|
|
lora_rank = config.get("lora_r", 16),
|
|
target_modules = config.get("target_modules"),
|
|
gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
|
|
optimizer = config.get("optim", "adamw_8bit"),
|
|
)
|
|
config["resolved_gpu_ids"] = resolved_gpu_ids
|
|
config["gpu_selection"] = gpu_selection
|
|
|
|
from .worker import run_training_process
|
|
|
|
try:
|
|
with native_path_secret_removed_for_child_start():
|
|
event_queue = _CTX.Queue()
|
|
stop_queue = _CTX.Queue()
|
|
|
|
proc = _CTX.Process(
|
|
target = run_without_native_path_secret,
|
|
args = (run_training_process,),
|
|
kwargs = {
|
|
"event_queue": event_queue,
|
|
"stop_queue": stop_queue,
|
|
"config": config,
|
|
},
|
|
daemon = True,
|
|
)
|
|
proc.start()
|
|
except Exception:
|
|
logger.error("Failed to start training subprocess", exc_info = True)
|
|
return False
|
|
|
|
logger.info("Training subprocess started (pid=%s)", proc.pid)
|
|
|
|
# Reset state (old pump thread dead, proc.start() succeeded).
|
|
self.current_job_id = job_id
|
|
self._should_stop = False
|
|
self._cancel_requested = False
|
|
self._progress = TrainingProgress(
|
|
is_training = True, status_message = "Initializing training..."
|
|
)
|
|
self.loss_history.clear()
|
|
self.lr_history.clear()
|
|
self.step_history.clear()
|
|
self.grad_norm_history.clear()
|
|
self.grad_norm_step_history.clear()
|
|
self.eval_loss_history.clear()
|
|
self.eval_step_history.clear()
|
|
self.eval_enabled = False
|
|
self._output_dir = None
|
|
self._metric_buffer.clear()
|
|
self._run_finalized = False
|
|
self._db_run_created = False
|
|
self._db_total_steps_set = False
|
|
self._db_config = _sanitize_db_config(config)
|
|
self._db_started_at = datetime.now(timezone.utc).isoformat()
|
|
|
|
# Assign subprocess handles after state reset.
|
|
self._event_queue = event_queue
|
|
self._stop_queue = stop_queue
|
|
self._proc = proc
|
|
|
|
# Eagerly create DB run row so it appears in history during model loading.
|
|
self._ensure_db_run_created()
|
|
|
|
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
|
|
self._pump_thread.start()
|
|
|
|
return True
|
|
|
|
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..."
|
|
)
|
|
return True
|
|
|
|
def force_terminate(self) -> None:
|
|
"""Force-kill the training subprocess so state can be reset immediately."""
|
|
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
|
|
cancelled = self._cancel_requested
|
|
output_dir = self._output_dir
|
|
|
|
if proc is not None:
|
|
proc.join(timeout = 5.0)
|
|
if proc.is_alive():
|
|
proc.kill()
|
|
proc.join(timeout = 2.0)
|
|
|
|
# Wait for pump thread to finish DB finalization (8s covers SQLite's 5s lock timeout).
|
|
if self._pump_thread is not None and self._pump_thread.is_alive():
|
|
self._pump_thread.join(timeout = 8.0)
|
|
|
|
if cancelled and output_dir:
|
|
try:
|
|
_cleanup_cancelled_checkpoints(output_dir)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to clean up cancelled-run checkpoints under %s",
|
|
output_dir,
|
|
)
|
|
|
|
def is_training_active(self) -> bool:
|
|
"""Check if training is currently active."""
|
|
with self._lock:
|
|
if self._proc is not None and self._proc.is_alive():
|
|
return True
|
|
|
|
if self._should_stop:
|
|
return False
|
|
|
|
p = self._progress
|
|
if p.is_training:
|
|
return True
|
|
if p.is_completed or p.error:
|
|
return False
|
|
|
|
# Infer activity from the status message.
|
|
status_lower = (p.status_message or "").lower()
|
|
if any(
|
|
k in status_lower
|
|
for k in [
|
|
"cancelled",
|
|
"canceled",
|
|
"stopped",
|
|
"completed",
|
|
"ready to train",
|
|
]
|
|
):
|
|
return False
|
|
if any(
|
|
k in status_lower
|
|
for k in [
|
|
"loading",
|
|
"preparing",
|
|
"training",
|
|
"configuring",
|
|
"tokenizing",
|
|
"starting",
|
|
"importing",
|
|
]
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
def get_training_status(self, theme: str = "light") -> Tuple:
|
|
"""Get current training status and loss plot."""
|
|
with self._lock:
|
|
progress = self._progress
|
|
|
|
if not (progress.is_training or progress.is_completed or progress.error):
|
|
return (None, progress)
|
|
|
|
plot = self._create_loss_plot(progress, theme)
|
|
return (plot, progress)
|
|
|
|
def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]:
|
|
"""Refresh plot with new theme."""
|
|
if theme and isinstance(theme, str) and theme in ["light", "dark"]:
|
|
self.current_theme = theme
|
|
if self.loss_history:
|
|
with self._lock:
|
|
progress = self._progress
|
|
return self._create_loss_plot(progress, self.current_theme)
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Compatibility shims — routes/training.py accesses these
|
|
# ------------------------------------------------------------------
|
|
|
|
class _TrainerShim:
|
|
"""Minimal shim so routes that access backend.trainer.* still work."""
|
|
|
|
def __init__(self, backend: "TrainingBackend"):
|
|
self._backend = backend
|
|
self.should_stop = False
|
|
|
|
@property
|
|
def training_progress(self):
|
|
return self._backend._progress
|
|
|
|
@training_progress.setter
|
|
def training_progress(self, value):
|
|
self._backend._progress = value
|
|
|
|
def get_training_progress(self):
|
|
return self._backend._progress
|
|
|
|
def _update_progress(self, **kwargs):
|
|
with self._backend._lock:
|
|
for key, value in kwargs.items():
|
|
if hasattr(self._backend._progress, key):
|
|
setattr(self._backend._progress, key, value)
|
|
|
|
@property
|
|
def trainer(self):
|
|
"""Compatibility shim for routes that access backend.trainer.*"""
|
|
return self._TrainerShim(self)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Event pump (background thread)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _pump_loop(self) -> None:
|
|
"""Background thread: consume events from subprocess → update state."""
|
|
while True:
|
|
if self._proc is None or self._event_queue is None:
|
|
return
|
|
|
|
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
|
|
if event is not None:
|
|
self._handle_event(event)
|
|
continue
|
|
|
|
if self._proc.is_alive():
|
|
continue
|
|
|
|
# Process exited — drain remaining events.
|
|
for e in self._drain_queue(self._event_queue):
|
|
self._handle_event(e)
|
|
|
|
# Mark done if no explicit complete/error was received.
|
|
with self._lock:
|
|
if self._progress.is_training:
|
|
if self._should_stop:
|
|
self._progress.is_training = False
|
|
self._progress.status_message = "Training stopped."
|
|
else:
|
|
self._progress.is_training = False
|
|
self._progress.error = (
|
|
self._progress.error or "Training process exited unexpectedly"
|
|
)
|
|
|
|
self._ensure_db_run_created()
|
|
self._finalize_run_in_db(
|
|
status = "stopped" if self._should_stop else "error",
|
|
error_message = None
|
|
if self._should_stop
|
|
else "Training process terminated unexpectedly",
|
|
)
|
|
return
|
|
|
|
def _handle_event(self, event: dict) -> None:
|
|
"""Apply a subprocess event to local state.
|
|
|
|
State updates happen inside self._lock; DB I/O happens after releasing
|
|
it so status-polling endpoints aren't blocked by slow SQLite writes.
|
|
"""
|
|
etype = event.get("type")
|
|
db_action: Optional[str] = None
|
|
db_action_kwargs: dict = {}
|
|
|
|
with self._lock:
|
|
if etype == "progress":
|
|
self._progress.step = event.get("step", self._progress.step)
|
|
self._progress.epoch = event.get("epoch", self._progress.epoch)
|
|
# loss/lr sanitized below.
|
|
_raw_loss = event.get("loss")
|
|
_raw_lr = event.get("learning_rate")
|
|
try:
|
|
_safe_loss = float(_raw_loss) if _raw_loss is not None else None
|
|
except (TypeError, ValueError):
|
|
logger.debug("Could not convert loss to float: %s", _raw_loss)
|
|
_safe_loss = None
|
|
_loss_is_nonfinite = _safe_loss is not None and not math.isfinite(_safe_loss)
|
|
if _loss_is_nonfinite:
|
|
# Drop the value rather than laundering it back to the last
|
|
# finite loss; clients see loss=None at this step so the NaN
|
|
# is not hidden behind a stale value. Training continues.
|
|
_safe_loss = None
|
|
if not getattr(self._progress, "_nonfinite_loss_warned", False):
|
|
self._progress._nonfinite_loss_warned = True
|
|
logger.warning(
|
|
"Training produced non-finite loss at step %s; "
|
|
"loss field will report null until it recovers.",
|
|
event.get("step", "?"),
|
|
)
|
|
try:
|
|
_safe_lr = float(_raw_lr) if _raw_lr is not None else None
|
|
except (TypeError, ValueError):
|
|
logger.debug("Could not convert learning_rate to float: %s", _raw_lr)
|
|
_safe_lr = None
|
|
if _safe_lr is not None and not math.isfinite(_safe_lr):
|
|
_safe_lr = None
|
|
if _safe_loss is not None:
|
|
self._progress.loss = _safe_loss
|
|
elif _loss_is_nonfinite:
|
|
# Clear stale finite loss so the API doesn't keep
|
|
# reporting the last good value while NaN is happening.
|
|
self._progress.loss = None
|
|
if _safe_lr is not None:
|
|
self._progress.learning_rate = _safe_lr
|
|
self._progress.total_steps = event.get("total_steps", self._progress.total_steps)
|
|
self._progress.elapsed_seconds = event.get("elapsed_seconds")
|
|
self._progress.eta_seconds = event.get("eta_seconds")
|
|
self._progress.grad_norm = event.get("grad_norm")
|
|
self._progress.num_tokens = event.get("num_tokens")
|
|
self._progress.eval_loss = event.get("eval_loss")
|
|
_peak = event.get("peak_memory_gb")
|
|
if _peak is not None:
|
|
try:
|
|
self._progress.peak_memory_gb = float(_peak)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
self._progress.is_training = True
|
|
status = event.get("status_message", "")
|
|
if status:
|
|
self._progress.status_message = status
|
|
|
|
# Update metric histories using sanitized values.
|
|
step = event.get("step", 0)
|
|
loss = _safe_loss
|
|
lr = _safe_lr
|
|
if step > 0 and loss is not None:
|
|
self.loss_history.append(loss)
|
|
self.lr_history.append(lr if lr is not None else 0.0)
|
|
self.step_history.append(step)
|
|
|
|
grad_norm = event.get("grad_norm")
|
|
gn = None
|
|
if grad_norm is not None:
|
|
try:
|
|
gn = float(grad_norm)
|
|
except (TypeError, ValueError):
|
|
gn = None
|
|
if step > 0 and gn is not None and math.isfinite(gn):
|
|
self.grad_norm_history.append(gn)
|
|
self.grad_norm_step_history.append(step)
|
|
else:
|
|
gn = None
|
|
|
|
eval_loss = event.get("eval_loss")
|
|
if eval_loss is not None:
|
|
try:
|
|
eval_loss = float(eval_loss)
|
|
except (TypeError, ValueError):
|
|
logger.debug("Could not convert eval_loss to float: %s", eval_loss)
|
|
eval_loss = None
|
|
if step > 0 and eval_loss is not None and math.isfinite(eval_loss):
|
|
self.eval_loss_history.append(eval_loss)
|
|
self.eval_step_history.append(step)
|
|
self.eval_enabled = True
|
|
else:
|
|
eval_loss = None
|
|
|
|
# Buffer metric for DB flush.
|
|
self._metric_buffer.append(
|
|
{
|
|
"step": step,
|
|
"loss": loss,
|
|
"learning_rate": lr,
|
|
"grad_norm": gn,
|
|
"eval_loss": eval_loss,
|
|
"epoch": event.get("epoch"),
|
|
"num_tokens": event.get("num_tokens"),
|
|
"elapsed_seconds": event.get("elapsed_seconds"),
|
|
}
|
|
)
|
|
|
|
# Pick the DB action to run after releasing the lock.
|
|
if not self._db_run_created and self.current_job_id and self._db_config:
|
|
db_action = "create_run"
|
|
db_action_kwargs = {
|
|
"job_id": self.current_job_id,
|
|
"model_name": self._db_config["model_name"],
|
|
"dataset_name": self._db_config.get("hf_dataset")
|
|
or next(iter(self._db_config.get("local_datasets") or []), "unknown"),
|
|
"config_json": _json.dumps(self._db_config),
|
|
"started_at": self._db_started_at or datetime.now(timezone.utc).isoformat(),
|
|
"total_steps": event.get("total_steps"),
|
|
}
|
|
elif (
|
|
event.get("total_steps")
|
|
and self._db_run_created
|
|
and not self._db_total_steps_set
|
|
):
|
|
db_action = "update_total_steps"
|
|
db_action_kwargs = {
|
|
"job_id": self.current_job_id,
|
|
"total_steps": event["total_steps"],
|
|
}
|
|
elif len(self._metric_buffer) >= self.FLUSH_THRESHOLD:
|
|
db_action = "flush"
|
|
|
|
elif etype == "eval_configured":
|
|
self.eval_enabled = True
|
|
|
|
elif etype == "status":
|
|
self._progress.status_message = event.get("message", "")
|
|
self._progress.is_training = True
|
|
|
|
elif etype == "complete":
|
|
self._progress.is_training = False
|
|
self._progress.is_completed = True
|
|
self._output_dir = event.get("output_dir")
|
|
msg = event.get("status_message", "Training completed")
|
|
self._progress.status_message = msg
|
|
if not self._db_run_created and self.current_job_id and self._db_config:
|
|
db_action = "create_and_finalize"
|
|
else:
|
|
db_action = "finalize"
|
|
db_action_kwargs = {
|
|
"status": "stopped" if self._should_stop else "completed",
|
|
"output_dir": self._output_dir,
|
|
}
|
|
|
|
elif etype == "error":
|
|
self._progress.is_training = False
|
|
self._progress.error = event.get("error", "Unknown error")
|
|
logger.error("Training error: %s", event.get("error"))
|
|
stack = event.get("stack", "")
|
|
if stack:
|
|
logger.error("Stack trace:\n%s", stack)
|
|
if not self._db_run_created and self.current_job_id and self._db_config:
|
|
db_action = "create_and_finalize"
|
|
else:
|
|
db_action = "finalize"
|
|
db_action_kwargs = {
|
|
"status": "stopped" if self._should_stop else "error",
|
|
"error_message": event.get("error", "Unknown error"),
|
|
}
|
|
|
|
# --- 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
|
|
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)
|
|
elif db_action == "create_and_finalize":
|
|
self._ensure_db_run_created()
|
|
self._finalize_run_in_db(**db_action_kwargs)
|
|
elif db_action == "update_total_steps":
|
|
try:
|
|
from storage.studio_db import update_run_total_steps
|
|
update_run_total_steps(db_action_kwargs["job_id"], db_action_kwargs["total_steps"])
|
|
self._db_total_steps_set = True
|
|
except Exception:
|
|
logger.warning("Failed to update total_steps in DB", exc_info = True)
|
|
elif db_action == "flush":
|
|
self._flush_metrics_to_db()
|
|
elif db_action == "finalize":
|
|
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
|
|
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"))
|
|
or "unknown"
|
|
)
|
|
create_run(
|
|
id = self.current_job_id,
|
|
model_name = self._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,
|
|
)
|
|
self._db_run_created = True
|
|
except Exception:
|
|
logger.warning("Failed to create DB run record for early failure", exc_info = True)
|
|
|
|
def _finalize_run_in_db(
|
|
self,
|
|
status: str,
|
|
error_message: Optional[str] = None,
|
|
output_dir: 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()
|
|
try:
|
|
from storage.studio_db import finish_run
|
|
from utils.downsample import downsample
|
|
|
|
sparkline = downsample(self.loss_history, 50)
|
|
finish_run(
|
|
id = self.current_job_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,
|
|
loss_sparkline = _json.dumps(sparkline),
|
|
output_dir = output_dir,
|
|
error_message = error_message,
|
|
)
|
|
self._run_finalized = True
|
|
except Exception:
|
|
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)
|
|
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,
|
|
)
|
|
except Exception:
|
|
# Leave buffer intact for retry on next flush
|
|
logger.warning("Failed to flush metrics to DB", exc_info = True)
|
|
|
|
@staticmethod
|
|
def _read_queue(q: Any, timeout_sec: float) -> Optional[dict]:
|
|
try:
|
|
return q.get(timeout = timeout_sec)
|
|
except queue.Empty:
|
|
return None
|
|
except (EOFError, OSError, ValueError):
|
|
return None
|
|
|
|
@staticmethod
|
|
def _drain_queue(q: Any) -> list:
|
|
events = []
|
|
while True:
|
|
try:
|
|
events.append(q.get_nowait())
|
|
except queue.Empty:
|
|
return events
|
|
except (EOFError, OSError, ValueError):
|
|
return events
|
|
|
|
# ------------------------------------------------------------------
|
|
# Plot generation
|
|
# ------------------------------------------------------------------
|
|
|
|
def _create_loss_plot(
|
|
self,
|
|
progress: TrainingProgress,
|
|
theme: str = "light",
|
|
) -> plt.Figure:
|
|
"""Create training loss plot with theme-aware styling."""
|
|
plt.close("all")
|
|
|
|
LIGHT_STYLE = {
|
|
"facecolor": "#ffffff",
|
|
"grid_color": "#d1d5db",
|
|
"line": "#16b88a",
|
|
"text": "#1f2937",
|
|
"empty_text": "#6b7280",
|
|
}
|
|
DARK_STYLE = {
|
|
"facecolor": "#292929",
|
|
"grid_color": "#404040",
|
|
"line": "#4ade80",
|
|
"text": "#e5e7eb",
|
|
"empty_text": "#9ca3af",
|
|
}
|
|
|
|
style = LIGHT_STYLE if theme == "light" else DARK_STYLE
|
|
|
|
fig, ax = plt.subplots(figsize = (PLOT_WIDTH, PLOT_HEIGHT))
|
|
fig.patch.set_facecolor(style["facecolor"])
|
|
ax.set_facecolor(style["facecolor"])
|
|
|
|
if self.loss_history:
|
|
steps = self.step_history
|
|
losses = self.loss_history
|
|
scatter_color = "#60a5fa"
|
|
ax.scatter(
|
|
steps,
|
|
losses,
|
|
s = 16,
|
|
alpha = 0.6,
|
|
color = scatter_color,
|
|
linewidths = 0,
|
|
label = "Training Loss (raw)",
|
|
)
|
|
|
|
MA_WINDOW = 20
|
|
window = min(MA_WINDOW, len(losses))
|
|
|
|
if window >= 2:
|
|
cumsum = [0.0]
|
|
for v in losses:
|
|
cumsum.append(cumsum[-1] + float(v))
|
|
|
|
ma = []
|
|
for i in range(len(losses)):
|
|
start = max(0, i - window + 1)
|
|
denom = i - start + 1
|
|
ma.append((cumsum[i + 1] - cumsum[start]) / denom)
|
|
|
|
ax.plot(
|
|
steps,
|
|
ma,
|
|
color = style["line"],
|
|
linewidth = 2.5,
|
|
alpha = 0.95,
|
|
label = f"Moving Avg ({ma[-1]:.4f})",
|
|
)
|
|
|
|
leg = ax.legend(frameon = False, fontsize = 9)
|
|
for t in leg.get_texts():
|
|
t.set_color(style["text"])
|
|
|
|
ax.set_xlabel("Steps", fontsize = 10, color = style["text"])
|
|
ax.set_ylabel("Loss", fontsize = 10, color = style["text"])
|
|
|
|
if progress.error:
|
|
title = f"Error: {progress.error}"
|
|
elif progress.is_completed:
|
|
loss_str = f"{progress.loss:.4f}" if progress.loss is not None else "--"
|
|
title = f"Training completed! Final loss: {loss_str}"
|
|
elif progress.status_message:
|
|
title = progress.status_message
|
|
elif progress.step > 0:
|
|
loss_str = f"{progress.loss:.4f}" if progress.loss is not None else "--"
|
|
title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {loss_str}"
|
|
else:
|
|
title = "Training Loss"
|
|
|
|
ax.set_title(title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"])
|
|
ax.grid(True, alpha = 0.4, linestyle = "--", color = style["grid_color"])
|
|
ax.tick_params(colors = style["text"], which = "both")
|
|
ax.spines["top"].set_visible(False)
|
|
ax.spines["right"].set_visible(False)
|
|
ax.spines["bottom"].set_color(style["text"])
|
|
ax.spines["left"].set_color(style["text"])
|
|
else:
|
|
display_msg = (
|
|
progress.status_message
|
|
if progress.status_message
|
|
else "Waiting for training data..."
|
|
)
|
|
ax.text(
|
|
0.5,
|
|
0.5,
|
|
display_msg,
|
|
ha = "center",
|
|
va = "center",
|
|
fontsize = 16,
|
|
color = style["empty_text"],
|
|
transform = ax.transAxes,
|
|
)
|
|
ax.set_xticks([])
|
|
ax.set_yticks([])
|
|
for spine in ax.spines.values():
|
|
spine.set_visible(False)
|
|
|
|
fig.tight_layout()
|
|
return fig
|
|
|
|
def _transfer_to_inference_backend(self) -> bool:
|
|
"""Transfer model to inference backend.
|
|
|
|
No-op: with subprocess training the model is freed on exit, so inference
|
|
must load from the saved checkpoint on disk.
|
|
"""
|
|
logger.info(
|
|
"_transfer_to_inference_backend: subprocess training — "
|
|
"model must be loaded from disk (output_dir=%s)",
|
|
self._output_dir,
|
|
)
|
|
return False
|
|
|
|
|
|
# ========== GLOBAL INSTANCE ==========
|
|
_training_backend = None
|
|
|
|
|
|
def get_training_backend() -> TrainingBackend:
|
|
"""Get global training backend instance"""
|
|
global _training_backend
|
|
if _training_backend is None:
|
|
_training_backend = TrainingBackend()
|
|
return _training_backend
|