843 lines
33 KiB
Python
843 lines
33 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 API routes
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from typing import Optional, Any
|
|
from loggers import get_logger
|
|
import asyncio
|
|
from datetime import datetime
|
|
import uuid as _uuid
|
|
|
|
# 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
|
|
from utils.models.model_config import load_model_defaults
|
|
from utils.paths import resolve_dataset_path
|
|
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
|
|
from utils.models.model_config import load_model_defaults
|
|
from utils.paths import resolve_dataset_path
|
|
|
|
# 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 = get_logger(__name__)
|
|
|
|
|
|
def _validate_local_dataset_paths(
|
|
paths: list[str], label: str = "Local dataset"
|
|
) -> list[str]:
|
|
"""Resolve and validate a list of local dataset paths. Returns validated absolute paths."""
|
|
validated = []
|
|
missing = []
|
|
for dataset_path in paths:
|
|
dataset_file = resolve_dataset_path(dataset_path)
|
|
if not dataset_file.exists():
|
|
missing.append(f"{dataset_path} (resolved: {dataset_file})")
|
|
continue
|
|
logger.info(f"Found {label.lower()} file: {dataset_file}")
|
|
validated.append(str(dataset_file))
|
|
|
|
if missing:
|
|
missing_detail = "; ".join(missing[:3])
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"{label} not found: {missing_detail}",
|
|
)
|
|
return validated
|
|
|
|
|
|
@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}")
|
|
|
|
# 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()
|
|
|
|
# Check if training is already active (before mutating any state)
|
|
if backend.is_training_active():
|
|
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
|
|
return TrainingJobResponse(
|
|
job_id = existing_job_id or "",
|
|
status = "error",
|
|
message = (
|
|
"Training is already in progress. "
|
|
"Stop current training before starting a new one."
|
|
),
|
|
error = "Training already active",
|
|
)
|
|
|
|
# Generate job ID — passed into start_training() which sets it on the
|
|
# backend only after confirming the old pump thread is dead.
|
|
job_id = (
|
|
f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}"
|
|
)
|
|
|
|
# Validate dataset paths if provided
|
|
if request.local_datasets:
|
|
request.local_datasets = _validate_local_dataset_paths(
|
|
request.local_datasets, "Local dataset"
|
|
)
|
|
if request.local_eval_datasets and request.eval_steps > 0:
|
|
request.local_eval_datasets = _validate_local_dataset_paths(
|
|
request.local_eval_datasets, "Local eval dataset"
|
|
)
|
|
|
|
# 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,
|
|
"local_eval_datasets": request.local_eval_datasets,
|
|
"format_type": request.format_type,
|
|
"subset": request.subset,
|
|
"train_split": request.train_split,
|
|
"eval_split": request.eval_split,
|
|
"eval_steps": request.eval_steps,
|
|
"dataset_slice_start": request.dataset_slice_start,
|
|
"dataset_slice_end": request.dataset_slice_end,
|
|
"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_image": request.is_dataset_image,
|
|
"is_dataset_audio": request.is_dataset_audio,
|
|
"is_embedding": request.is_embedding,
|
|
"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 "",
|
|
"trust_remote_code": request.trust_remote_code,
|
|
}
|
|
|
|
# Training page has no trust_remote_code toggle — the value comes from
|
|
# YAML model defaults applied when the user selects a model. As a safety
|
|
# net, consult the YAML directly so models that need it always get it.
|
|
if not training_kwargs["trust_remote_code"]:
|
|
model_defaults = load_model_defaults(request.model_name)
|
|
yaml_trust = model_defaults.get("training", {}).get(
|
|
"trust_remote_code", False
|
|
)
|
|
if yaml_trust:
|
|
logger.info(
|
|
f"YAML config sets trust_remote_code=True for {request.model_name}"
|
|
)
|
|
training_kwargs["trust_remote_code"] = True
|
|
|
|
# Free GPU memory: shut down any running inference/export subprocesses
|
|
# before training starts (they'd compete for VRAM otherwise)
|
|
try:
|
|
from core.inference import get_inference_backend
|
|
|
|
inf_backend = get_inference_backend()
|
|
if inf_backend.active_model_name:
|
|
logger.info(
|
|
"Unloading inference model '%s' to free GPU memory for training",
|
|
inf_backend.active_model_name,
|
|
)
|
|
inf_backend._shutdown_subprocess()
|
|
inf_backend.active_model_name = None
|
|
inf_backend.models.clear()
|
|
except Exception as e:
|
|
logger.warning("Could not unload inference model: %s", e)
|
|
|
|
try:
|
|
from core.export import get_export_backend
|
|
|
|
exp_backend = get_export_backend()
|
|
if exp_backend.current_checkpoint:
|
|
logger.info(
|
|
"Shutting down export subprocess to free GPU memory for training"
|
|
)
|
|
exp_backend._shutdown_subprocess()
|
|
exp_backend.current_checkpoint = None
|
|
exp_backend.is_vision = False
|
|
exp_backend.is_peft = False
|
|
except Exception as e:
|
|
logger.warning("Could not shut down export subprocess: %s", e)
|
|
|
|
# start_training now spawns a subprocess (non-blocking)
|
|
success = backend.start_training(job_id = job_id, **training_kwargs)
|
|
|
|
if not success:
|
|
progress_error = backend.trainer.training_progress.error
|
|
return TrainingJobResponse(
|
|
job_id = backend.current_job_id or "",
|
|
status = "error",
|
|
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 subprocess",
|
|
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()
|
|
is_active = backend.is_training_active()
|
|
logger.info("Stop requested: save=%s is_active=%s", body.save, is_active)
|
|
|
|
if not is_active:
|
|
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()
|
|
is_active = backend.is_training_active()
|
|
|
|
if is_active:
|
|
if backend._cancel_requested:
|
|
# Cancel (save=False) was requested — force-terminate so we can reset immediately
|
|
logger.info(
|
|
"Force-terminating subprocess for immediate reset (cancel path)"
|
|
)
|
|
backend.force_terminate()
|
|
else:
|
|
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._should_stop = False # Clear stop flag so status returns to idle
|
|
backend.trainer._update_progress(
|
|
is_training = False,
|
|
is_completed = False,
|
|
error = None,
|
|
status_message = "Ready to train",
|
|
step = 0,
|
|
loss = None,
|
|
epoch = 0,
|
|
total_steps = 0,
|
|
)
|
|
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", "") or ""
|
|
|
|
# Check if training is active
|
|
is_active = backend.is_training_active()
|
|
|
|
# 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, "_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 or "importing" 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"
|
|
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", None),
|
|
"learning_rate": getattr(progress, "learning_rate", None),
|
|
}
|
|
|
|
# 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", "") or ""
|
|
|
|
# ── Helpers ──────────────────────────────────────────────
|
|
def build_progress(
|
|
step: int,
|
|
loss: Optional[float],
|
|
learning_rate: Optional[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 None
|
|
)
|
|
lr_val = (
|
|
backend.lr_history[i] if i < len(backend.lr_history) else None
|
|
)
|
|
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 = None,
|
|
learning_rate = None,
|
|
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 None
|
|
)
|
|
final_lr = backend.lr_history[-1] if backend.lr_history else None
|
|
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, None, None, 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 None
|
|
)
|
|
current_lr = backend.lr_history[-1] if backend.lr_history else None
|
|
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,
|
|
None,
|
|
None,
|
|
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, None, None, 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, None, None, 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 None
|
|
final_lr = backend.lr_history[-1] if backend.lr_history else None
|
|
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",
|
|
},
|
|
)
|