""" Training API routes """ import sys from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import StreamingResponse 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 backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) # Import backend functions try: from core.training import get_training_backend except ImportError: # Fallback: try to import from parent directory parent_backend = backend_path.parent / "backend" if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from core.training import get_training_backend # Auth from auth.authentication import get_current_subject from models import ( TrainingStartRequest, TrainingJobResponse, TrainingStatus, TrainingProgress, ) from models.responses import TrainingStopResponse, TrainingMetricsResponse from pydantic import BaseModel as PydanticBaseModel class TrainingStopRequest(PydanticBaseModel): save: bool = True router = APIRouter() logger = logging.getLogger(__name__) # Configure logger if not logger.handlers: handler = logging.StreamHandler() handler.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') handler.setFormatter(formatter) logger.addHandler(handler) logger.setLevel(logging.INFO) @router.get("/hardware") async def get_hardware_utilization( current_subject: str = Depends(get_current_subject), ): """ Get a live snapshot of GPU hardware utilization. Designed to be polled by the frontend during training. Returns GPU utilization %, temperature, VRAM usage, and power draw via nvidia-smi for maximum accuracy. """ from utils.hardware import get_gpu_utilization return get_gpu_utilization() @router.post("/start") async def start_training( request: TrainingStartRequest, current_subject: str = Depends(get_current_subject), ): """ Start a training job. This endpoint initiates training in the background and returns immediately. Use the /status endpoint to check training progress. """ try: logger.info(f"Starting training job with model: {request.model_name}") backend = get_training_backend() # Generate job ID and attach to backend for later status/progress calls job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}" backend.current_job_id = job_id # Check if training is already active if backend.is_training_active(): existing_job_id: Optional[str] = getattr(backend, "current_job_id", "") return TrainingJobResponse( job_id=existing_job_id or job_id, status="error", message=( "Training is already in progress. " "Stop current training before starting a new one." ), error="Training already active", ) # Validate dataset paths if provided if request.local_datasets: validated_datasets = [] # Get the backend directory (where this file is located) backend_dir = Path(__file__).parent.parent assets_datasets_dir = backend_dir / "assets" / "datasets" for dataset_path in request.local_datasets: dataset_file = Path(dataset_path) # If not absolute, try multiple locations if not dataset_file.is_absolute(): # First try: relative to current working directory candidate = Path.cwd() / dataset_path if not candidate.exists(): # Second try: relative to assets/datasets folder candidate = assets_datasets_dir / dataset_path if not candidate.exists(): # Third try: just the filename in assets/datasets candidate = assets_datasets_dir / dataset_file.name dataset_file = candidate if not dataset_file.exists(): logger.warning( f"Dataset file not found: {dataset_path} (resolved: {dataset_file})" ) else: logger.info(f"Found dataset file: {dataset_file}") validated_datasets.append(str(dataset_file)) request.local_datasets = validated_datasets # Convert request to kwargs for backend training_kwargs = { "model_name": request.model_name, "training_type": request.training_type, "hf_token": request.hf_token or "", "load_in_4bit": request.load_in_4bit, "max_seq_length": request.max_seq_length, "hf_dataset": request.hf_dataset or "", "local_datasets": request.local_datasets, "format_type": request.format_type, "subset": request.subset, "train_split": request.train_split, "eval_split": request.eval_split, "eval_steps": request.eval_steps, "custom_format_mapping": request.custom_format_mapping, "num_epochs": request.num_epochs, "learning_rate": request.learning_rate, "batch_size": request.batch_size, "gradient_accumulation_steps": request.gradient_accumulation_steps, "warmup_steps": request.warmup_steps, "warmup_ratio": request.warmup_ratio, "max_steps": request.max_steps, "save_steps": request.save_steps, "weight_decay": request.weight_decay, "random_seed": request.random_seed, "packing": request.packing, "optim": request.optim, "lr_scheduler_type": request.lr_scheduler_type, "use_lora": request.use_lora, "lora_r": request.lora_r, "lora_alpha": request.lora_alpha, "lora_dropout": request.lora_dropout, "target_modules": request.target_modules if request.target_modules else None, "gradient_checkpointing": request.gradient_checkpointing.strip() if request.gradient_checkpointing and request.gradient_checkpointing.strip() else "unsloth", "use_rslora": request.use_rslora, "use_loftq": request.use_loftq, "train_on_completions": request.train_on_completions, "finetune_vision_layers": request.finetune_vision_layers, "finetune_language_layers": request.finetune_language_layers, "finetune_attention_modules": request.finetune_attention_modules, "finetune_mlp_modules": request.finetune_mlp_modules, "is_dataset_multimodal": request.is_dataset_multimodal, "enable_wandb": request.enable_wandb, "wandb_token": request.wandb_token or "", "wandb_project": request.wandb_project or "", "enable_tensorboard": request.enable_tensorboard, "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 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}") return TrainingJobResponse( job_id=job_id, status="error", message=( "Training thread failed to start. " "Check server logs for details." ), error="Thread not alive", ) return TrainingJobResponse( job_id=job_id, status="queued", message="Training job queued and starting in background", error=None, ) except Exception as e: logger.error(f"Error starting training: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to start training: {str(e)}", ) @router.post("/stop", response_model=TrainingStopResponse) async def stop_training( body: TrainingStopRequest = TrainingStopRequest(), current_subject: str = Depends(get_current_subject), ): """ Stop the currently running training job. Body: save (bool): If True (default), save the model at the current checkpoint. """ 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), ) if not is_active and not thread_alive: return TrainingStopResponse( status="idle", message="No training job is currently running" ) # Call backend stop method backend.stop_training(save=body.save) return TrainingStopResponse( status="stopped", message="Stop requested. Training will stop at the next safe step." ) except Exception as e: logger.error(f"Error stopping training: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to stop training: {str(e)}" ) @router.post("/reset") async def reset_training( current_subject: str = Depends(get_current_subject), ): """ Reset training state so the user can return to configuration. """ 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), ) 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.loss_history = [] backend.lr_history = [] backend.step_history = [] backend.grad_norm_history = [] backend.grad_norm_step_history = [] return {"status": "ok"} except HTTPException: raise except Exception as e: logger.error(f"Error resetting training: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to reset training: {str(e)}", ) @router.get("/status") async def get_training_status( current_subject: str = Depends(get_current_subject), ): """ Get the current training status. """ try: backend = get_training_backend() job_id: str = getattr(backend, "current_job_id", "") # 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() except Exception: progress = None status_message = ( getattr(progress, "status_message", None) if progress else None ) or "Ready to train" 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) # Derive high-level phase if error_message: phase = "error" elif is_active: msg_lower = status_message.lower() if "loading" in msg_lower: phase = "loading_model" elif any( k in msg_lower for k in ["preparing", "initializing", "configuring"] ): phase = "configuring" else: phase = "training" elif trainer_stopped: phase = "stopped" elif progress and getattr(progress, "is_completed", False): phase = "completed" elif has_thread: phase = "loading_model" else: phase = "idle" details = None if progress: details = { "epoch": getattr(progress, "epoch", 0), "step": getattr(progress, "step", 0), "total_steps": getattr(progress, "total_steps", 0), "loss": getattr(progress, "loss", 0.0), "learning_rate": getattr(progress, "learning_rate", 0.0), } # Build metric history for chart recovery after SSE reconnection metric_history = None if backend.step_history: metric_history = { "steps": list(backend.step_history), "loss": list(backend.loss_history), "lr": list(backend.lr_history), "grad_norm": list(getattr(backend, "grad_norm_history", [])), "grad_norm_steps": list(getattr(backend, "grad_norm_step_history", [])), "eval_loss": list(backend.eval_loss_history), "eval_steps": list(backend.eval_step_history), } return TrainingStatus( job_id=job_id, phase=phase, is_training_running=is_active, eval_enabled=backend.eval_enabled, message=status_message, error=error_message, details=details, metric_history=metric_history, ) except Exception as e: logger.error(f"Error getting training status: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to get training status: {str(e)}" ) @router.get("/metrics", response_model=TrainingMetricsResponse) async def get_training_metrics( current_subject: str = Depends(get_current_subject), ): """ Get training metrics (loss, learning rate, steps). """ try: backend = get_training_backend() # Get metrics from backend loss_history = backend.loss_history lr_history = backend.lr_history step_history = backend.step_history grad_norm_history = getattr(backend, "grad_norm_history", []) grad_norm_step_history = getattr(backend, "grad_norm_step_history", []) # Get current values current_loss = loss_history[-1] if loss_history else None current_lr = lr_history[-1] if lr_history else None current_step = step_history[-1] if step_history else None return TrainingMetricsResponse( loss_history=loss_history, lr_history=lr_history, step_history=step_history, grad_norm_history=grad_norm_history, grad_norm_step_history=grad_norm_step_history, current_loss=current_loss, current_lr=current_lr, current_step=current_step, ) except Exception as e: logger.error(f"Error getting training metrics: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to get training metrics: {str(e)}" ) @router.get("/progress") async def stream_training_progress( request: Request, current_subject: str = Depends(get_current_subject), ): """ Stream training progress updates using Server-Sent Events (SSE). This endpoint provides real-time updates on training progress. Supports reconnection via the SSE spec: - Sends `id:` with each event so the browser tracks position. - Sends `retry:` to control reconnection interval. - Sends named `event:` types (progress, heartbeat, complete, error). - Reads `Last-Event-ID` header on reconnect to replay missed steps. """ # Read Last-Event-ID header for reconnection resume last_event_id = request.headers.get("last-event-id") resume_from_step: Optional[int] = None if last_event_id is not None: try: resume_from_step = int(last_event_id) logger.info(f"SSE reconnect: resuming from step {resume_from_step}") except ValueError: logger.warning(f"Invalid Last-Event-ID: {last_event_id}") async def event_generator(): backend = get_training_backend() job_id: str = getattr(backend, "current_job_id", "") # ── Helpers ────────────────────────────────────────────── def build_progress( step: int, loss: float, learning_rate: float, total_steps: int, epoch: Optional[float] = None, progress: Optional[Any] = None, grad_norm_override: Optional[float] = None, eval_loss_override: Optional[float] = None, ) -> TrainingProgress: total = max(total_steps, 0) if step < 0 or total == 0: progress_percent = 0.0 else: progress_percent = ( float(step) / float(total) * 100.0 if total > 0 else 0.0 ) # Get actual values from progress object if available elapsed_seconds = getattr(progress, 'elapsed_seconds', None) if progress else None eta_seconds = getattr(progress, 'eta_seconds', None) if progress else None grad_norm = grad_norm_override if grad_norm is None and progress: grad_norm = getattr(progress, 'grad_norm', None) num_tokens = getattr(progress, 'num_tokens', None) if progress else None eval_loss = eval_loss_override if eval_loss is None and progress: eval_loss = getattr(progress, 'eval_loss', None) return TrainingProgress( job_id=job_id, step=step, total_steps=total, loss=loss, learning_rate=learning_rate, progress_percent=progress_percent, epoch=epoch, elapsed_seconds=elapsed_seconds, eta_seconds=eta_seconds, grad_norm=grad_norm, num_tokens=num_tokens, eval_loss=eval_loss, ) def format_sse( data: str, event: str = "progress", event_id: Optional[int] = None, ) -> str: """Format a single SSE message with id/event/data fields.""" lines = [] if event_id is not None: lines.append(f"id: {event_id}") lines.append(f"event: {event}") lines.append(f"data: {data}") lines.append("") # trailing blank line lines.append("") # double newline terminates the event return "\n".join(lines) # ── Retry directive ────────────────────────────────────── # Tell the browser to reconnect after 3 seconds if the connection drops yield "retry: 3000\n\n" # ── Replay missed steps on reconnect ───────────────────── if resume_from_step is not None and backend.step_history: replayed = 0 grad_norm_by_step = { step_val: grad_val for step_val, grad_val in zip( getattr(backend, "grad_norm_step_history", []), getattr(backend, "grad_norm_history", []), ) } for i, step_val in enumerate(backend.step_history): if step_val > resume_from_step: loss_val = backend.loss_history[i] if i < len(backend.loss_history) else 0.0 lr_val = backend.lr_history[i] if i < len(backend.lr_history) else 0.0 tp_replay = getattr( getattr(backend, "trainer", None), "training_progress", None ) total_replay = getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None payload = build_progress( step_val, loss_val, lr_val, total_replay, epoch_replay, progress=tp_replay, grad_norm_override=grad_norm_by_step.get(step_val), ) yield format_sse(payload.model_dump_json(), event="progress", event_id=step_val) replayed += 1 if replayed: logger.info(f"SSE reconnect: replayed {replayed} missed steps") # ── Initial status (only on fresh connections) ─────────── if resume_from_step is None: is_active = backend.is_training_active() tp = getattr(getattr(backend, "trainer", None), "training_progress", None) initial_total_steps = getattr(tp, "total_steps", 0) if tp else 0 initial_epoch = getattr(tp, "epoch", None) if tp else None initial_progress = build_progress( step=0, loss=0.0, learning_rate=0.0, total_steps=initial_total_steps, epoch=initial_epoch, progress=tp, ) yield format_sse(initial_progress.model_dump_json(), event="progress", event_id=0) # If not active, send final state and exit if not is_active: if backend.step_history: final_step = backend.step_history[-1] final_loss = backend.loss_history[-1] if backend.loss_history else 0.0 final_lr = backend.lr_history[-1] if backend.lr_history else 0.0 final_total_steps = ( getattr(tp, "total_steps", final_step) if tp else final_step ) final_epoch = getattr(tp, "epoch", None) if tp else None payload = build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch, progress=tp) yield format_sse(payload.model_dump_json(), event="complete", event_id=final_step) else: yield format_sse( build_progress(-1, 0.0, 0.0, 0, progress=tp).model_dump_json(), event="complete", event_id=0, ) return # ── Live polling loop ──────────────────────────────────── last_step = resume_from_step if resume_from_step is not None else -1 no_update_count = 0 max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation) while backend.is_training_active(): try: if backend.step_history: current_step = backend.step_history[-1] current_loss = backend.loss_history[-1] if backend.loss_history else 0.0 current_lr = backend.lr_history[-1] if backend.lr_history else 0.0 tp_inner = getattr( getattr(backend, "trainer", None), "training_progress", None ) current_total_steps = ( getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step ) current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None # Only send if step changed if current_step != last_step: progress_payload = build_progress( current_step, current_loss, current_lr, current_total_steps, current_epoch, progress=tp_inner, ) yield format_sse( progress_payload.model_dump_json(), event="progress", event_id=current_step, ) last_step = current_step no_update_count = 0 else: no_update_count += 1 # Send heartbeat every 10 seconds if no_update_count % 10 == 0: heartbeat_payload = build_progress( current_step, current_loss, current_lr, current_total_steps, current_epoch, progress=tp_inner, ) yield format_sse( heartbeat_payload.model_dump_json(), event="heartbeat", event_id=current_step, ) else: # No steps yet, but training is active (model loading, etc.) no_update_count += 1 if no_update_count % 5 == 0: # Pull total_steps and status from trainer so # the frontend can show "Tokenizing…" etc. tp_prep = getattr( getattr(backend, "trainer", None), "training_progress", None, ) prep_total = ( getattr(tp_prep, "total_steps", 0) if tp_prep else 0 ) preparing_payload = build_progress( 0, 0.0, 0.0, prep_total, progress=tp_prep, ) yield format_sse( preparing_payload.model_dump_json(), event="heartbeat", event_id=0, ) # Timeout check if no_update_count > max_no_updates: logger.warning("Progress stream timeout - no updates received") tp_timeout = getattr(getattr(backend, "trainer", None), "training_progress", None) timeout_payload = build_progress(last_step, 0.0, 0.0, 0, progress=tp_timeout) yield format_sse( timeout_payload.model_dump_json(), event="error", event_id=last_step if last_step >= 0 else 0, ) break await asyncio.sleep(1) # Poll every second except Exception as e: logger.error(f"Error in progress stream: {e}", exc_info=True) tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None) error_payload = build_progress(0, 0.0, 0.0, 0, progress=tp_error) yield format_sse( error_payload.model_dump_json(), event="error", event_id=last_step if last_step >= 0 else 0, ) break # ── Final "complete" event ─────────────────────────────── final_step = backend.step_history[-1] if backend.step_history else last_step final_loss = backend.loss_history[-1] if backend.loss_history else 0.0 final_lr = backend.lr_history[-1] if backend.lr_history else 0.0 final_tp = getattr( getattr(backend, "trainer", None), "training_progress", None ) final_total_steps = ( getattr(final_tp, "total_steps", final_step) if final_tp else final_step ) final_epoch = getattr(final_tp, "epoch", None) if final_tp else None final_payload = build_progress( final_step, final_loss, final_lr, final_total_steps, final_epoch, progress=final_tp, ) yield format_sse( final_payload.model_dump_json(), event="complete", event_id=final_step if final_step >= 0 else 0, ) return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", } )