From 40bfe42974f7ee43f569f6701f41c798e46e3df2 Mon Sep 17 00:00:00 2001 From: sshah229 Date: Tue, 10 Feb 2026 16:16:17 -0700 Subject: [PATCH] added the pydantic models and routes for export --- studio/backend/main.py | 1 + studio/backend/models/__init__.py | 21 ++ studio/backend/models/export.py | 141 ++++++++++++++ studio/backend/routes/__init__.py | 10 +- studio/backend/routes/export.py | 310 ++++++++++++++++++++++++++++++ 5 files changed, 482 insertions(+), 1 deletion(-) create mode 100644 studio/backend/models/export.py create mode 100644 studio/backend/routes/export.py diff --git a/studio/backend/main.py b/studio/backend/main.py index e957c668e3..34c61bed6b 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -67,6 +67,7 @@ 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"]) app.include_router(datasets_router, prefix="/api/datasets", tags=["datasets"]) +app.include_router(export_router, prefix="/api/export", tags=["export"]) # ============ Health and System Endpoints ============ diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index 584fb9f7c2..8d947b1b9c 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -19,6 +19,17 @@ from .auth import ( RefreshTokenRequest, AuthStatusResponse, ) +from .export import ( + CheckpointInfo, + CheckpointListResponse, + LoadCheckpointRequest, + ExportStatusResponse, + ExportOperationResponse, + ExportMergedModelRequest, + ExportBaseModelRequest, + ExportGGUFRequest, + ExportLoRAAdapterRequest, +) from .users import Token from .datasets import ( CheckFormatRequest, @@ -55,6 +66,16 @@ __all__ = [ "AuthLoginRequest", "RefreshTokenRequest", "AuthStatusResponse", + # Export schemas + "CheckpointInfo", + "CheckpointListResponse", + "LoadCheckpointRequest", + "ExportStatusResponse", + "ExportOperationResponse", + "ExportMergedModelRequest", + "ExportBaseModelRequest", + "ExportGGUFRequest", + "ExportLoRAAdapterRequest", "Token", # Dataset schemas "CheckFormatRequest", diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py new file mode 100644 index 0000000000..d108c00751 --- /dev/null +++ b/studio/backend/models/export.py @@ -0,0 +1,141 @@ +""" +Pydantic schemas for Export API. +""" +from pydantic import BaseModel, Field +from typing import List, Optional, Literal, Dict, Any + + +class CheckpointInfo(BaseModel): + """Information about a discovered checkpoint directory.""" + + display_name: str = Field(..., description="User-friendly checkpoint name (folder name)") + path: str = Field(..., description="Full path to the checkpoint directory") + + +class CheckpointListResponse(BaseModel): + """Response for listing available checkpoints in an outputs directory.""" + + outputs_dir: str = Field(..., description="Directory that was scanned") + checkpoints: List[CheckpointInfo] = Field( + default_factory=list, + description="List of discovered checkpoints", + ) + + +class LoadCheckpointRequest(BaseModel): + """Request for loading a checkpoint into the export backend.""" + + checkpoint_path: str = Field(..., description="Path to the checkpoint directory") + max_seq_length: int = Field( + 2048, + ge=128, + le=32768, + description="Maximum sequence length for loading the model", + ) + load_in_4bit: bool = Field( + True, + description="Whether to load the model in 4-bit quantization", + ) + + +class ExportStatusResponse(BaseModel): + """Current export backend status.""" + + current_checkpoint: Optional[str] = Field( + None, + description="Path to the currently loaded checkpoint, if any", + ) + is_vision: bool = Field( + False, + description="True if the loaded checkpoint is a vision model", + ) + is_peft: bool = Field( + False, + description="True if the loaded checkpoint is a PEFT (LoRA) model", + ) + + +class ExportOperationResponse(BaseModel): + """Generic response for export operations.""" + + success: bool = Field(..., description="True if the operation succeeded") + message: str = Field(..., description="Human-readable status or error message") + details: Optional[Dict[str, Any]] = Field( + default=None, + description="Optional extra details about the operation", + ) + + +class ExportCommonOptions(BaseModel): + """Common options for export operations that save locally and/or push to Hub.""" + + save_directory: str = Field( + ..., + description="Local directory where the exported artifacts will be written", + ) + push_to_hub: bool = Field( + False, + description="If True, also push the exported model to the Hugging Face Hub", + ) + repo_id: Optional[str] = Field( + None, + description="Hugging Face Hub repository ID (username/model-name)", + ) + hf_token: Optional[str] = Field( + None, + description="Hugging Face access token used for Hub operations", + ) + private: bool = Field( + False, + description="If True, create a private repository on the Hub (where applicable)", + ) + + +class ExportMergedModelRequest(ExportCommonOptions): + """Request for exporting a merged PEFT model.""" + + format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field( + "16-bit (FP16)", + description="Export precision / format for the merged model", + ) + + +class ExportBaseModelRequest(ExportCommonOptions): + """Request for exporting a non-PEFT (base) model.""" + + # Uses fields from ExportCommonOptions only + pass + + +class ExportGGUFRequest(BaseModel): + """Request for exporting the current model to GGUF format.""" + + save_directory: str = Field( + ..., + description="Directory where GGUF files will be saved", + ) + quantization_method: str = Field( + "Q4_K_M", + description='GGUF quantization method (e.g. "Q4_K_M")', + ) + push_to_hub: bool = Field( + False, + description="If True, also push GGUF artifacts to the Hugging Face Hub", + ) + repo_id: Optional[str] = Field( + None, + description="Hugging Face Hub repository ID for GGUF upload", + ) + hf_token: Optional[str] = Field( + None, + description="Hugging Face token for GGUF upload", + ) + + +class ExportLoRAAdapterRequest(ExportCommonOptions): + """Request for exporting only the LoRA adapter (not merged).""" + + # Uses fields from ExportCommonOptions only + pass + + diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index 5a16125a64..04c7ff7f1a 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -7,5 +7,13 @@ 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 +from routes.export import router as export_router -__all__ = ["training_router", "models_router", "inference_router", "datasets_router", "auth_router"] \ No newline at end of file +__all__ = [ + "training_router", + "models_router", + "inference_router", + "datasets_router", + "auth_router", + "export_router", +] \ No newline at end of file diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py new file mode 100644 index 0000000000..b72dea48b8 --- /dev/null +++ b/studio/backend/routes/export.py @@ -0,0 +1,310 @@ +""" +Export API routes: checkpoint discovery and model export operations. +""" + +import sys +from pathlib import Path +from fastapi import APIRouter, Depends, HTTPException, Query +import logging + +# 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)) + +# Auth +from auth.authentication import get_current_subject + +# Import backend functions +try: + from core.export import get_export_backend +except ImportError: + parent_backend = backend_path.parent / "backend" + if str(parent_backend) not in sys.path: + sys.path.insert(0, str(parent_backend)) + from core.export import get_export_backend + +# Import Pydantic models +from models import ( + CheckpointInfo, + CheckpointListResponse, + LoadCheckpointRequest, + ExportStatusResponse, + ExportOperationResponse, + ExportMergedModelRequest, + ExportBaseModelRequest, + ExportGGUFRequest, + ExportLoRAAdapterRequest, +) + +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("/checkpoints", response_model=CheckpointListResponse) +async def list_checkpoints( + outputs_dir: str = Query( + default="./outputs", + description="Directory to scan for checkpoints", + ), + current_subject: str = Depends(get_current_subject), +): + """ + List available checkpoints in the outputs directory. + + Wraps ExportBackend.scan_checkpoints. + """ + try: + backend = get_export_backend() + raw_checkpoints = backend.scan_checkpoints(outputs_dir=outputs_dir) + + checkpoints = [ + CheckpointInfo(display_name=display_name, path=path) + for display_name, path in raw_checkpoints + ] + + return CheckpointListResponse( + outputs_dir=outputs_dir, + checkpoints=checkpoints, + ) + except Exception as e: + logger.error(f"Error listing checkpoints: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to list checkpoints: {str(e)}", + ) + + +@router.post("/load-checkpoint", response_model=ExportOperationResponse) +async def load_checkpoint( + request: LoadCheckpointRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Load a checkpoint into the export backend. + + Wraps ExportBackend.load_checkpoint. + """ + try: + backend = get_export_backend() + success, message = backend.load_checkpoint( + checkpoint_path=request.checkpoint_path, + max_seq_length=request.max_seq_length, + load_in_4bit=request.load_in_4bit, + ) + + if not success: + raise HTTPException(status_code=400, detail=message) + + return ExportOperationResponse(success=True, message=message) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error loading checkpoint: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to load checkpoint: {str(e)}", + ) + + +@router.post("/cleanup", response_model=ExportOperationResponse) +async def cleanup_export_memory( + current_subject: str = Depends(get_current_subject), +): + """ + Cleanup export-related models from memory (GPU/CPU). + + Wraps ExportBackend.cleanup_memory. + """ + try: + backend = get_export_backend() + success = backend.cleanup_memory() + + if not success: + raise HTTPException( + status_code=500, + detail="Memory cleanup failed. See server logs for details.", + ) + + return ExportOperationResponse( + success=True, + message="Memory cleanup completed successfully", + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error during export memory cleanup: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to cleanup export memory: {str(e)}", + ) + + +@router.get("/status", response_model=ExportStatusResponse) +async def get_export_status( + current_subject: str = Depends(get_current_subject), +): + """ + Get current export backend status (loaded checkpoint, model type, PEFT flag). + """ + try: + backend = get_export_backend() + return ExportStatusResponse( + current_checkpoint=backend.current_checkpoint, + is_vision=bool(getattr(backend, "is_vision", False)), + is_peft=bool(getattr(backend, "is_peft", False)), + ) + except Exception as e: + logger.error(f"Error getting export status: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to get export status: {str(e)}", + ) + + +@router.post("/export/merged", response_model=ExportOperationResponse) +async def export_merged_model( + request: ExportMergedModelRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Export a merged PEFT model (e.g., 16-bit or 4-bit) and optionally push to Hub. + + Wraps ExportBackend.export_merged_model. + """ + try: + backend = get_export_backend() + success, message = backend.export_merged_model( + save_directory=request.save_directory, + format_type=request.format_type, + push_to_hub=request.push_to_hub, + repo_id=request.repo_id, + hf_token=request.hf_token, + private=request.private, + ) + + if not success: + raise HTTPException(status_code=400, detail=message) + + return ExportOperationResponse(success=True, message=message) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting merged model: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to export merged model: {str(e)}", + ) + + +@router.post("/export/base", response_model=ExportOperationResponse) +async def export_base_model( + request: ExportBaseModelRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Export a non-PEFT base model and optionally push to Hub. + + Wraps ExportBackend.export_base_model. + """ + try: + backend = get_export_backend() + success, message = backend.export_base_model( + save_directory=request.save_directory, + push_to_hub=request.push_to_hub, + repo_id=request.repo_id, + hf_token=request.hf_token, + private=request.private, + ) + + if not success: + raise HTTPException(status_code=400, detail=message) + + return ExportOperationResponse(success=True, message=message) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting base model: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to export base model: {str(e)}", + ) + + +@router.post("/export/gguf", response_model=ExportOperationResponse) +async def export_gguf( + request: ExportGGUFRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Export the current model to GGUF format and optionally push to Hub. + + Wraps ExportBackend.export_gguf. + """ + try: + backend = get_export_backend() + success, message = backend.export_gguf( + save_directory=request.save_directory, + quantization_method=request.quantization_method, + push_to_hub=request.push_to_hub, + repo_id=request.repo_id, + hf_token=request.hf_token, + ) + + if not success: + raise HTTPException(status_code=400, detail=message) + + return ExportOperationResponse(success=True, message=message) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting GGUF model: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to export GGUF model: {str(e)}", + ) + + +@router.post("/export/lora", response_model=ExportOperationResponse) +async def export_lora_adapter( + request: ExportLoRAAdapterRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Export only the LoRA adapter (if the loaded model is PEFT). + + Wraps ExportBackend.export_lora_adapter. + """ + try: + backend = get_export_backend() + success, message = backend.export_lora_adapter( + save_directory=request.save_directory, + push_to_hub=request.push_to_hub, + repo_id=request.repo_id, + hf_token=request.hf_token, + private=request.private, + ) + + if not success: + raise HTTPException(status_code=400, detail=message) + + return ExportOperationResponse(success=True, message=message) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting LoRA adapter: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to export LoRA adapter: {str(e)}", + ) + +