diff --git a/setup.ps1 b/setup.ps1 index 8a2ec56033..0b28456b44 100644 --- a/setup.ps1 +++ b/setup.ps1 @@ -691,6 +691,22 @@ python "$PSScriptRoot\install_python_stack.py" # Restore ErrorActionPreference after pip/python work $ErrorActionPreference = $prevEAP +# ── Pre-install transformers 5.x into .venv_t5/ ── +# Models like GLM-4.7-Flash need transformers>=5.1.0. Instead of pip-installing +# at runtime (slow, ~10-15s), we pre-install into a separate directory. +# The training subprocess just prepends .venv_t5/ to sys.path — instant switch. +Write-Host "" +Write-Host " Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan +$VenvT5Dir = Join-Path $PSScriptRoot ".venv_t5" +if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir } +New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null +$prevEAP_t5 = $ErrorActionPreference +$ErrorActionPreference = "Continue" +pip install --target $VenvT5Dir --no-deps "transformers==5.1.0" 2>&1 | Out-Null +pip install --target $VenvT5Dir --no-deps "huggingface_hub>=1.3.0" 2>&1 | Out-Null +$ErrorActionPreference = $prevEAP_t5 +Write-Host "[OK] Transformers 5.x pre-installed to .venv_t5/" -ForegroundColor Green + # ========================================================================== # PHASE 3.5: Install OpenSSL dev (for HTTPS support in llama-server) # ========================================================================== diff --git a/setup.sh b/setup.sh index f1622ed7a7..2c2179acb0 100755 --- a/setup.sh +++ b/setup.sh @@ -180,10 +180,23 @@ else # Local: create venv (always start fresh to preserve correct install order) rm -rf .venv rm -rf .venv_overlay # Clean up stale transformers version overlay + rm -rf .venv_t5 # Will be rebuilt below "$BEST_PY" -m venv .venv source .venv/bin/activate install_python_stack - + + # ── 6b. Pre-install transformers 5.x into .venv_t5/ ── + # Models like GLM-4.7-Flash need transformers>=5.1.0. Instead of pip-installing + # at runtime (slow, ~10-15s), we pre-install into a separate directory. + # The training subprocess just prepends .venv_t5/ to sys.path — instant switch. + echo "" + echo " Pre-installing transformers 5.x for newer model support..." + VENV_T5_DIR="$SCRIPT_DIR/.venv_t5" + mkdir -p "$VENV_T5_DIR" + run_quiet "pip install transformers 5.x" pip install --target "$VENV_T5_DIR" --no-deps "transformers==5.1.0" + run_quiet "pip install huggingface_hub for t5" pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub>=1.3.0" + echo "✅ Transformers 5.x pre-installed to .venv_t5/" + # ── 7. WSL: pre-install GGUF build dependencies ── # On WSL, sudo requires a password and can't be entered during GGUF export # (runs in a non-interactive subprocess). Install build deps here instead. diff --git a/studio/backend/core/training/__init__.py b/studio/backend/core/training/__init__.py index 8fc2a6c721..0e0299caf6 100644 --- a/studio/backend/core/training/__init__.py +++ b/studio/backend/core/training/__init__.py @@ -1,12 +1,9 @@ """ Training submodule - Training backends and trainer classes """ -from .trainer import UnslothTrainer, get_trainer, TrainingProgress -from .training import TrainingBackend, get_training_backend +from .training import TrainingBackend, TrainingProgress, get_training_backend __all__ = [ - 'UnslothTrainer', - 'get_trainer', 'TrainingProgress', 'TrainingBackend', 'get_training_backend', diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 153f4335e3..cf2238ba30 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -1,620 +1,526 @@ """ -Training backend for FastAPI integration -""" -import matplotlib.pyplot as plt -from typing import Any, Generator, Tuple -import logging -import math +Training backend — subprocess orchestrator. -from .trainer import get_trainer, TrainingProgress -from utils.hardware import clear_gpu_cache +Each training job runs in a fresh subprocess (mp.get_context("spawn")), +solving the transformers version-switching problem. The old in-process +UnslothTrainer singleton is only used inside the subprocess (worker.py). + +This file orchestrates the subprocess lifecycle, pumps events from the +worker's mp.Queue, and exposes the same API surface to routes/training.py. + +Pattern follows core/data_recipe/jobs/manager.py. +""" +import math +import multiprocessing as mp +import queue +import threading +import time +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Tuple, Any + +import matplotlib.pyplot as plt logger = logging.getLogger(__name__) +_CTX = mp.get_context("spawn") + # Plot styling constants -PLOT_WIDTH = 8 # Inches -PLOT_HEIGHT = 3.5 # Inches +PLOT_WIDTH = 8 +PLOT_HEIGHT = 3.5 + + +@dataclass +class TrainingProgress: + """Mirror of trainer.TrainingProgress — kept here so the parent process + never needs to import the heavy ML modules.""" + epoch: float = 0 + step: int = 0 + total_steps: int = 0 + loss: float = 0.0 + learning_rate: float = 0.0 + 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 class TrainingBackend: """ - Training orchestration backend. - Handles both text and vision models, LoRA and full finetuning. + Training orchestration backend — subprocess-based. + Launches a fresh subprocess per training job, communicates via mp.Queue. """ def __init__(self): - self.trainer = get_trainer() + # 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() - # Training Metrics - self.loss_history = [] - self.lr_history = [] - self.step_history = [] - self.grad_norm_history = [] - self.grad_norm_step_history = [] - self.eval_loss_history = [] - self.eval_step_history = [] + # Progress state (updated by pump thread from subprocess events) + self._progress = TrainingProgress() + self._should_stop = 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 + + logger.info("TrainingBackend initialized (subprocess mode)") + + # ------------------------------------------------------------------ + # Public API (called by routes/training.py) + # ------------------------------------------------------------------ + + def start_training(self, **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 was started successfully. + """ + with self._lock: + if self._proc is not None and self._proc.is_alive(): + logger.warning("Training subprocess already running") + return False + + # Reset state + self._should_stop = 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.current_theme = "light" + self._output_dir = None - self.trainer.add_progress_callback(self._on_progress_update) + # Resolve project root (studio/backend/core/training/ → project root) + project_root = str(Path(__file__).resolve().parent.parent.parent.parent.parent) - logger.info("TrainingBackend initialized") + # Build config dict for the subprocess + config = { + "project_root": project_root, + "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), + "hf_dataset": kwargs.get("hf_dataset", ""), + "local_datasets": kwargs.get("local_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_multimodal": kwargs.get("is_dataset_multimodal", False), + "num_epochs": kwargs.get("num_epochs", 3), + "learning_rate": kwargs.get("learning_rate", "2e-4"), + "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.01), + "random_seed": kwargs.get("random_seed", 3407), + "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"), + } - def _on_progress_update(self, progress: TrainingProgress): - """Callback for progress updates""" - if progress.step >= 0 and progress.loss > 0: - self.loss_history.append(progress.loss) - self.lr_history.append(progress.learning_rate) - self.step_history.append(progress.step) - if progress.step >= 0 and progress.grad_norm is not None: - try: - grad_norm = float(progress.grad_norm) - except (TypeError, ValueError): - grad_norm = None - if grad_norm is not None and math.isfinite(grad_norm): - self.grad_norm_history.append(grad_norm) - self.grad_norm_step_history.append(progress.step) - if progress.eval_loss is not None: - self.eval_loss_history.append(progress.eval_loss) - self.eval_step_history.append(progress.step) + # Derive load_in_4bit from training_type + if config["training_type"] != "LoRA/QLoRA": + config["load_in_4bit"] = False - def start_training(self, - # Model parameters - model_name: str, - training_type: str, # NEW: "LoRA/QLoRA" or "Full Finetuning" - hf_token: str, - load_in_4bit: bool, - max_seq_length: int, + # Spawn subprocess + from .worker import run_training_process - # Dataset parameters - hf_dataset: str, - local_datasets: list, - format_type: str, # CHANGED: was data_template + self._event_queue = _CTX.Queue() + self._stop_queue = _CTX.Queue() - # Training parameters - num_epochs: int, - learning_rate: str, - batch_size: int, - gradient_accumulation_steps: int, - warmup_steps: int, # May be None even without default - warmup_ratio: float, # May be None even without default - max_steps: int, - save_steps: int, - weight_decay: float, - random_seed: int, - packing: bool, - optim: str, - lr_scheduler_type: str, + self._proc = _CTX.Process( + target=run_training_process, + kwargs={ + "event_queue": self._event_queue, + "stop_queue": self._stop_queue, + "config": config, + }, + daemon=True, + ) + self._proc.start() + logger.info("Training subprocess started (pid=%s)", self._proc.pid) - # LoRA parameters - use_lora: bool, # Should be derived from training_type - lora_r: int, - lora_alpha: int, - lora_dropout: float, - target_modules: list, - gradient_checkpointing: str, - use_rslora: bool, - use_loftq: bool, - train_on_completions: bool, + # Start event pump thread + self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True) + self._pump_thread.start() - # NEW: Vision-specific LoRA parameters - finetune_vision_layers: bool, - finetune_language_layers: bool, - finetune_attention_modules: bool, - finetune_mlp_modules: bool, - - # Logging parameters - enable_wandb: bool, - wandb_token: str, - wandb_project: str, - enable_tensorboard: bool, - tensorboard_dir: str, - - # Optional parameters - custom_format_mapping: dict = None, - subset: str = None, - train_split: str = "train", - eval_split: str = None, - eval_steps: float = 0.00, - is_dataset_multimodal: bool = False, - dataset_slice_start: int = None, - dataset_slice_end: int = None) -> bool: - """ - Start training. - - Returns: - True if training started successfully, False otherwise. - """ - try: - # Wait for any previous training thread to finish - old_thread = getattr(self.trainer, "training_thread", None) - if old_thread and old_thread.is_alive(): - logger.info("Waiting for previous training thread to finish...") - old_thread.join(timeout=30) - - # Explicitly free old SFTTrainer and CUDA resources before loading new model. - # Without this, forked multiprocessing workers (num_proc tokenization) inherit - # stale CUDA state from the previous run, causing extreme slowdowns or crashes. - if self.trainer.trainer is not None: - logger.info("Cleaning up previous SFTTrainer...") - self.trainer.trainer = None - if self.trainer.model is not None: - self.trainer.model = None - if self.trainer.tokenizer is not None: - self.trainer.tokenizer = None - # Flush all pending async CUDA ops so forked tokenization processes - # don't inherit stale async state that causes pool join to hang. - import torch as _torch - if _torch.cuda.is_available(): - _torch.cuda.synchronize() - import gc - gc.collect() - clear_gpu_cache() - - # Reset stop flag and clear history - self.trainer.should_stop = False - self.trainer.save_on_stop = True - self.loss_history = [] - self.lr_history = [] - self.step_history = [] - self.grad_norm_history = [] - self.grad_norm_step_history = [] - self.eval_loss_history = [] - self.eval_step_history = [] - self.eval_enabled = False - import time - output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}" - - # Derive use_lora from training_type - use_lora_actual = (training_type == "LoRA/QLoRA") - if use_lora_actual: print("using Lora") - else: print("using full finetuning") - logger.info(f"Starting training - Type: {training_type}, Model: {model_name}") - - # ========== LOAD MODEL ========== - logger.info("Loading model...") - success = self.trainer.load_model( - model_name=model_name, - max_seq_length=max_seq_length, - load_in_4bit=load_in_4bit if use_lora_actual else False, # Only 4bit for LoRA - hf_token=hf_token if hf_token.strip() else None, - is_dataset_multimodal=is_dataset_multimodal, - ) - - if not success or self.trainer.should_stop: - logger.error("Failed to load model or stopped by user") - return False - - # ========== PREPARE MODEL FOR TRAINING ========== - if use_lora_actual: - logger.info("Preparing model with LoRA...") - success = self.trainer.prepare_model_for_training( - use_lora=True, - # Vision-specific parameters - finetune_vision_layers=finetune_vision_layers, - finetune_language_layers=finetune_language_layers, - finetune_attention_modules=finetune_attention_modules, - finetune_mlp_modules=finetune_mlp_modules, - # Standard LoRA parameters - target_modules=target_modules, - lora_r=lora_r, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - use_gradient_checkpointing=gradient_checkpointing, - use_rslora=use_rslora, - use_loftq=use_loftq - ) - else: - logger.info("Preparing model for full finetuning...") - success = self.trainer.prepare_model_for_training( - use_lora=False # Full finetuning - ) - - if not success or self.trainer.should_stop: - logger.error("Failed to prepare model or stopped by user") - return False - - # ========== LOAD DATASET ========== - logger.info("Loading dataset...") - #breakpoint() - dataset_result = self.trainer.load_and_format_dataset( - dataset_source=hf_dataset if hf_dataset.strip() else None, - format_type=format_type, - local_datasets=local_datasets if local_datasets else None, - custom_format_mapping=custom_format_mapping, - subset=subset, - train_split=train_split, - eval_split=eval_split, - eval_steps=eval_steps, - dataset_slice_start=dataset_slice_start, - dataset_slice_end=dataset_slice_end, - ) - - # Unpack: load_and_format_dataset returns (dataset, eval_dataset) - if isinstance(dataset_result, tuple): - dataset, eval_dataset = dataset_result - else: - dataset = dataset_result - eval_dataset = None - - # Track whether eval is enabled for status reporting - self.eval_enabled = eval_dataset is not None - - if dataset is None or self.trainer.should_stop: - logger.error("Failed to load dataset or stopped by user") - return False - - # ========== START TRAINING ========== - # Convert learning rate string to float - try: - lr_value = float(learning_rate) - except ValueError: - logger.error(f"Invalid learning rate: {learning_rate}") - self.trainer._update_progress( - error=f"Invalid learning rate: {learning_rate}", - is_training=False - ) - return - - logger.info("Starting training worker thread...") - success = self.trainer.start_training( - dataset=dataset, - eval_dataset=eval_dataset, - eval_steps=eval_steps, - output_dir=output_dir, - num_epochs=num_epochs, - learning_rate=lr_value, - batch_size=batch_size, - gradient_accumulation_steps=gradient_accumulation_steps, - warmup_steps=warmup_steps, - warmup_ratio=warmup_ratio, - max_steps=max_steps if max_steps > 0 else 0, - save_steps=save_steps if save_steps > 0 else 0, - weight_decay=weight_decay, - random_seed=random_seed, - packing=packing, - train_on_completions=train_on_completions, - enable_wandb=enable_wandb, - wandb_project=wandb_project, - wandb_token=wandb_token if wandb_token.strip() else None, - enable_tensorboard=enable_tensorboard, - tensorboard_dir=tensorboard_dir, - max_seq_length=max_seq_length, - optim=optim, - lr_scheduler_type=lr_scheduler_type, - ) - - if not success: - logger.error("Failed to start training") - return False - - return True - - except Exception as e: - logger.error(f"Error in start_training: {e}", exc_info=True) - self.trainer._update_progress( - error=str(e), - is_training=False - ) - return False + return True def stop_training(self, save: bool = True) -> bool: - """ - Stop ongoing training. + """Send stop signal to the training subprocess.""" + self._should_stop = 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 - Args: - save: If True, save the model at the current checkpoint. + def is_training_active(self) -> bool: + """Check if training is currently active.""" + with self._lock: + # Subprocess alive = active + if self._proc is not None and self._proc.is_alive(): + return True + + # Stop was requested and process exited → inactive + if self._should_stop: + return False + + # Check progress state + p = self._progress + if p.is_training: + return True + if p.is_completed or p.error: + return False + + # Check status message for activity indicators + 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 - Returns: - True if training was successfully stopped. - """ - try: - logger.info(f"Stopping training (save={save})...") - self.trainer.stop_training(save=save) - return True - except Exception as e: - logger.error(f"Error stopping training: {e}") return False def get_training_status(self, theme: str = "light") -> Tuple: - """ - Get current training status and loss plot. + """Get current training status and loss plot.""" + with self._lock: + progress = self._progress - Args: - theme: "light" or "dark" for plot styling + if not (progress.is_training or progress.is_completed or progress.error): + return (None, progress) - Returns: - Tuple of (plot, progress) - """ + plot = self._create_loss_plot(progress, theme) + return (plot, progress) - try: - progress = self.trainer.get_training_progress() - - # If not training and not completed, return no updates - if not (progress.is_training or progress.is_completed or progress.error): - return (None, progress) - - # Generate plot - plot = self._create_loss_plot(progress, theme) - return (plot, progress) - - except Exception as e: - logger.error(f"Error getting training status: {e}") - return (None, None) - - def refresh_plot_for_theme(self, theme: str) -> plt.Figure: - """ - Refresh plot with new theme. - - Args: - theme: "light" or "dark" - - Returns: - Updated matplotlib figure - """ + 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 - - # Always generate plot if we have loss history if self.loss_history: - progress = self.trainer.get_training_progress() + with self._lock: + progress = self._progress return self._create_loss_plot(progress, self.current_theme) - return None - def is_training_active(self) -> bool: - """ - Check if training is currently active (from load_model start to completion/error). - - Returns: - True if training is in progress, False otherwise - """ + # ------------------------------------------------------------------ + # 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 + + # Try to read an event + event = self._read_queue(self._event_queue, timeout_sec=0.25) + if event is not None: + self._handle_event(event) + continue + + # No event — check if process is still alive + 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 as 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" + return + + def _handle_event(self, event: dict) -> None: + """Apply a subprocess event to local state.""" + etype = event.get("type") + + with self._lock: + if etype == "progress": + self._progress.step = event.get("step", self._progress.step) + self._progress.epoch = event.get("epoch", self._progress.epoch) + self._progress.loss = event.get("loss", self._progress.loss) + self._progress.learning_rate = event.get("learning_rate", self._progress.learning_rate) + 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") + self._progress.is_training = True + status = event.get("status_message", "") + if status: + self._progress.status_message = status + + # Update metric histories + step = event.get("step", 0) + loss = event.get("loss", 0.0) + lr = event.get("learning_rate", 0.0) + if step >= 0 and loss > 0: + self.loss_history.append(loss) + self.lr_history.append(lr) + self.step_history.append(step) + + grad_norm = event.get("grad_norm") + if grad_norm is not None: + try: + gn = float(grad_norm) + except (TypeError, ValueError): + gn = None + if gn is not None and math.isfinite(gn): + self.grad_norm_history.append(gn) + self.grad_norm_step_history.append(step) + + eval_loss = event.get("eval_loss") + if eval_loss is not None: + self.eval_loss_history.append(eval_loss) + self.eval_step_history.append(step) + 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 + + 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) + + @staticmethod + def _read_queue(q: Any, timeout_sec: float) -> Optional[dict]: try: - training_thread = getattr(self.trainer, "training_thread", None) - if training_thread and training_thread.is_alive(): - return True + return q.get(timeout=timeout_sec) + except queue.Empty: + return None + except (EOFError, OSError, ValueError): + return None - # Stop requested and worker already exited => inactive. - # This allows UI to show stopped state + "Back to configuration". - if getattr(self.trainer, "should_stop", False): - return False + @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 - progress = self.trainer.get_training_progress() - # Training is active if is_training is True - # Also check if we're in loading/preparation phase (status_message indicates activity) - is_active = progress.is_training - # Also consider it active if we have a status message indicating loading/preparation - # but haven't completed or errored yet - if not is_active and not progress.is_completed and not progress.error: - status = progress.status_message or "" - status_lower = status.lower() - if any( - keyword in status_lower - for keyword in ["cancelled", "canceled", "stopped", "completed", "ready to train"] - ): - return False - if any( - keyword in status_lower - for keyword in [ - "loading", - "preparing", - "training", - "configuring", - "tokenizing", - "starting", - ] - ): - is_active = True - return is_active - except Exception as e: - logger.error(f"Error checking training state: {e}") - return False + # ------------------------------------------------------------------ + # Plot generation (unchanged from original) + # ------------------------------------------------------------------ def _create_loss_plot(self, progress: TrainingProgress, theme: str = "light") -> plt.Figure: - """ - Create training loss plot with theme-aware styling. + """Create training loss plot with theme-aware styling.""" + plt.close('all') - Args: - progress: Current training progress - theme: "light" or "dark" + 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" + } - Returns: - Matplotlib figure - """ - plt.close('all') + style = LIGHT_STYLE if theme == "light" else DARK_STYLE - # Theme-specific styling - 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" - } + fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT)) + fig.patch.set_facecolor(style["facecolor"]) + ax.set_facecolor(style["facecolor"]) - style = LIGHT_STYLE if theme == "light" else DARK_STYLE + 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)") - fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT)) - fig.patch.set_facecolor(style["facecolor"]) - ax.set_facecolor(style["facecolor"]) + MA_WINDOW = 20 + window = min(MA_WINDOW, len(losses)) - if self.loss_history: - steps = self.step_history - losses = self.loss_history - scatter_color = "#60a5fa" - # Scatter plot for raw loss points - ax.scatter( - steps, - losses, - s=16, - alpha=0.6, - color=scatter_color, - linewidths=0, - label="Training Loss (raw)", - ) + if window >= 2: + cumsum = [0.0] + for v in losses: + cumsum.append(cumsum[-1] + float(v)) - # Moving average line overlay (trailing window) - MA_WINDOW = 20 # adjust smoothing aggressiveness - window = min(MA_WINDOW, len(losses)) + 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) - if window >= 2: - cumsum = [0.0] - for v in losses: - cumsum.append(cumsum[-1] + float(v)) + ax.plot(steps, ma, color=style["line"], linewidth=2.5, alpha=0.95, + label=f"Moving Avg ({ma[-1]:.4f})") - 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) + leg = ax.legend(frameon=False, fontsize=9) + for t in leg.get_texts(): + t.set_color(style["text"]) - ax.plot( - steps, - ma, - color=style["line"], - linewidth=2.5, - alpha=0.95, - label=f"Moving Avg ({ma[-1]:.4f})", - ) + ax.set_xlabel('Steps', fontsize=10, color=style["text"]) + ax.set_ylabel('Loss', fontsize=10, color=style["text"]) - 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"]) - - # Build status message for title - if progress.error: - title = f"Error: {progress.error}" - elif progress.is_completed: - title = f"Training completed! Final loss: {progress.loss:.4f}" - elif progress.status_message: - title = progress.status_message - elif progress.step > 0: - title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {progress.loss:.4f}" - else: - title = "Training Loss" - - ax.set_title(title, fontsize=11, fontweight='bold', - pad=10, color=style["text"]) - - # Style grid and spines - 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"]) + if progress.error: + title = f"Error: {progress.error}" + elif progress.is_completed: + title = f"Training completed! Final loss: {progress.loss:.4f}" + elif progress.status_message: + title = progress.status_message + elif progress.step > 0: + title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {progress.loss:.4f}" 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) + title = "Training Loss" - fig.tight_layout() - return fig + 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. + + With subprocess-based training, the model lives in the subprocess + and is freed when it exits. Inference must load from the saved + checkpoint on disk. This is a no-op placeholder. """ - Transfer the trained model to InferenceBackend. - Called automatically when training completes. - """ - print("=" * 60) - print("DEBUG: _transfer_to_inference_backend() CALLED") - print("=" * 60) - - try: - from ..inference import get_inference_backend - - session = self.current_training_session - - # Check if already transferred - if session.get('transferred', False): - print("DEBUG: Already transferred, returning True") - logger.info("Model already transferred, skipping") - return True - - # Validate session data - if not session.get('base_model_name') or not session.get('output_dir'): - logger.warning("Training session incomplete, cannot transfer") - logger.warning(f"Session data: {session}") - return False - - inference_backend = get_inference_backend() - - base_model_name = session['base_model_name'] - output_dir = session['output_dir'] - is_lora = session['is_lora'] - is_vlm = session['is_vlm'] - - logger.info(f"=" * 60) - logger.info(f"TRANSFERRING MODEL TO INFERENCE BACKEND") - logger.info(f"=" * 60) - logger.info(f" Base model: {base_model_name}") - logger.info(f" Output dir: {output_dir}") - logger.info(f" Is LoRA: {is_lora}") - logger.info(f" Is VLM: {is_vlm}") - - # Transfer the model object directly from trainer memory. - # If is_lora is True, self.trainer.model is a PeftModel (Base + Adapter). - # If is_lora is False, it is the finetuned Base Model. - inference_backend.models[base_model_name] = { - "model": self.trainer.model, - "tokenizer": self.trainer.tokenizer, - "is_vision": is_vlm, - "is_lora": is_lora, - "model_path": base_model_name, - "base_model": None, - "loaded_adapters": {}, - # Unsloth/PEFT training keeps the active adapter named 'default' in memory - "active_adapter": "default" if is_lora else None, - } - - # For vision models, also transfer processor - if is_vlm: - if hasattr(self.trainer, 'tokenizer'): - inference_backend.models[base_model_name]["processor"] = self.trainer.tokenizer - logger.info(" Transferred processor for vision model") - - # Load chat template info - inference_backend._load_chat_template_info(base_model_name) - - # If it was LoRA, register the output path. - # This ensures the Eval UI dropdown (which lists files) knows that - # the model currently in memory corresponds to this specific output directory. - if is_lora: - inference_backend.models[base_model_name]["last_trained_adapter"] = output_dir - logger.info(f"Marked trained LoRA adapter: {output_dir}") - - # Set as active model - inference_backend.active_model_name = base_model_name - logger.info(f"Set active model: {base_model_name}") - - return True - - except Exception as e: - logger.error(f"Error transferring model to inference backend: {e}") - import traceback - traceback.print_exc() - return False + 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 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py new file mode 100644 index 0000000000..de7499f2a1 --- /dev/null +++ b/studio/backend/core/training/worker.py @@ -0,0 +1,339 @@ +""" +Training subprocess entry point. + +Each training job runs in a fresh subprocess (mp.get_context("spawn")). +This gives us a clean Python interpreter with no stale module state — +solving the transformers version-switching problem completely. + +Pattern follows core/data_recipe/jobs/worker.py. +""" +from __future__ import annotations + +import logging +import os +import sys +import time +import traceback +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + + +def _activate_transformers_version(model_name: str, project_root: str) -> None: + """Activate the correct transformers version BEFORE any ML imports. + + If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ + directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). + """ + # Ensure backend is on path for utils imports + backend_path = os.path.join(project_root, "studio", "backend") + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.transformers_version import needs_transformers_5, _resolve_base_model + + resolved = _resolve_base_model(model_name) + if needs_transformers_5(resolved): + venv_t5 = os.path.join(project_root, ".venv_t5") + if os.path.isdir(venv_t5): + sys.path.insert(0, venv_t5) + logger.info("Activated transformers 5.x from %s", venv_t5) + else: + # Fallback: pip install at runtime (slower, ~10-15s) + logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5) + import subprocess as sp + os.makedirs(venv_t5, exist_ok=True) + sp.run( + [sys.executable, "-m", "pip", "install", "--target", venv_t5, + "--no-deps", "transformers==5.1.0"], + stdout=sp.PIPE, stderr=sp.STDOUT, + ) + sp.run( + [sys.executable, "-m", "pip", "install", "--target", venv_t5, + "--no-deps", "huggingface_hub>=1.3.0"], + stdout=sp.PIPE, stderr=sp.STDOUT, + ) + if os.path.isdir(venv_t5): + sys.path.insert(0, venv_t5) + else: + logger.info("Using default transformers (4.57.x) for %s", model_name) + + +def run_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, +) -> None: + """Subprocess entrypoint. Fresh Python — no stale module state. + + Args: + event_queue: mp.Queue for sending progress/status/error events to parent. + stop_queue: mp.Queue for receiving stop commands from parent. + config: Training configuration dict with all parameters. + """ + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + project_root = config["project_root"] + model_name = config["model_name"] + + # ── 1. Activate correct transformers version BEFORE any ML imports ── + try: + _activate_transformers_version(model_name, project_root) + except Exception as exc: + event_queue.put({ + "type": "error", + "error": f"Failed to activate transformers version: {exc}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 2. Now import ML libraries (fresh in this clean process) ── + try: + _send_status(event_queue, "Importing ML libraries...") + + backend_path = os.path.join(project_root, "studio", "backend") + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from core.training.trainer import UnslothTrainer, TrainingProgress + + import transformers + logger.info("Subprocess loaded transformers %s", transformers.__version__) + except Exception as exc: + event_queue.put({ + "type": "error", + "error": f"Failed to import ML libraries: {exc}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 3. Create a fresh trainer instance ── + trainer = UnslothTrainer() + + # Wire up progress callback → event_queue + def _on_progress(progress: TrainingProgress): + if progress.step >= 0 and progress.loss > 0: + event_queue.put({ + "type": "progress", + "step": progress.step, + "epoch": progress.epoch, + "loss": progress.loss, + "learning_rate": progress.learning_rate, + "total_steps": progress.total_steps, + "elapsed_seconds": progress.elapsed_seconds, + "eta_seconds": progress.eta_seconds, + "grad_norm": progress.grad_norm, + "num_tokens": progress.num_tokens, + "eval_loss": progress.eval_loss, + "status_message": progress.status_message, + "ts": time.time(), + }) + if progress.status_message: + _send_status(event_queue, progress.status_message) + + trainer.add_progress_callback(_on_progress) + + # Wire up stop_queue polling to trainer.should_stop + import threading + import queue as _queue + + def _poll_stop(): + while True: + try: + msg = stop_queue.get(timeout=1.0) + if msg and msg.get("type") == "stop": + save = msg.get("save", True) + trainer.should_stop = True + trainer.save_on_stop = save + logger.info("Stop signal received (save=%s)", save) + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target=_poll_stop, daemon=True) + stop_thread.start() + + # ── 4. Execute the training pipeline ── + try: + hf_token = config.get("hf_token", "") + hf_token = hf_token if hf_token and hf_token.strip() else None + + # Load model + _send_status(event_queue, "Loading model...") + success = trainer.load_model( + model_name=model_name, + max_seq_length=config["max_seq_length"], + load_in_4bit=config["load_in_4bit"], + hf_token=hf_token, + is_dataset_multimodal=config.get("is_dataset_multimodal", False), + ) + if not success or trainer.should_stop: + if trainer.should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + else: + event_queue.put({ + "type": "error", + "error": trainer.training_progress.error or "Failed to load model", + "stack": "", "ts": time.time(), + }) + return + + # Prepare model (LoRA or full finetuning) + training_type = config.get("training_type", "LoRA/QLoRA") + use_lora = (training_type == "LoRA/QLoRA") + if use_lora: + _send_status(event_queue, "Configuring LoRA adapters...") + success = trainer.prepare_model_for_training( + use_lora=True, + finetune_vision_layers=config.get("finetune_vision_layers", True), + finetune_language_layers=config.get("finetune_language_layers", True), + finetune_attention_modules=config.get("finetune_attention_modules", True), + finetune_mlp_modules=config.get("finetune_mlp_modules", True), + target_modules=config.get("target_modules"), + lora_r=config.get("lora_r", 16), + lora_alpha=config.get("lora_alpha", 16), + lora_dropout=config.get("lora_dropout", 0.0), + use_gradient_checkpointing=config.get("gradient_checkpointing", "unsloth"), + use_rslora=config.get("use_rslora", False), + use_loftq=config.get("use_loftq", False), + ) + else: + _send_status(event_queue, "Preparing model for full finetuning...") + success = trainer.prepare_model_for_training(use_lora=False) + + if not success or trainer.should_stop: + if trainer.should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + else: + event_queue.put({ + "type": "error", + "error": trainer.training_progress.error or "Failed to prepare model", + "stack": "", "ts": time.time(), + }) + return + + # Load dataset + _send_status(event_queue, "Loading and formatting dataset...") + hf_dataset = config.get("hf_dataset", "") + dataset_result = trainer.load_and_format_dataset( + dataset_source=hf_dataset if hf_dataset and hf_dataset.strip() else None, + format_type=config.get("format_type", ""), + local_datasets=config.get("local_datasets") or None, + custom_format_mapping=config.get("custom_format_mapping"), + subset=config.get("subset"), + train_split=config.get("train_split", "train"), + eval_split=config.get("eval_split"), + eval_steps=config.get("eval_steps", 0.00), + dataset_slice_start=config.get("dataset_slice_start"), + dataset_slice_end=config.get("dataset_slice_end"), + ) + + if isinstance(dataset_result, tuple): + dataset, eval_dataset = dataset_result + else: + dataset = dataset_result + eval_dataset = None + + # Disable eval if eval_steps <= 0 + eval_steps = config.get("eval_steps", 0.00) + if eval_steps is not None and float(eval_steps) <= 0: + eval_dataset = None + + if dataset is None or trainer.should_stop: + if trainer.should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + else: + event_queue.put({ + "type": "error", + "error": trainer.training_progress.error or "Failed to load dataset", + "stack": "", "ts": time.time(), + }) + return + + # Convert learning rate + try: + lr_value = float(config.get("learning_rate", "2e-4")) + except ValueError: + event_queue.put({ + "type": "error", + "error": f"Invalid learning rate: {config.get('learning_rate')}", + "stack": "", "ts": time.time(), + }) + return + + # Generate output dir + output_dir = config.get("output_dir") + if not output_dir: + output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}" + + # Start training (directly — no inner thread, we ARE the subprocess) + _send_status(event_queue, "Starting training...") + max_steps = config.get("max_steps", 0) + save_steps = config.get("save_steps", 0) + + trainer._train_worker( + dataset, + output_dir=output_dir, + num_epochs=config.get("num_epochs", 3), + learning_rate=lr_value, + batch_size=config.get("batch_size", 2), + gradient_accumulation_steps=config.get("gradient_accumulation_steps", 4), + warmup_steps=config.get("warmup_steps"), + warmup_ratio=config.get("warmup_ratio"), + max_steps=max_steps if max_steps and max_steps > 0 else 0, + save_steps=save_steps if save_steps and save_steps > 0 else 0, + weight_decay=config.get("weight_decay", 0.01), + random_seed=config.get("random_seed", 3407), + packing=config.get("packing", False), + train_on_completions=config.get("train_on_completions", False), + enable_wandb=config.get("enable_wandb", False), + wandb_project=config.get("wandb_project", "unsloth-training"), + wandb_token=config.get("wandb_token"), + enable_tensorboard=config.get("enable_tensorboard", False), + tensorboard_dir=config.get("tensorboard_dir", "runs"), + eval_dataset=eval_dataset, + eval_steps=eval_steps, + max_seq_length=config.get("max_seq_length", 2048), + optim=config.get("optim", "adamw_8bit"), + lr_scheduler_type=config.get("lr_scheduler_type", "linear"), + ) + + # Check final state + progress = trainer.get_training_progress() + if progress.error: + event_queue.put({ + "type": "error", + "error": progress.error, + "stack": "", + "ts": time.time(), + }) + else: + event_queue.put({ + "type": "complete", + "output_dir": output_dir, + "status_message": progress.status_message or "Training completed", + "ts": time.time(), + }) + + except Exception as exc: + event_queue.put({ + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + + +def _send_status(event_queue: Any, message: str) -> None: + """Send a status update to the parent process.""" + event_queue.put({ + "type": "status", + "message": message, + "ts": time.time(), + }) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c73cb661f9..5df2a777de 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -9,7 +9,6 @@ from typing import Dict, Optional, Any import logging import asyncio from datetime import datetime -import threading # Add backend directory to path # The backend code should be in the same directory structure @@ -86,9 +85,9 @@ async def start_training( try: logger.info(f"Starting training job with model: {request.model_name}") - # Ensure correct transformers version for this model architecture - from utils.transformers_version import ensure_transformers_version - ensure_transformers_version(request.model_name) + # NOTE: No in-process ensure_transformers_version() call here. + # The subprocess (worker.py) activates the correct version in a + # fresh Python interpreter before importing any ML libraries. backend = get_training_backend() @@ -193,84 +192,22 @@ async def start_training( "tensorboard_dir": request.tensorboard_dir or "", } - # Set initial "preparing" state - try: - backend.trainer._update_progress( - status_message="Initializing training...", - is_training=False, - ) - except Exception: - pass + # start_training now spawns a subprocess (non-blocking) + success = backend.start_training(**training_kwargs) - def run_training(): - try: - logger.info( - f"Starting training job {job_id} with model {request.model_name}" - ) - - # Update status to show we're loading model - try: - backend.trainer._update_progress(status_message="Loading model...") - except Exception as e: - logger.error(f"Error updating progress: {e}") - - # start_training returns bool (not generator) - run_result = backend.start_training(**training_kwargs) - logger.info( - "Training job %s backend.start_training returned type=%s value=%r", - job_id, - type(run_result).__name__, - run_result, - ) - if not run_result: - progress_error = backend.trainer.training_progress.error - raise RuntimeError(progress_error or "Training failed to start") - - logger.info(f"Training job {job_id} started successfully") - - except Exception as e: - logger.error(f"Training error in job {job_id}: {e}", exc_info=True) - try: - backend.trainer._update_progress( - error=str(e), - is_training=False, - ) - except Exception as update_error: - logger.error(f"Failed to update progress: {update_error}") - - # Start training in a daemon thread - training_thread = threading.Thread( - target=run_training, - daemon=True, - name=f"Training-{job_id}", - ) - training_thread.start() - - # Store thread reference for status checking - backend._training_thread = training_thread - - # Give it a moment to start - import time - - time.sleep(0.5) - - # Verify training thread is alive - if not training_thread.is_alive(): - logger.warning(f"Training thread died immediately for job {job_id}") + if not success: + progress_error = backend.trainer.training_progress.error return TrainingJobResponse( job_id=job_id, status="error", - message=( - "Training thread failed to start. " - "Check server logs for details." - ), - error="Thread not alive", + message=progress_error or "Failed to start training subprocess", + error=progress_error or "subprocess_start_failed", ) return TrainingJobResponse( job_id=job_id, status="queued", - message="Training job queued and starting in background", + message="Training job queued and starting in subprocess", error=None, ) @@ -295,18 +232,10 @@ async def stop_training( """ try: backend = get_training_backend() - trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None) - thread_alive = bool(trainer_thread and trainer_thread.is_alive()) is_active = backend.is_training_active() - logger.info( - "Stop requested: save=%s is_active=%s thread_alive=%s should_stop=%s", - body.save, - is_active, - thread_alive, - getattr(getattr(backend, "trainer", None), "should_stop", None), - ) + logger.info("Stop requested: save=%s is_active=%s", body.save, is_active) - if not is_active and not thread_alive: + if not is_active: return TrainingStopResponse( status="idle", message="No training job is currently running" @@ -337,25 +266,21 @@ async def reset_training( """ try: backend = get_training_backend() - trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None) - thread_alive = bool(trainer_thread and trainer_thread.is_alive()) is_active = backend.is_training_active() - if is_active or thread_alive: - logger.warning( - "Rejected reset while training active: is_active=%s thread_alive=%s should_stop=%s", - is_active, - thread_alive, - getattr(getattr(backend, "trainer", None), "should_stop", None), - ) + if is_active: + logger.warning("Rejected reset while training active: is_active=%s", is_active) raise HTTPException( status_code=409, detail="Training is still running. Stop training and wait for it to finish before resetting.", ) logger.info("Reset training state: clearing runtime + metric history") - backend.trainer.should_stop = False - backend.trainer.training_progress = backend.trainer.training_progress.__class__() + backend.trainer._update_progress( + is_training=False, is_completed=False, error=None, + status_message="Ready to train", step=0, loss=0.0, epoch=0, + total_steps=0, + ) backend.loss_history = [] backend.lr_history = [] backend.step_history = [] @@ -386,13 +311,6 @@ async def get_training_status( # Check if training is active is_active = backend.is_training_active() - # Check if there's a training thread running (preparation phase) - has_thread = ( - hasattr(backend, "_training_thread") - and backend._training_thread - and backend._training_thread.is_alive() - ) - # Get progress info from trainer try: progress = backend.trainer.get_training_progress() @@ -405,14 +323,14 @@ async def get_training_status( error_message = getattr(progress, "error", None) if progress else None # Check if training was stopped by user - trainer_stopped = getattr(backend.trainer, "should_stop", False) + trainer_stopped = getattr(backend, "_should_stop", False) # Derive high-level phase if error_message: phase = "error" elif is_active: msg_lower = status_message.lower() - if "loading" in msg_lower: + if "loading" in msg_lower or "importing" in msg_lower: phase = "loading_model" elif any( k in msg_lower for k in ["preparing", "initializing", "configuring"] @@ -424,8 +342,6 @@ async def get_training_status( phase = "stopped" elif progress and getattr(progress, "is_completed", False): phase = "completed" - elif has_thread: - phase = "loading_model" else: phase = "idle"