From f0da131c8bdb27edf031a1bda1fcde5169562ced Mon Sep 17 00:00:00 2001
From: Quentin Pradet <quentin.pradet@elastic.co>
Date: Mon, 10 Jun 2024 16:50:13 +0400
Subject: [PATCH] GA the update trained model action (#2576)

* GA the update trained model action

* Fix lint
---
 output/schema/schema.json                     | 20 ++++++++++++++++---
 output/typescript/types.ts                    |  1 +
 .../ml.update_trained_model_deployment.json   | 11 ++++++++--
 .../MlUpdateTrainedModelDeploymentRequest.ts  | 12 +++++++++++
 4 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/output/schema/schema.json b/output/schema/schema.json
index 72d263caab..478edc203f 100644
--- a/output/schema/schema.json
+++ b/output/schema/schema.json
@@ -12956,7 +12956,7 @@
         "name": "Request",
         "namespace": "ml.update_trained_model_deployment"
       },
-      "requestBodyRequired": true,
+      "requestBodyRequired": false,
       "requestMediaType": [
         "application/json"
       ],
@@ -165793,8 +165793,22 @@
           }
         }
       ],
-      "query": [],
-      "specLocation": "ml/update_trained_model_deployment/MlUpdateTrainedModelDeploymentRequest.ts#L24-L50"
+      "query": [
+        {
+          "description": "The number of model allocations on each node where the model is deployed.\nAll allocations on a node share the same copy of the model in memory but use\na separate set of threads to evaluate the model.\nIncreasing this value generally increases the throughput.\nIf this setting is greater than the number of hardware threads\nit will automatically be changed to a value less than the number of hardware threads.",
+          "name": "number_of_allocations",
+          "required": false,
+          "serverDefault": 1,
+          "type": {
+            "kind": "instance_of",
+            "type": {
+              "name": "integer",
+              "namespace": "_types"
+            }
+          }
+        }
+      ],
+      "specLocation": "ml/update_trained_model_deployment/MlUpdateTrainedModelDeploymentRequest.ts#L24-L62"
     },
     {
       "body": {
diff --git a/output/typescript/types.ts b/output/typescript/types.ts
index 514fdba55f..7be4e13eab 100644
--- a/output/typescript/types.ts
+++ b/output/typescript/types.ts
@@ -15366,6 +15366,7 @@ export interface MlUpdateModelSnapshotResponse {
 
 export interface MlUpdateTrainedModelDeploymentRequest extends RequestBase {
   model_id: Id
+  number_of_allocations?: integer
   body?: {
     number_of_allocations?: integer
   }
diff --git a/specification/_json_spec/ml.update_trained_model_deployment.json b/specification/_json_spec/ml.update_trained_model_deployment.json
index 8ab14538f7..3fa5c31a0f 100644
--- a/specification/_json_spec/ml.update_trained_model_deployment.json
+++ b/specification/_json_spec/ml.update_trained_model_deployment.json
@@ -4,7 +4,7 @@
       "url": "https://www.elastic.co/guide/en/elasticsearch/reference/current/update-trained-model-deployment.html",
       "description": "Updates certain properties of trained model deployment."
     },
-    "stability": "beta",
+    "stability": "stable",
     "visibility": "public",
     "headers": {
       "accept": ["application/json"],
@@ -24,9 +24,16 @@
         }
       ]
     },
+    "params": {
+      "number_of_allocations": {
+        "type": "int",
+        "required": false,
+        "description": "Update the model deployment to this number of allocations."
+      }
+    },
     "body": {
       "description": "The updated trained model deployment settings",
-      "required": true
+      "required": false
     }
   }
 }
diff --git a/specification/ml/update_trained_model_deployment/MlUpdateTrainedModelDeploymentRequest.ts b/specification/ml/update_trained_model_deployment/MlUpdateTrainedModelDeploymentRequest.ts
index 4f5de2d821..bdd68b11d1 100644
--- a/specification/ml/update_trained_model_deployment/MlUpdateTrainedModelDeploymentRequest.ts
+++ b/specification/ml/update_trained_model_deployment/MlUpdateTrainedModelDeploymentRequest.ts
@@ -35,6 +35,18 @@ export interface Request extends RequestBase {
      */
     model_id: Id
   }
+  query_parameters: {
+    /**
+     * The number of model allocations on each node where the model is deployed.
+     * All allocations on a node share the same copy of the model in memory but use
+     * a separate set of threads to evaluate the model.
+     * Increasing this value generally increases the throughput.
+     * If this setting is greater than the number of hardware threads
+     * it will automatically be changed to a value less than the number of hardware threads.
+     * @server_default 1
+     */
+    number_of_allocations?: integer
+  }
   body: {
     /**
      * The number of model allocations on each node where the model is deployed.