ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: move checkpoint scanning to utils/models and /checkpoints endpoint to models router

Browse source
This commit is contained in:
Roland Tannous 2026-02-16 09:32:11 +00:00
commit f0298edeb8
8 changed files with 138 additions and 116 deletions
Download patch file Download diff file Expand all files Collapse all files

47
studio/backend/core/export/export.py
View file

@ -87,51 +87,8 @@ class ExportBackend:
Returns:
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
models = []
outputs_path = Path(outputs_dir)
if not outputs_path.exists():
logger.warning(f"Outputs directory not found: {outputs_dir}")
return models
try:
for item in outputs_path.iterdir():
if not item.is_dir():
continue
config_file = item / "config.json"
adapter_config = item / "adapter_config.json"
if not (config_file.exists() or adapter_config.exists()):
continue
# This is a valid training run
checkpoints = []
# Add the final model checkpoint
checkpoints.append((item.name, str(item)))
# Scan for intermediate checkpoints (checkpoint-N subdirs)
for sub in sorted(item.iterdir()):
if not sub.is_dir() or not sub.name.startswith("checkpoint-"):
continue
sub_config = sub / "config.json"
sub_adapter = sub / "adapter_config.json"
if sub_config.exists() or sub_adapter.exists():
checkpoints.append((sub.name, str(sub)))
models.append((item.name, checkpoints))
logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
# Sort by modification time (newest first)
models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True)
logger.info(f"Found {len(models)} training runs in {outputs_dir}")
return models
except Exception as e:
logger.error(f"Error scanning checkpoints: {e}")
return []
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir=outputs_dir)
def load_checkpoint(self,
checkpoint_path: str,

6
studio/backend/models/__init__.py
View file

@ -8,6 +8,9 @@ from .training import (
TrainingProgress,
)
from .models import (
CheckpointInfo,
ModelCheckpoints,
CheckpointListResponse,
ModelDetails,
LoRAInfo,
LoRAScanResponse,
@ -20,9 +23,6 @@ from .auth import (
AuthStatusResponse,
)
from .export import (
CheckpointInfo,
ModelCheckpoints,
CheckpointListResponse,
LoadCheckpointRequest,
ExportStatusResponse,
ExportOperationResponse,

27
studio/backend/models/export.py
View file

@ -5,33 +5,6 @@ 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 ModelCheckpoints(BaseModel):
"""A training run and its associated checkpoints."""
name: str = Field(..., description="Training run folder name")
checkpoints: List[CheckpointInfo] = Field(
default_factory=list,
description="List of checkpoints for this training run (final + intermediate)",
)
class CheckpointListResponse(BaseModel):
"""Response for listing available checkpoints in an outputs directory."""
outputs_dir: str = Field(..., description="Directory that was scanned")
models: List[ModelCheckpoints] = Field(
default_factory=list,
description="List of training runs with their checkpoints",
)
class LoadCheckpointRequest(BaseModel):
"""Request for loading a checkpoint into the export backend."""

27
studio/backend/models/models.py
View file

@ -5,6 +5,33 @@ from pydantic import BaseModel, Field
from typing import Optional, List, 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 ModelCheckpoints(BaseModel):
"""A training run and its associated checkpoints."""
name: str = Field(..., description="Training run folder name")
checkpoints: List[CheckpointInfo] = Field(
default_factory=list,
description="List of checkpoints for this training run (final + intermediate)",
)
class CheckpointListResponse(BaseModel):
"""Response for listing available checkpoints in an outputs directory."""
outputs_dir: str = Field(..., description="Directory that was scanned")
models: List[ModelCheckpoints] = Field(
default_factory=list,
description="List of training runs with their checkpoints",
)
class ModelDetails(BaseModel):
"""Detailed model configuration and metadata - can be used for both list and detail views"""
id: str = Field(..., description="Model identifier")

41
studio/backend/routes/export.py
View file

@ -26,9 +26,6 @@ except ImportError:
# Import Pydantic models
from models import (
CheckpointInfo,
CheckpointListResponse,
ModelCheckpoints,
LoadCheckpointRequest,
ExportStatusResponse,
ExportOperationResponse,
@ -51,44 +48,6 @@ if not logger.handlers:
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_models = backend.scan_checkpoints(outputs_dir=outputs_dir)
models = [
ModelCheckpoints(
name=model_name,
checkpoints=[
CheckpointInfo(display_name=display_name, path=path)
for display_name, path in checkpoints
],
)
for model_name, checkpoints in raw_models
]
return CheckpointListResponse(
outputs_dir=outputs_dir,
models=models,
)
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)

42
studio/backend/routes/models.py
View file

@ -21,6 +21,7 @@ try:
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
scan_checkpoints,
ModelConfig,
)
from core.inference import get_inference_backend
@ -34,11 +35,15 @@ except ImportError:
load_model_defaults,
get_base_model_from_lora,
is_vision_model,
scan_checkpoints,
ModelConfig,
)
from core.inference import get_inference_backend
from models import (
CheckpointInfo,
CheckpointListResponse,
ModelCheckpoints,
ModelDetails,
LoRAScanResponse,
LoRAInfo,
@ -266,3 +271,40 @@ async def check_vision_model(
detail=f"Failed to check vision model: {str(e)}"
)
@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.
Scans the outputs folder for training runs and their checkpoints.
"""
try:
raw_models = scan_checkpoints(outputs_dir=outputs_dir)
models = [
ModelCheckpoints(
name=model_name,
checkpoints=[
CheckpointInfo(display_name=display_name, path=path)
for display_name, path in checkpoints
],
)
for model_name, checkpoints in raw_models
]
return CheckpointListResponse(
outputs_dir=outputs_dir,
models=models,
)
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)}",
)

2
studio/backend/utils/models/__init__.py
View file

@ -11,6 +11,7 @@ from .model_config import (
MODEL_NAME_MAPPING,
UI_STATUS_INDICATORS,
)
from .checkpoints import scan_checkpoints
__all__ = [
'ModelConfig',
@ -21,4 +22,5 @@ __all__ = [
'load_model_config',
'MODEL_NAME_MAPPING',
'UI_STATUS_INDICATORS',
'scan_checkpoints',
]

62
studio/backend/utils/models/checkpoints.py Normal file
View file

@ -0,0 +1,62 @@
"""
Checkpoint scanning utilities for discovering training runs and their checkpoints.
"""
import logging
from pathlib import Path
from typing import List, Tuple
logger = logging.getLogger(__name__)
def scan_checkpoints(outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tuple[str, str]]]]:
"""
Scan outputs folder for training runs and their checkpoints.
Returns:
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
models = []
outputs_path = Path(outputs_dir)
if not outputs_path.exists():
logger.warning(f"Outputs directory not found: {outputs_dir}")
return models
try:
for item in outputs_path.iterdir():
if not item.is_dir():
continue
config_file = item / "config.json"
adapter_config = item / "adapter_config.json"
if not (config_file.exists() or adapter_config.exists()):
continue
# This is a valid training run
checkpoints = []
# Add the final model checkpoint
checkpoints.append((item.name, str(item)))
# Scan for intermediate checkpoints (checkpoint-N subdirs)
for sub in sorted(item.iterdir()):
if not sub.is_dir() or not sub.name.startswith("checkpoint-"):
continue
sub_config = sub / "config.json"
sub_adapter = sub / "adapter_config.json"
if sub_config.exists() or sub_adapter.exists():
checkpoints.append((sub.name, str(sub)))
models.append((item.name, checkpoints))
logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
# Sort by modification time (newest first)
models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True)
logger.info(f"Found {len(models)} training runs in {outputs_dir}")
return models
except Exception as e:
logger.error(f"Error scanning checkpoints: {e}")
return []