Improve exception for trained model deployment scale up timeout (#128218)

* Improve exception for trained model deployment scale up timeout

* Update docs/changelog/128218.yaml

* Rename exception and update exception message

---------

Co-authored-by: Elastic Machine <[email protected]>

Author: dan-rubinstein

docs/changelog/128218.yaml

@@ -0,0 +1,5 @@
+pr: 128218
+summary: Improve exception for trained model deployment scale up timeout
+area: Machine Learning
+type: enhancement
+issues: []

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportStartTrainedModelDeploymentAction.java

@@ -64,6 +64,7 @@
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.core.ml.utils.TransportVersionUtils;
 import org.elasticsearch.xpack.ml.MachineLearning;
+import org.elasticsearch.xpack.ml.inference.assignment.ModelDeploymentTimeoutException;
 import org.elasticsearch.xpack.ml.inference.assignment.TrainedModelAssignmentService;
 import org.elasticsearch.xpack.ml.inference.persistence.TrainedModelDefinitionDoc;
 import org.elasticsearch.xpack.ml.notifications.InferenceAuditor;
@@ -350,11 +351,14 @@ public void onFailure(Exception e) {
                 @Override
                 public void onTimeout(TimeValue timeout) {
                     onFailure(
-                        new ElasticsearchStatusException(
-                            "Timed out after [{}] waiting for model deployment to start. "
-                                + "Use the trained model stats API to track the state of the deployment.",
-                            RestStatus.REQUEST_TIMEOUT,
-                            request.getTimeout() // use the full request timeout in the error message
+                        new ModelDeploymentTimeoutException(
+                            format(
+                                "Timed out after [%s] waiting for trained model deployment [%s] to start. "
+                                    + "Use the trained model stats API to track the state of the deployment "
+                                    + "and try again once it has started.",
+                                request.getTimeout(),
+                                request.getDeploymentId()
+                            )
                         )
                     );
                 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/InferenceWaitForAllocation.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.xpack.core.ml.action.InferModelAction;
@@ -20,6 +21,7 @@
 import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
 import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignmentMetadata;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
+import org.elasticsearch.xpack.ml.inference.assignment.ModelDeploymentTimeoutException;
 import org.elasticsearch.xpack.ml.inference.assignment.TrainedModelAssignmentService;
 
 import java.util.HashMap;
@@ -188,5 +190,19 @@ public void onFailure(Exception e) {
             pendingRequestCount.decrementAndGet();
             request.listener().onFailure(e);
         }
+
+        @Override
+        public void onTimeout(TimeValue timeout) {
+            onFailure(
+                new ModelDeploymentTimeoutException(
+                    format(
+                        "Timed out after [%s] waiting for trained model deployment [%s] to start. "
+                            + "Use the trained model stats API to track the state of the deployment and try again once it has started.",
+                        timeout,
+                        request.deploymentId()
+                    )
+                )
+            );
+        }
     }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/ModelDeploymentTimeoutException.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.ml.inference.assignment;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.rest.RestStatus;
+
+public class ModelDeploymentTimeoutException extends ElasticsearchStatusException {
+    public ModelDeploymentTimeoutException(String message) {
+        super(message, RestStatus.REQUEST_TIMEOUT);
+    }
+}

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentService.java

@@ -9,7 +9,6 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionType;
@@ -25,7 +24,6 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.node.NodeClosedException;
-import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.ConnectTransportException;
 import org.elasticsearch.xpack.core.ml.action.CreateTrainedModelAssignmentAction;
@@ -37,6 +35,7 @@
 import java.util.Objects;
 import java.util.function.Predicate;
 
+import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
 
 public class TrainedModelAssignmentService {
@@ -120,7 +119,15 @@ public void onTimeout(TimeValue timeout) {
 
     public interface WaitForAssignmentListener extends ActionListener<TrainedModelAssignment> {
         default void onTimeout(TimeValue timeout) {
-            onFailure(new ElasticsearchStatusException("Starting deployment timed out after [{}]", RestStatus.REQUEST_TIMEOUT, timeout));
+            onFailure(
+                new ModelDeploymentTimeoutException(
+                    format(
+                        "Timed out after [%s] waiting for trained model deployment to start. "
+                            + "Use the trained model stats API to track the state of the deployment and try again once it has started.",
+                        timeout
+                    )
+                )
+            );
         }
     }