Improve exception for trained model deployment scale up timeout

Author: dan-rubinstein

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.ModelDeploymentScaleUpException;
 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 ModelDeploymentScaleUpException(
+                            format(
+                                "Timed out after [%s] waiting for trained model deployment [%s] to start. "
+                                    + "Please ensure the trained model deployment has started and try again.",
+                                request.getTimeout(),
+                                request.getDeploymentId()
+                            ),
+                            RestStatus.REQUEST_TIMEOUT
                         )
                     );
                 }

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.ModelDeploymentScaleUpException;
 import org.elasticsearch.xpack.ml.inference.assignment.TrainedModelAssignmentService;
 
 import java.util.HashMap;
@@ -188,5 +190,20 @@ public void onFailure(Exception e) {
             pendingRequestCount.decrementAndGet();
             request.listener().onFailure(e);
         }
+
+        @Override
+        public void onTimeout(TimeValue timeout) {
+            onFailure(
+                new ModelDeploymentScaleUpException(
+                    format(
+                        "Timed out after [%s] waiting for trained model deployment [%s] to start. "
+                            + "Please ensure the trained model deployment has started and try again.",
+                        timeout,
+                        request.deploymentId()
+                    ),
+                    RestStatus.REQUEST_TIMEOUT
+                )
+            );
+        }
     }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/ModelDeploymentScaleUpException.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 ModelDeploymentScaleUpException extends ElasticsearchStatusException {
+    public ModelDeploymentScaleUpException(String message, RestStatus status) {
+        super(message, status);
+    }
+}

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;
@@ -37,6 +36,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 +120,16 @@ 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 ModelDeploymentScaleUpException(
+                    format(
+                        "Timed out after [%s] waiting for trained model deployment to start. "
+                            + "Please ensure the trained model deployment has started and try again.",
+                        timeout
+                    ),
+                    RestStatus.REQUEST_TIMEOUT
+                )
+            );
         }
     }