807 lines
32 KiB
Python
807 lines
32 KiB
Python
"""
|
|
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",
|
|
}
|
|
)
|