Merge pull request #227 from togethercomputer/sbassam/update-rl-specs

MOSH-1979: Sync RL API OpenAPI spec with latest shaping protos to incorporate training checkpoint saving/resuming

Author: sbassam

openapi.yaml

@@ -7493,6 +7493,64 @@ paths:
           schema:
             description: Operation ID
             type: string
+  /rl/training-sessions/{session_id}/operations/training-checkpoint:
+    post:
+      summary: Save training checkpoint
+      description: Submits an operation that will asynchronously save the full training state (adapter + optimizer + step).
+      operationId: createTrainingCheckpoint
+      tags: [RL]
+      responses:
+        "200":
+          description: Save training checkpoint operation details
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/RL.TrainingCheckpointOperation'
+        default:
+          description: An unexpected error response.
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorData'
+      parameters:
+        - name: session_id
+          in: path
+          required: true
+          schema:
+            description: Training session ID
+            type: string
+  /rl/training-sessions/{session_id}/operations/training-checkpoint/{operation_id}:
+    get:
+      summary: Get save training checkpoint operation
+      description: Retrieves the current status and result of a save training checkpoint operation.
+      operationId: getTrainingCheckpointOperation
+      tags: [RL]
+      responses:
+        "200":
+          description: Save training checkpoint operation details
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/RL.TrainingCheckpointOperation'
+        default:
+          description: An unexpected error response.
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorData'
+      parameters:
+        - name: session_id
+          in: path
+          required: true
+          schema:
+            description: Training session ID
+            type: string
+        - name: operation_id
+          in: path
+          required: true
+          schema:
+            description: Operation ID
+            type: string
   /rl/checkpoints/{id}/download:
     get:
       summary: Download checkpoint
@@ -7977,10 +8035,10 @@ components:
           description: Base model to use for the training session
           type: string
           example: meta-llama/Meta-Llama-3-8B-Instruct
-        checkpoint_id:
-          description: Checkpoint ID to use for the training session
+        resume_from_checkpoint_id:
+          description: Checkpoint ID to resume from
           type: string
-          example: checkpoint-123
+          example: 123e4567-e89b-12d3-a456-426614174000
         lora_config:
           $ref: '#/components/schemas/RL.LoraConfig'
     RL.TrainingSessionStatus:
@@ -8016,6 +8074,16 @@ components:
             type: object
             $ref: '#/components/schemas/RL.InferenceCheckpoint'
           description: List of saved inference checkpoints for this session
+        training_checkpoints:
+          type: array
+          items:
+            type: object
+            $ref: '#/components/schemas/RL.TrainingCheckpoint'
+          description: List of saved training checkpoints for this session
+        resume_from_checkpoint_id:
+          type: string
+          example: 123e4567-e89b-12d3-a456-426614174000
+          description: Checkpoint ID this session was resumed from
         step:
           description: Current training step
           type: string
@@ -8152,6 +8220,46 @@ components:
           format: date-time
           example: "2026-01-02T00:00:00Z"
           description: Timestamp when the model was registered
+    RL.TrainingCheckpoint:
+      type: object
+      description: Saved training checkpoint
+      properties:
+        id:
+          type: string
+          example: 123e4567-e89b-12d3-a456-426614174000
+          description: Unique identifier for the checkpoint
+        step:
+          type: string
+          format: uint64
+          example: 42
+          description: Training step at time of save
+        created_at:
+          type: string
+          format: date-time
+          example: "2026-01-02T00:00:00Z"
+          description: Timestamp when the checkpoint was created
+    RL.TrainingCheckpointResult:
+      type: object
+      properties:
+        checkpoint_id:
+          type: string
+          example: 550e8400-e29b-41d4-a716-446655440000
+          description: ID of the saved training checkpoint (use for resume via Start)
+    RL.TrainingCheckpointOperation:
+      type: object
+      properties:
+        id:
+          type: string
+          example: 550e8400-e29b-41d4-a716-446655440000
+          description: Operation ID
+        status:
+          $ref: '#/components/schemas/RL.TrainingOperationStatus'
+          example: TRAINING_OPERATION_STATUS_PENDING
+          description: Operation status
+        output:
+          $ref: '#/components/schemas/RL.TrainingCheckpointResult'
+        error:
+          $ref: '#/components/schemas/RL.TrainingOperationError'
     RL.CheckpointVariant:
       type: string
       enum: