Merge pull request #106 from unslothai/fix/adding-checkpointing-data-to-api

Fixing the get checkpoint api

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

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

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

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)