From 482e934e09c38850b7ee2b51daf0b1382f7e2022 Mon Sep 17 00:00:00 2001 From: sshah229 Date: Wed, 4 Feb 2026 05:49:29 -0700 Subject: [PATCH 1/2] Refactored the training and model routes and added the jwt authentication --- studio/backend/auth/jwt.py | 69 ++++ studio/backend/core/__init__.py | 4 +- studio/backend/models/__init__.py | 30 +- studio/backend/routes/models.py | 77 +++- studio/backend/routes/training.py | 371 ++++++++++++------ .../backend/utils/datasets/dataset_utils.py | 3 +- 6 files changed, 393 insertions(+), 161 deletions(-) create mode 100644 studio/backend/auth/jwt.py diff --git a/studio/backend/auth/jwt.py b/studio/backend/auth/jwt.py new file mode 100644 index 0000000000..2637efa821 --- /dev/null +++ b/studio/backend/auth/jwt.py @@ -0,0 +1,69 @@ +import secrets +from datetime import UTC, datetime, timedelta +from typing import Optional + +from fastapi import Depends, HTTPException, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from jose import JWTError, jwt + + +# Ephemeral in-memory secret: +# - Generated fresh on each backend process start +# - Never written to disk +# - Not configurable by the user +SECRET_KEY = secrets.token_urlsafe(64) +ALGORITHM = "HS256" +ACCESS_TOKEN_EXPIRE_MINUTES = 60 + +security = HTTPBearer() # Reads Authorization: Bearer + + +def create_access_token( + subject: str, + expires_delta: Optional[timedelta] = None, +) -> str: + """ + Create a signed JWT for the given subject (e.g. "local-user"). + + Tokens are valid only for the lifetime of this process, because the + SECRET_KEY is regenerated each time the backend restarts. + """ + to_encode = {"sub": subject} + expire = datetime.now(UTC) + ( + expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES) + ) + to_encode.update({"exp": expire}) + return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) + + +async def get_current_subject( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> str: + """ + FastAPI dependency to validate the JWT and return the subject. + + Use this as a dependency on routes that should be protected, e.g.: + + @router.get("/secure") + async def secure_endpoint(current_subject: str = Depends(get_current_subject)): + ... + """ + token = credentials.credentials + try: + payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) + subject: Optional[str] = payload.get("sub") + if subject is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid token payload", + ) + return subject + except JWTError: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or expired token", + ) +token = create_access_token("local-user") +print(token) + + diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 46ee6c14b1..79f2cdd628 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -3,10 +3,10 @@ Unified core module for Unsloth backend """ # Inference -from .inference import InferenceBackend, get_inference_backend +from .inference.inference import InferenceBackend, get_inference_backend # Training -from .training import UnslothTrainer, get_trainer, TrainingBackend, get_training_backend, create_training_handlers, TrainingProgress +from .training.training import UnslothTrainer, get_trainer, TrainingBackend, get_training_backend, create_training_handlers, TrainingProgress # Configuration (from utils) from utils.models import is_vision_model, ModelConfig, scan_trained_loras, load_model_defaults, get_base_model_from_lora diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index 07836f2168..a47da0242c 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -3,35 +3,25 @@ Pydantic models for API request/response schemas """ from .training import ( TrainingStartRequest, - TrainingStartResponse, - TrainingStatusResponse, - TrainingMetricsResponse, - TrainingProgressResponse, + TrainingJobResponse, + TrainingStatus, + TrainingProgress, ) from .models import ( - ModelSearchRequest, - ModelSearchResponse, - ModelListResponse, - ModelConfigResponse, - LoRAScanResponse, + ModelDetails, LoRAInfo, - ModelInfo, + LoRAScanResponse, ) __all__ = [ # Training schemas "TrainingStartRequest", - "TrainingStartResponse", - "TrainingStatusResponse", - "TrainingMetricsResponse", - "TrainingProgressResponse", + "TrainingJobResponse", + "TrainingStatus", + "TrainingProgress", # Model management schemas - "ModelSearchRequest", - "ModelSearchResponse", - "ModelListResponse", - "ModelConfigResponse", - "LoRAScanResponse", + "ModelDetails", "LoRAInfo", - "ModelInfo", + "LoRAScanResponse", ] diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 40db17eb04..276c69e4a6 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3,15 +3,19 @@ Model Management API routes """ import sys from pathlib import Path -from fastapi import APIRouter, HTTPException, Query -from typing import Optional +from fastapi import APIRouter, Depends, HTTPException, Query +from typing import List, Optional import logging +from pydantic import BaseModel + # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) +from auth.jwt import get_current_subject + # Import backend functions try: from utils.utils import search_hf_models @@ -38,16 +42,42 @@ except ImportError: ) from core.inference import get_inference_backend -from models.models import ( - ModelSearchRequest, - ModelSearchResponse, - ModelInfo, - ModelListResponse, - ModelConfigResponse, +from models import ( + ModelDetails, LoRAScanResponse, LoRAInfo, ) + +class ModelInfo(BaseModel): + """Basic model info used in search/list responses""" + + id: str + name: Optional[str] = None + is_vision: Optional[bool] = False + is_lora: Optional[bool] = False + + +class ModelSearchRequest(BaseModel): + """Request body for model search""" + + query: str + hf_token: Optional[str] = None + + +class ModelSearchResponse(BaseModel): + """Response schema for model search""" + + models: List[ModelInfo] + total: int + + +class ModelListResponse(BaseModel): + """Response schema for listing models""" + + models: List[ModelInfo] + default_models: List[str] + router = APIRouter() logger = logging.getLogger(__name__) @@ -62,7 +92,10 @@ if not logger.handlers: @router.post("/search") -async def search_models(request: ModelSearchRequest): +async def search_models( + request: ModelSearchRequest, + current_subject: str = Depends(get_current_subject), +): """ Search for models on HuggingFace Hub. @@ -117,7 +150,9 @@ async def search_models(request: ModelSearchRequest): @router.get("/list") -async def list_models(): +async def list_models( + current_subject: str = Depends(get_current_subject), +): """ List available models (default models and loaded models). @@ -174,7 +209,10 @@ async def list_models(): @router.get("/config/{model_name:path}") -async def get_model_config(model_name: str): +async def get_model_config( + model_name: str, + current_subject: str = Depends(get_current_subject), +): """ Get configuration for a specific model. @@ -200,12 +238,12 @@ async def get_model_config(model_name: str): # If ModelConfig creation fails, use defaults pass - return ModelConfigResponse( + return ModelDetails( model_name=model_name, config=config_dict, is_vision=is_vision, is_lora=is_lora, - base_model=base_model + base_model=base_model, ) except Exception as e: @@ -218,7 +256,8 @@ async def get_model_config(model_name: str): @router.get("/loras") async def scan_loras( - outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters") + outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters"), + current_subject: str = Depends(get_current_subject), ): """ Scan for trained LoRA adapters in the outputs directory. @@ -256,7 +295,10 @@ async def scan_loras( @router.get("/loras/{lora_path:path}/base-model") -async def get_lora_base_model(lora_path: str): +async def get_lora_base_model( + lora_path: str, + current_subject: str = Depends(get_current_subject), +): """ Get the base model for a LoRA adapter. @@ -287,7 +329,10 @@ async def get_lora_base_model(lora_path: str): @router.get("/check-vision/{model_name:path}") -async def check_vision_model(model_name: str): +async def check_vision_model( + model_name: str, + current_subject: str = Depends(get_current_subject), +): """ Check if a model is a vision model. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4eaa0643d1..8e4642eef3 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -3,9 +3,9 @@ Training API routes """ import sys from pathlib import Path -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import StreamingResponse -from typing import Dict +from typing import Dict, Optional import logging import asyncio from datetime import datetime @@ -27,12 +27,14 @@ except ImportError: sys.path.insert(0, str(parent_backend)) from core.training import get_training_backend -from models.training import ( +# Auth +from auth.jwt import get_current_subject + +from models import ( TrainingStartRequest, - TrainingStartResponse, - TrainingStatusResponse, - TrainingMetricsResponse, - TrainingProgressResponse, + TrainingJobResponse, + TrainingStatus, + TrainingProgress, ) router = APIRouter() @@ -49,35 +51,47 @@ if not logger.handlers: @router.post("/start") -async def start_training(request: TrainingStartRequest): +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(): - return TrainingStartResponse( + 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" + 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 @@ -89,14 +103,16 @@ async def start_training(request: TrainingStartRequest): # 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})") + 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, @@ -125,7 +141,9 @@ async def start_training(request: TrainingStartRequest): "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", + "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, @@ -139,84 +157,95 @@ async def start_training(request: TrainingStartRequest): "enable_tensorboard": request.enable_tensorboard, "tensorboard_dir": request.tensorboard_dir or "", } - - # Generate job ID - job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}" - + # Set initial "preparing" state try: backend.trainer._update_progress( status_message="Initializing training...", - is_training=False + is_training=False, ) - except: + except Exception: pass - + def run_training(): try: - logger.info(f"Starting training job {job_id} with model {request.model_name}") - + 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}") - + # Consume the generator - this actually runs the training update_count = 0 - for update_tuple in backend.start_training(**training_kwargs): + for _update_tuple in backend.start_training(**training_kwargs): update_count += 1 if update_count % 10 == 0: logger.info(f"Training progress update #{update_count}") - + logger.info(f"Training job {job_id} completed 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 + 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 = 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 TrainingStartResponse( + return TrainingJobResponse( + job_id=job_id, status="error", - message="Training thread failed to start. Check server logs for details.", - error="Thread not alive" + message=( + "Training thread failed to start. " + "Check server logs for details." + ), + error="Thread not alive", ) - - return TrainingStartResponse( - status="started", + + return TrainingJobResponse( job_id=job_id, - message="Training job started successfully" + 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)}" + detail=f"Failed to start training: {str(e)}", ) @router.post("/stop") -async def stop_training(): +async def stop_training( + current_subject: str = Depends(get_current_subject), +): """ Stop the currently running training job. """ @@ -246,55 +275,75 @@ async def stop_training(): @router.get("/status") -async def get_training_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 + 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() - status_message = progress.status_message or "Ready to train" - except: + except Exception: progress = None - status_message = "Unknown" - - if is_active: - # Actual training is running - trainer = backend.trainer - current_step = getattr(trainer.training_progress, 'step', None) or (progress.step if progress else None) - total_steps = getattr(trainer.training_progress, 'total_steps', None) or (progress.total_steps if progress else None) - - return TrainingStatusResponse( - status="training", - is_active=True, - message=status_message or "Training is in progress", - current_step=current_step, - total_steps=total_steps - ) - elif has_thread or (progress and status_message and any(keyword in status_message.lower() for keyword in ["loading", "preparing", "initializing"])): - # Training thread is running but not yet in active training phase - return TrainingStatusResponse( - status="preparing", - is_active=False, - message=status_message or "Preparing training...", - current_step=None, - total_steps=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 + + # 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 progress and getattr(progress, "is_completed", False): + phase = "completed" + elif has_thread: + phase = "loading_model" else: - return TrainingStatusResponse( - status="idle", - is_active=False, - message="No training job is currently running" - ) + 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), + } + + return TrainingStatus( + job_id=job_id, + phase=phase, + is_training_running=is_active, + message=status_message, + error=error_message, + details=details, + ) except Exception as e: logger.error(f"Error getting training status: {e}", exc_info=True) @@ -305,7 +354,9 @@ async def get_training_status(): @router.get("/metrics") -async def get_training_metrics(): +async def get_training_metrics( + current_subject: str = Depends(get_current_subject), +): """ Get training metrics (loss, learning rate, steps). """ @@ -316,20 +367,21 @@ async def get_training_metrics(): loss_history = backend.loss_history lr_history = backend.lr_history step_history = backend.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, - current_loss=current_loss, - current_lr=current_lr, - current_step=current_step - ) + + # Keep metrics as a simple JSON payload instead of a Pydantic model + return { + "loss_history": loss_history, + "lr_history": lr_history, + "step_history": 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) @@ -340,7 +392,9 @@ async def get_training_metrics(): @router.get("/progress") -async def stream_training_progress(): +async def stream_training_progress( + current_subject: str = Depends(get_current_subject), +): """ Stream training progress updates using Server-Sent Events (SSE). @@ -348,12 +402,53 @@ async def stream_training_progress(): """ async def event_generator(): backend = get_training_backend() - + job_id: str = getattr(backend, "current_job_id", "") + + # Helper to build a TrainingProgress payload from raw values + def build_progress( + step: int, + loss: float, + learning_rate: float, + total_steps: int, + epoch: Optional[int] = 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 + ) + + 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=None, + eta_seconds=None, + grad_norm=None, + num_tokens=None, + ) + # Send initial status is_active = backend.is_training_active() - initial_message = 'Connecting...' if is_active else 'No training in progress' - yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message=initial_message).model_dump_json()}\n\n" - + 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, + ) + yield f"data: {initial_progress.model_dump_json()}\n\n" + # If not active, check if there's any history if not is_active: if backend.step_history: @@ -361,9 +456,13 @@ async def stream_training_progress(): 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 - yield f"data: {TrainingProgressResponse(step=final_step, loss=final_loss, learning_rate=final_lr, status_message='Training completed').model_dump_json()}\n\n" + final_total_steps = ( + getattr(tp, "total_steps", final_step) if tp else final_step + ) + final_epoch = getattr(tp, "epoch", None) if tp else None + yield f"data: {build_progress(final_step, final_loss, final_lr, final_total_steps, final_epoch).model_dump_json()}\n\n" else: - yield f"data: {TrainingProgressResponse(step=-1, loss=0.0, learning_rate=0.0, status_message='No training in progress').model_dump_json()}\n\n" + yield f"data: {build_progress(-1, 0.0, 0.0, 0).model_dump_json()}\n\n" return # Poll for updates while training is active @@ -378,53 +477,81 @@ async def stream_training_progress(): 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 = TrainingProgressResponse( - step=current_step, - loss=current_loss, - learning_rate=current_lr, - status_message=f"Training step {current_step}" + progress_payload = build_progress( + current_step, + current_loss, + current_lr, + current_total_steps, + current_epoch, ) - yield f"data: {progress.model_dump_json()}\n\n" + yield f"data: {progress_payload.model_dump_json()}\n\n" last_step = current_step no_update_count = 0 else: no_update_count += 1 # Send heartbeat every 10 seconds if no_update_count % 10 == 0: - progress = TrainingProgressResponse( - step=current_step, - loss=current_loss, - learning_rate=current_lr, - status_message=f"Training step {current_step} (waiting for next update...)" + heartbeat_payload = build_progress( + current_step, + current_loss, + current_lr, + current_total_steps, + current_epoch, ) - yield f"data: {progress.model_dump_json()}\n\n" + yield f"data: {heartbeat_payload.model_dump_json()}\n\n" else: # No steps yet, but training is active no_update_count += 1 if no_update_count % 5 == 0: - yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message='Preparing training...').model_dump_json()}\n\n" + preparing_payload = build_progress(0, 0.0, 0.0, 0) + yield f"data: {preparing_payload.model_dump_json()}\n\n" # Timeout check if no_update_count > max_no_updates: logger.warning("Progress stream timeout - no updates received") - yield f"data: {TrainingProgressResponse(step=last_step, loss=0.0, learning_rate=0.0, status_message='Progress timeout - training may have stopped').model_dump_json()}\n\n" + timeout_payload = build_progress(last_step, 0.0, 0.0, 0) + yield f"data: {timeout_payload.model_dump_json()}\n\n" break await asyncio.sleep(1) # Poll every second except Exception as e: logger.error(f"Error in progress stream: {e}", exc_info=True) - yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message=f'Error: {str(e)}').model_dump_json()}\n\n" + error_payload = build_progress(0, 0.0, 0.0, 0) + yield f"data: {error_payload.model_dump_json()}\n\n" break - + # Send final status 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 - yield f"data: {TrainingProgressResponse(step=final_step, loss=final_loss, learning_rate=final_lr, status_message='Training completed').model_dump_json()}\n\n" + 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, + ) + yield f"data: {final_payload.model_dump_json()}\n\n" return StreamingResponse( event_generator(), diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index a1137ae1f6..4690ec31bd 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -37,7 +37,8 @@ from .chat_templates import ( ) from .vlm_processing import generate_smart_vlm_instruction from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator -from .model_mappings import TEMPLATE_TO_MODEL_MAPPER, RESPONSE_MARKERS +from .model_mappings import TEMPLATE_TO_MODEL_MAPPER +# , RESPONSE_MARKERS def check_dataset_format(dataset, is_vlm: bool = False) -> dict: From daa20821c110d580042d8a0cbc688f2b43e285f2 Mon Sep 17 00:00:00 2001 From: sshah229 Date: Fri, 6 Feb 2026 02:56:54 -0700 Subject: [PATCH 2/2] refactored the code for username/password and added pydantic models and routes for the same --- studio/backend/auth/__init__.py | 24 +++++++ studio/backend/auth/auth.db | Bin 0 -> 12288 bytes studio/backend/auth/hashing.py | 40 ++++++++++++ studio/backend/auth/jwt.py | 34 +++++++--- studio/backend/auth/storage.py | 101 ++++++++++++++++++++++++++++++ studio/backend/main.py | 3 +- studio/backend/models/__init__.py | 9 +++ studio/backend/models/auth.py | 22 +++++++ studio/backend/routes/__init__.py | 3 +- studio/backend/routes/auth.py | 88 ++++++++++++++++++++++++++ 10 files changed, 312 insertions(+), 12 deletions(-) create mode 100644 studio/backend/auth/auth.db create mode 100644 studio/backend/auth/hashing.py create mode 100644 studio/backend/auth/storage.py create mode 100644 studio/backend/models/auth.py create mode 100644 studio/backend/routes/auth.py diff --git a/studio/backend/auth/__init__.py b/studio/backend/auth/__init__.py index e69de29bb2..4ea6ea0a8c 100644 --- a/studio/backend/auth/__init__.py +++ b/studio/backend/auth/__init__.py @@ -0,0 +1,24 @@ +""" +Authentication module for JWT-based auth with SQLite storage. +""" +from .jwt import create_access_token, get_current_subject, reload_secret +from .storage import ( + is_initialized, + create_initial_user, + get_user_and_secret, + load_jwt_secret, +) +from .hashing import hash_password, verify_password + +__all__ = [ + "create_access_token", + "get_current_subject", + "reload_secret", + "is_initialized", + "create_initial_user", + "get_user_and_secret", + "load_jwt_secret", + "hash_password", + "verify_password", +] + diff --git a/studio/backend/auth/auth.db b/studio/backend/auth/auth.db new file mode 100644 index 0000000000000000000000000000000000000000..b525e4b09621984252c0620d4efe40cc4fbfc4f2 GIT binary patch literal 12288 zcmeI$L66cv6ae5UYh2uf#T!TFsEL?P+bOLl3+Wo5pj(UV9t@c>9qCGGX}jQZLy!KX zJ^P#VXxYsYcTMDO-pjo7P5V0Oe7EWQQzy&=EaEg`IabkABuVHs#t5NOkq1TIg;W%e z?k2@`JKq;2^m6&VQvHQW$Jgliy82_^3-BBSKmY_l00ck)1V8`;KmY_l00bTgd_9y7 zYxTPH?PJdTS->vy>G(1W(p^z%ySC-o*t1%Wjd!>4vm^XRgdBGUp8dvl@jJKEx7>5w zv(I1L>)9?lXHkGX`^>|`LFaU6&siuN|93a6)FY< z2!H?xfB*=900@8p2!H?xfB*=9z+VEJwp2#7{-&}q8IQubZ^#qcRAfdNCqfvq!E_QB zx=9pDIBUv+>Y6G03ggt2Wt|W~H9=W3pemy7mB~^c-VnBcZ2!H?xfB*=900@8p2!H?xfB*>mUjpTW K+LN2V7QX>ySfjfD literal 0 HcmV?d00001 diff --git a/studio/backend/auth/hashing.py b/studio/backend/auth/hashing.py new file mode 100644 index 0000000000..c5d629a2a2 --- /dev/null +++ b/studio/backend/auth/hashing.py @@ -0,0 +1,40 @@ +""" +Password hashing utilities using PBKDF2. +""" +import hashlib +import hmac +import secrets +from typing import Tuple + + +def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]: + """ + Hash a password using PBKDF2-HMAC-SHA256. + + Returns (salt, hex_hash) tuple. + """ + if salt is None: + salt = secrets.token_hex(16) + dk = hashlib.pbkdf2_hmac( + "sha256", + password.encode("utf-8"), + salt.encode("utf-8"), + 100_000, # 100k iterations + ) + return salt, dk.hex() + + +def verify_password(password: str, salt: str, hashed: str) -> bool: + """ + Verify a password against a stored salt and hash. + + Uses constant-time comparison to prevent timing attacks. + """ + dk = hashlib.pbkdf2_hmac( + "sha256", + password.encode("utf-8"), + salt.encode("utf-8"), + 100_000, + ) + return hmac.compare_digest(dk.hex(), hashed) + diff --git a/studio/backend/auth/jwt.py b/studio/backend/auth/jwt.py index 2637efa821..33fc68e1e9 100644 --- a/studio/backend/auth/jwt.py +++ b/studio/backend/auth/jwt.py @@ -6,15 +6,20 @@ from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from jose import JWTError, jwt +from .storage import load_jwt_secret -# Ephemeral in-memory secret: -# - Generated fresh on each backend process start -# - Never written to disk -# - Not configurable by the user -SECRET_KEY = secrets.token_urlsafe(64) ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 60 +# Load stable secret from SQLite (set during first-time setup) +# This will raise RuntimeError if auth hasn't been initialized yet +try: + SECRET_KEY = load_jwt_secret() +except RuntimeError: + # Fallback: use a temporary secret until setup is complete + # This allows the app to start, but protected routes will fail until setup + SECRET_KEY = secrets.token_urlsafe(64) + security = HTTPBearer() # Reads Authorization: Bearer @@ -23,10 +28,9 @@ def create_access_token( expires_delta: Optional[timedelta] = None, ) -> str: """ - Create a signed JWT for the given subject (e.g. "local-user"). + Create a signed JWT for the given subject (e.g. username). - Tokens are valid only for the lifetime of this process, because the - SECRET_KEY is regenerated each time the backend restarts. + Tokens are valid across restarts because SECRET_KEY is stored in SQLite. """ to_encode = {"sub": subject} expire = datetime.now(UTC) + ( @@ -36,6 +40,16 @@ def create_access_token( return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) +def reload_secret() -> None: + """ + Reload the JWT secret from SQLite. + + Call this after setup to ensure new tokens use the persistent secret. + """ + global SECRET_KEY + SECRET_KEY = load_jwt_secret() + + async def get_current_subject( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: @@ -63,7 +77,7 @@ async def get_current_subject( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or expired token", ) -token = create_access_token("local-user") -print(token) +# token = create_access_token("local-user") +# print(token) diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py new file mode 100644 index 0000000000..0e2c9d7388 --- /dev/null +++ b/studio/backend/auth/storage.py @@ -0,0 +1,101 @@ +""" +SQLite storage for authentication data (user credentials + JWT secret). +""" +import sqlite3 +from pathlib import Path +from typing import Optional, Tuple + +DB_PATH = Path(__file__).parent / "auth.db" + + +def get_connection() -> sqlite3.Connection: + """Get a connection to the auth database, creating tables if needed.""" + conn = sqlite3.connect(DB_PATH) + conn.row_factory = sqlite3.Row + conn.execute( + """ + CREATE TABLE IF NOT EXISTS auth_user ( + id INTEGER PRIMARY KEY, + username TEXT UNIQUE NOT NULL, + password_salt TEXT NOT NULL, + password_hash TEXT NOT NULL, + jwt_secret TEXT NOT NULL + ); + """ + ) + conn.commit() + return conn + + +def is_initialized() -> bool: + """Check if auth has been set up (user exists in DB).""" + conn = get_connection() + cur = conn.execute("SELECT COUNT(*) AS c FROM auth_user") + row = cur.fetchone() + conn.close() + return bool(row["c"]) + + +def create_initial_user(username: str, password: str, jwt_secret: str) -> None: + """ + Create the initial admin user in the database. + + Raises sqlite3.IntegrityError if username already exists. + """ + from .hashing import hash_password + + salt, pwd_hash = hash_password(password) + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO auth_user (username, password_salt, password_hash, jwt_secret) + VALUES (?, ?, ?, ?) + """, + (username, salt, pwd_hash, jwt_secret), + ) + conn.commit() + finally: + conn.close() + + +def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str]]: + """ + Get user's password salt, hash, and JWT secret. + + Returns (password_salt, password_hash, jwt_secret) or None if user not found. + """ + conn = get_connection() + try: + cur = conn.execute( + """ + SELECT password_salt, password_hash, jwt_secret + FROM auth_user + WHERE username = ? + """, + (username,), + ) + row = cur.fetchone() + if not row: + return None + return row["password_salt"], row["password_hash"], row["jwt_secret"] + finally: + conn.close() + + +def load_jwt_secret() -> str: + """ + Load the JWT secret from the database. + + Raises RuntimeError if auth is not initialized. + """ + conn = get_connection() + try: + cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1") + row = cur.fetchone() + if not row: + raise RuntimeError("Auth is not initialized. Please set up a password first.") + return row["jwt_secret"] + finally: + conn.close() + diff --git a/studio/backend/main.py b/studio/backend/main.py index f335839f8d..f93f6c2820 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -9,7 +9,7 @@ from pathlib import Path from datetime import datetime # Import routers -from routes import training_router, models_router, inference_router, datasets_router +from routes import training_router, models_router, inference_router, datasets_router, auth_router # Create FastAPI app app = FastAPI( @@ -30,6 +30,7 @@ app.add_middleware( # ============ Register API Routes ============ # Register routers +app.include_router(auth_router, prefix="/api/auth", tags=["auth"]) app.include_router(training_router, prefix="/api/train", tags=["training"]) app.include_router(models_router, prefix="/api/models", tags=["models"]) app.include_router(inference_router, prefix="/api/inference", tags=["inference"]) diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index a47da0242c..55e827bd5a 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -12,6 +12,11 @@ from .models import ( LoRAInfo, LoRAScanResponse, ) +from .auth import ( + AuthSetupRequest, + AuthLoginRequest, + AuthStatusResponse, +) __all__ = [ # Training schemas @@ -23,5 +28,9 @@ __all__ = [ "ModelDetails", "LoRAInfo", "LoRAScanResponse", + # Auth schemas + "AuthSetupRequest", + "AuthLoginRequest", + "AuthStatusResponse", ] diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py new file mode 100644 index 0000000000..f71956db58 --- /dev/null +++ b/studio/backend/models/auth.py @@ -0,0 +1,22 @@ +""" +Pydantic schemas for Authentication API +""" +from pydantic import BaseModel, Field + + +class AuthSetupRequest(BaseModel): + """First-time setup: create the initial admin user + password.""" + username: str = Field(..., description="Admin username") + password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)") + + +class AuthLoginRequest(BaseModel): + """Login payload: username/password to obtain a JWT.""" + username: str = Field(..., description="Username") + password: str = Field(..., description="Password") + + +class AuthStatusResponse(BaseModel): + """Indicate whether auth has been initialized.""" + initialized: bool = Field(..., description="True if auth setup has been completed") + diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index a865794d54..5a16125a64 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -6,5 +6,6 @@ from routes.training import router as training_router from routes.models import router as models_router from routes.inference import router as inference_router from routes.datasets import router as datasets_router +from routes.auth import router as auth_router -__all__ = ["training_router", "models_router", "inference_router", "datasets_router"] \ No newline at end of file +__all__ = ["training_router", "models_router", "inference_router", "datasets_router", "auth_router"] \ No newline at end of file diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py new file mode 100644 index 0000000000..3f21eceeba --- /dev/null +++ b/studio/backend/routes/auth.py @@ -0,0 +1,88 @@ +""" +Authentication API routes +""" +from fastapi import APIRouter, HTTPException, status +import secrets + +from models.auth import ( + AuthSetupRequest, + AuthLoginRequest, + AuthStatusResponse, +) +from models.users import Token +from auth import storage, hashing +from auth.jwt import create_access_token, reload_secret + +router = APIRouter() + + +@router.get("/status", response_model=AuthStatusResponse) +async def auth_status() -> AuthStatusResponse: + """ + Check whether auth has already been initialized. + + - initialized = False -> frontend should show "Set admin password" screen. + - initialized = True -> frontend should show normal login. + """ + return AuthStatusResponse(initialized=storage.is_initialized()) + + +@router.post("/setup", response_model=Token, status_code=status.HTTP_201_CREATED) +async def setup_auth(payload: AuthSetupRequest) -> Token: + """ + First-time setup: create the admin user and a JWT secret. + + Can only be called once. Subsequent calls will return 400. + """ + if storage.is_initialized(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Auth is already initialized.", + ) + + # Generate a strong random JWT secret for this installation + jwt_secret = secrets.token_urlsafe(64) + + # Save username/password hash and secret in SQLite + try: + storage.create_initial_user( + username=payload.username, + password=payload.password, + jwt_secret=jwt_secret, + ) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to create user: {str(e)}", + ) + + # Reload JWT secret from DB (so jwt.py picks it up) + reload_secret() + + # Issue a token for the new user + access_token = create_access_token(subject=payload.username) + return Token(access_token=access_token, token_type="bearer") + + +@router.post("/login", response_model=Token) +async def login(payload: AuthLoginRequest) -> Token: + """ + Login with username/password and receive a JWT. + """ + record = storage.get_user_and_secret(payload.username) + if record is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + ) + + salt, pwd_hash, _jwt_secret = record + if not hashing.verify_password(payload.password, salt, pwd_hash): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + ) + + access_token = create_access_token(subject=payload.username) + return Token(access_token=access_token, token_type="bearer") +