Get offline batch inference details using task API in m (#8305)

* get offline batch inference details using task API in ml Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * Doc review Signed-off-by: Fanit Kolchina <kolchfa@amazon.com> * Typo fix Signed-off-by: Fanit Kolchina <kolchfa@amazon.com> * Apply suggestions from code review Co-authored-by: Nathan Bower <nbower@amazon.com> Signed-off-by: kolchfa-aws <105444904+kolchfa-aws@users.noreply.github.com> * Update _ml-commons-plugin/api/model-apis/batch-predict.md Signed-off-by: kolchfa-aws <105444904+kolchfa-aws@users.noreply.github.com> * Update _ml-commons-plugin/api/model-apis/batch-predict.md Signed-off-by: kolchfa-aws <105444904+kolchfa-aws@users.noreply.github.com> * Add parameter values Signed-off-by: Fanit Kolchina <kolchfa@amazon.com> * Extra spaces Signed-off-by: Fanit Kolchina <kolchfa@amazon.com> --------- Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> Signed-off-by: Fanit Kolchina <kolchfa@amazon.com> Signed-off-by: kolchfa-aws <105444904+kolchfa-aws@users.noreply.github.com> Co-authored-by: Fanit Kolchina <kolchfa@amazon.com> Co-authored-by: kolchfa-aws <105444904+kolchfa-aws@users.noreply.github.com> Co-authored-by: Nathan Bower <nbower@amazon.com>
opensearch-project · Sep 17, 2024 · db292d9 · db292d9
1 parent 709a54f
commit db292d9
Showing 1 changed file with 102 additions and 38 deletions.
diff --git a/_ml-commons-plugin/api/model-apis/batch-predict.md b/_ml-commons-plugin/api/model-apis/batch-predict.md
@@ -31,7 +31,13 @@ POST /_plugins/_ml/models/<model_id>/_batch_predict
 
 ## Prerequisites
 
-Before using the Batch Predict API, you need to create a connector to the externally hosted model. For example, to create a connector to an OpenAI `text-embedding-ada-002` model, send the following request:
+Before using the Batch Predict API, you need to create a connector to the externally hosted model. For each action, specify the `action_type` parameter that describes the action:
+
+- `batch_predict`: Runs the batch predict operation.
+- `batch_predict_status`: Checks the batch predict operation status.
+- `cancel_batch_predict`: Cancels the batch predict operation.
+
+For example, to create a connector to an OpenAI `text-embedding-ada-002` model, send the following request. The `cancel_batch_predict` action is optional and supports canceling the batch job running on OpenAI:
 
 ```json
 POST /_plugins/_ml/connectors/_create
@@ -68,6 +74,22 @@ POST /_plugins/_ml/connectors/_create
         "Authorization": "Bearer ${credential.openAI_key}"
       },
       "request_body": "{ \"input_file_id\": \"${parameters.input_file_id}\", \"endpoint\": \"${parameters.endpoint}\", \"completion_window\": \"24h\" }"
+    },
+    {
+      "action_type": "batch_predict_status",
+      "method": "GET",
+      "url": "https://api.openai.com/v1/batches/${parameters.id}",
+      "headers": {
+        "Authorization": "Bearer ${credential.openAI_key}"
+      }
+    },
+    {
+      "action_type": "cancel_batch_predict",
+      "method": "POST",
+      "url": "https://api.openai.com/v1/batches/${parameters.id}/cancel",
+      "headers": {
+        "Authorization": "Bearer ${credential.openAI_key}"
+      }
     }
   ]
 }
@@ -123,45 +145,87 @@ POST /_plugins/_ml/models/lyjxwZABNrAVdFa9zrcZ/_batch_predict
 
 #### Example response
 
+The response contains the task ID for the batch predict operation:
+
 ```json
 {
-  "inference_results": [
-    {
-      "output": [
-        {
-          "name": "response",
-          "dataAsMap": {
-            "id": "batch_<your file id>",
-            "object": "batch",
-            "endpoint": "/v1/embeddings",
-            "errors": null,
-            "input_file_id": "file-<your input file id>",
-            "completion_window": "24h",
-            "status": "validating",
-            "output_file_id": null,
-            "error_file_id": null,
-            "created_at": 1722037257,
-            "in_progress_at": null,
-            "expires_at": 1722123657,
-            "finalizing_at": null,
-            "completed_at": null,
-            "failed_at": null,
-            "expired_at": null,
-            "cancelling_at": null,
-            "cancelled_at": null,
-            "request_counts": {
-              "total": 0,
-              "completed": 0,
-              "failed": 0
-            },
-            "metadata": null
-          }
-        }
-      ],
-      "status_code": 200
-    }
-  ]
+  "task_id": "KYZSv5EBqL2d0mFvs80C",
+  "status": "CREATED"
 }
 ```
 
-For the definition of each field in the result, see [OpenAI Batch API](https://platform.openai.com/docs/guides/batch). Once the batch inference is complete, you can download the output by calling the [OpenAI Files API](https://platform.openai.com/docs/api-reference/files) and providing the file name specified in the `id` field of the response.
+To check the status of the batch predict job, provide the task ID to the [Tasks API]({{site.url}}{{site.baseurl}}/ml-commons-plugin/api/tasks-apis/get-task/). You can find the job details in the `remote_job` field in the task. Once the prediction is complete, the task `state` changes to `COMPLETED`.
+
+#### Example request
+
+```json
+GET /_plugins/_ml/tasks/KYZSv5EBqL2d0mFvs80C
+```
+{% include copy-curl.html %}
+
+#### Example response
+
+The response contains the batch predict operation details in the `remote_job` field:
+
+```json
+{
+  "model_id": "JYZRv5EBqL2d0mFvKs1E",
+  "task_type": "BATCH_PREDICTION",
+  "function_name": "REMOTE",
+  "state": "RUNNING",
+  "input_type": "REMOTE",
+  "worker_node": [
+    "Ee5OCIq0RAy05hqQsNI1rg"
+  ],
+  "create_time": 1725491751455,
+  "last_update_time": 1725491751455,
+  "is_async": false,
+  "remote_job": {
+    "cancelled_at": null,
+    "metadata": null,
+    "request_counts": {
+      "total": 3,
+      "completed": 3,
+      "failed": 0
+    },
+    "input_file_id": "file-XXXXXXXXXXXX",
+    "output_file_id": "file-XXXXXXXXXXXXX",
+    "error_file_id": null,
+    "created_at": 1725491753,
+    "in_progress_at": 1725491753,
+    "expired_at": null,
+    "finalizing_at": 1725491757,
+    "completed_at": null,
+    "endpoint": "/v1/embeddings",
+    "expires_at": 1725578153,
+    "cancelling_at": null,
+    "completion_window": "24h",
+    "id": "batch_XXXXXXXXXXXXXXX",
+    "failed_at": null,
+    "errors": null,
+    "object": "batch",
+    "status": "in_progress"
+  }
+}
+```
+
+For the definition of each field in the result, see [OpenAI Batch API](https://platform.openai.com/docs/guides/batch). Once the batch inference is complete, you can download the output by calling the [OpenAI Files API](https://platform.openai.com/docs/api-reference/files) and providing the file name specified in the `id` field of the response.
+
+### Canceling a batch predict job
+
+You can also cancel the batch predict operation running on the remote platform using the task ID returned by the batch predict request. To add this capability, set the `action_type` to `cancel_batch_predict` in the connector configuration when creating the connector.  
+
+#### Example request
+
+```json
+POST /_plugins/_ml/tasks/KYZSv5EBqL2d0mFvs80C/_cancel_batch
+```
+{% include copy-curl.html %}
+
+#### Example response
+
+```json
+{
+  "status": "OK"
+}
+```