Add first draft of resuming training from latest checkpoint

transformerlab/routers/experiment/jobs.py

@@ -696,3 +696,4 @@ async def get_artifacts(job_id: str, request: Request):
     artifacts.sort(key=lambda x: x["filename"], reverse=True)
 
     return {"artifacts": artifacts}
+

transformerlab/routers/remote.py

@@ -1,9 +1,12 @@
 import os
 import httpx
-from fastapi import APIRouter, Form, Request, File, UploadFile
+import json
+import re
+from fastapi import APIRouter, Form, Request, File, UploadFile, HTTPException
 from typing import Optional, List
 from transformerlab.services import job_service
 from transformerlab.services.job_service import job_update_status
+from lab.dirs import get_workspace_dir, get_job_checkpoints_dir
 
 
 router = APIRouter(prefix="/remote", tags=["remote"])
@@ -190,6 +193,22 @@
                         job_id, "cluster_name", response_data["cluster_name"], experimentId
                     )
 
+                # Save the launch_config for potential resume
+                launch_config = {
+                    "command": command,
+                    "setup": setup,
+                    "cluster_name": cluster_name,
+                    "cpus": cpus,
+                    "memory": memory,
+                    "disk_space": disk_space,
+                    "accelerators": accelerators,
+                    "num_nodes": num_nodes,
+                    "uploaded_dir_path": uploaded_dir_path,
+                }
+                job_service.job_update_job_data_insert_key_value(
+                    job_id, "launch_config", json.dumps(launch_config), experimentId
+                )
+
                 return {
                     "status": "success",
                     "data": response_data,
@@ -284,7 +303,6 @@
     Upload a directory to the remote Lattice orchestrator for later use in cluster launches.
     Files are stored locally first, then sent to orchestrator.
     """
-    from lab.dirs import get_workspace_dir
 
     # Validate environment variables
     result = validate_gpu_orchestrator_env_vars()
@@ -559,3 +577,115 @@
     except Exception as e:
         print(f"Error checking remote job status: {str(e)}")
         return {"status": "error", "message": "Error checking remote job status"}
+
+
+@router.post("/{job_id}/resume_from_checkpoint")
+async def resume_from_checkpoint(
+    job_id: str,
+    request: Request,
+):
+    """Resume training from a checkpoint by creating a new job with the checkpoint parameter"""
+    try:
+        body = await request.json()
+        checkpoint_name = body.get("checkpoint_name")
+        experimentId = body.get("experimentId")
+
+        if not checkpoint_name:
+            raise HTTPException(status_code=400, detail="checkpoint_name is required")
+        if not experimentId:
+            raise HTTPException(status_code=400, detail="experimentId is required")
+
+        # Get the original job
+        original_job = job_service.job_get(job_id)
+        if not original_job:
+            raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
+
+        original_job_data = original_job.get("job_data", {})
+
+        # Get the launch_config from the original job
+        launch_config_str = original_job_data.get("launch_config")
+        if not launch_config_str:
+            raise HTTPException(status_code=400, detail="launch_config not found in original job")
+
+        # Parse the launch_config JSON string
+        try:
+            launch_config = json.loads(launch_config_str)
+        except json.JSONDecodeError:
+            raise HTTPException(status_code=400, detail="Invalid launch_config format")
+
+        # Get the checkpoint directory path
+        checkpoints_dir = get_job_checkpoints_dir(job_id)
+        checkpoint_path = os.path.join(checkpoints_dir, checkpoint_name)
+
+        if not os.path.exists(checkpoint_path):
+            raise HTTPException(status_code=404, detail=f"Checkpoint {checkpoint_name} not found at {checkpoint_path}")
+
+        # Modify the command to include the checkpoint parameter
+        original_command = launch_config.get("command", "")
+        if not original_command:
+            raise HTTPException(status_code=400, detail="Original command not found in launch_config")
+
+        # Parse the command to find the python execution part and add --resume_from_checkpoint
+        command_parts = original_command.split(' && ')
+        python_part_index = None
+        for i, part in enumerate(command_parts):
+            if part.strip().startswith('python '):
+                python_part_index = i
+                break
+
+        if python_part_index is None:
+            raise HTTPException(status_code=400, detail="Could not find python execution command in original command")
+
+        # Modify the python part to include --resume_from_checkpoint
+        python_command = command_parts[python_part_index].strip()
+
+        # Check if --resume_from_checkpoint already exists
+        if "--resume_from_checkpoint" in python_command:
+            # Replace existing checkpoint parameter
+            resume_python_command = re.sub(
+                r'--resume_from_checkpoint\s+\S+',
+                f'--resume_from_checkpoint {checkpoint_path}',
+                python_command
+            )
+        else:
+            # Add the checkpoint parameter
+            resume_python_command = f"{python_command} --resume_from_checkpoint {checkpoint_path}"
+
+        # Replace the python part in the command_parts
+        command_parts[python_part_index] = resume_python_command
+
+        # Rejoin the command parts
+        resume_command = ' && '.join(command_parts)        
+        # Call launch_remote to create and launch the new job
+        launch_result = await launch_remote(
+            request=request,
+            experimentId=experimentId,
+            cluster_name=launch_config.get("cluster_name"),
+            command=resume_command,
+            task_name=f"Resume from {checkpoint_name}",
+            cpus=launch_config.get("cpus"),
+            memory=launch_config.get("memory"),
+            disk_space=launch_config.get("disk_space"),
+            accelerators=launch_config.get("accelerators"),
+            num_nodes=launch_config.get("num_nodes"),
+            setup=launch_config.get("setup"),
+            uploaded_dir_path=launch_config.get("uploaded_dir_path"),
+        )
+
+        if launch_result.get("status") == "success":
+            new_job_id = launch_result.get("job_id")
+            # Update the new job's job_data to mark it as resumed
+            job_service.job_update_job_data_insert_key_value(new_job_id, "resumed_from_checkpoint", checkpoint_name, experimentId)
+            job_service.job_update_job_data_insert_key_value(new_job_id, "parent_job_id", job_id, experimentId)
+            return {
+                "status": "success",
+                "job_id": new_job_id,
+                "message": f"Training resumed from checkpoint {checkpoint_name}",
+                "data": launch_result
+            }
+        else:
+            raise HTTPException(status_code=500, detail=f"Failed to launch remote job: {launch_result.get('message')}")
+
+    except Exception as e:
+        print(f"Error checking remote job status: {str(e)}")
+        return {"status": "error", "message": "Error checking remote job status"}