diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 8977586ac3..7a767659ac 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -80,39 +80,54 @@ class ExportBackend: logger.error(f"Error during memory cleanup: {e}") return False - def scan_checkpoints(self, outputs_dir: str = "./outputs") -> List[Tuple[str, str]]: + def scan_checkpoints(self, outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tuple[str, str]]]]: """ - Scan outputs folder for model checkpoints. + Scan outputs folder for training runs and their checkpoints. Returns: - List of tuples: [(display_name, checkpoint_path), ...] + List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...] """ - checkpoints = [] + models = [] outputs_path = Path(outputs_dir) if not outputs_path.exists(): logger.warning(f"Outputs directory not found: {outputs_dir}") - return checkpoints + return models try: for item in outputs_path.iterdir(): - if item.is_dir(): - # Check if this directory contains a model - config_file = item / "config.json" - adapter_config = item / "adapter_config.json" + if not item.is_dir(): + continue - if config_file.exists() or adapter_config.exists(): - # This is a valid checkpoint - display_name = item.name - checkpoint_path = str(item) - checkpoints.append((display_name, checkpoint_path)) - logger.debug(f"Found checkpoint: {display_name}") + 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) - checkpoints.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True) + models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True) - logger.info(f"Found {len(checkpoints)} checkpoints in {outputs_dir}") - return checkpoints + logger.info(f"Found {len(models)} training runs in {outputs_dir}") + return models except Exception as e: logger.error(f"Error scanning checkpoints: {e}") diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index 8d947b1b9c..aad07fd8b0 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -21,6 +21,7 @@ from .auth import ( ) from .export import ( CheckpointInfo, + ModelCheckpoints, CheckpointListResponse, LoadCheckpointRequest, ExportStatusResponse, @@ -68,6 +69,7 @@ __all__ = [ "AuthStatusResponse", # Export schemas "CheckpointInfo", + "ModelCheckpoints", "CheckpointListResponse", "LoadCheckpointRequest", "ExportStatusResponse", diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index d108c00751..779ddb1416 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -12,13 +12,23 @@ class CheckpointInfo(BaseModel): 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") - checkpoints: List[CheckpointInfo] = Field( + models: List[ModelCheckpoints] = Field( default_factory=list, - description="List of discovered checkpoints", + description="List of training runs with their checkpoints", ) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index b72dea48b8..34df78e276 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -28,6 +28,7 @@ except ImportError: from models import ( CheckpointInfo, CheckpointListResponse, + ModelCheckpoints, LoadCheckpointRequest, ExportStatusResponse, ExportOperationResponse, @@ -65,16 +66,22 @@ async def list_checkpoints( """ try: backend = get_export_backend() - raw_checkpoints = backend.scan_checkpoints(outputs_dir=outputs_dir) + raw_models = backend.scan_checkpoints(outputs_dir=outputs_dir) - checkpoints = [ - CheckpointInfo(display_name=display_name, path=path) - for display_name, path in raw_checkpoints + 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, - checkpoints=checkpoints, + models=models, ) except Exception as e: logger.error(f"Error listing checkpoints: {e}", exc_info=True)