added the pydantic models and routes for export

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 ============

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",

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
+
+

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"]
+__all__ = [
+    "training_router",
+    "models_router",
+    "inference_router",
+    "datasets_router",
+    "auth_router",
+    "export_router",
+]

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)}",
+        )
+
+