[ML] Fix timeout attaching to missing deployment (elastic#115517)

Fixes a timeout in the Inference API where if connecting to an existing deployment and that deployment does not exist the listener was not called.
mark-vieira · Oct 24, 2024 · b2ab9df · b2ab9df
1 parent d5265be
commit b2ab9df
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/...tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/CreateFromDeploymentIT.java b/...tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/CreateFromDeploymentIT.java
@@ -109,6 +109,14 @@ public void testModelIdDoesNotMatch() throws IOException {
         );
     }
 
+    public void testDeploymentDoesNotExist() {
+        var deploymentId = "missing_deployment";
+
+        var inferenceId = "inference_on_missing_deployment";
+        var e = expectThrows(ResponseException.class, () -> putModel(inferenceId, endpointConfig(deploymentId), TaskType.SPARSE_EMBEDDING));
+        assertThat(e.getMessage(), containsString("Cannot find deployment [missing_deployment]"));
+    }
+
     public void testNumAllocationsIsUpdated() throws IOException {
         var modelId = "update_num_allocations";
         var deploymentId = modelId;

diff --git a/...rg/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java b/...rg/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java
@@ -34,7 +34,6 @@
 import org.elasticsearch.xpack.core.inference.results.SparseEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.action.GetDeploymentStatsAction;
 import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsAction;
-import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsStatsAction;
 import org.elasticsearch.xpack.core.ml.action.InferModelAction;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentStats;
@@ -913,7 +912,7 @@ private void validateAgainstDeployment(
                     listener.onFailure(
                         new ElasticsearchStatusException(
                             "Deployment [{}] uses model [{}] which does not match the model [{}] in the request.",
-                            RestStatus.BAD_REQUEST, // TODO better message
+                            RestStatus.BAD_REQUEST,
                             deploymentId,
                             response.get().getModelId(),
                             modelId
@@ -933,21 +932,22 @@ private void validateAgainstDeployment(
                 checkTaskTypeForMlNodeModel(response.get().getModelId(), taskType, l.delegateFailureAndWrap((l2, compatibleTaskType) -> {
                     l2.onResponse(updatedSettings);
                 }));
+            } else {
+                listener.onFailure(new ElasticsearchStatusException("Cannot find deployment [{}]", RestStatus.NOT_FOUND, deploymentId));
             }
         }));
     }
 
     private void getDeployment(String deploymentId, ActionListener<Optional<AssignmentStats>> listener) {
         client.execute(
-            GetTrainedModelsStatsAction.INSTANCE,
-            new GetTrainedModelsStatsAction.Request(deploymentId),
+            GetDeploymentStatsAction.INSTANCE,
+            new GetDeploymentStatsAction.Request(deploymentId),
             listener.delegateFailureAndWrap((l, response) -> {
                 l.onResponse(
-                    response.getResources()
+                    response.getStats()
                         .results()
                         .stream()
-                        .filter(s -> s.getDeploymentStats() != null && s.getDeploymentStats().getDeploymentId().equals(deploymentId))
-                        .map(GetTrainedModelsStatsAction.Response.TrainedModelStats::getDeploymentStats)
+                        .filter(s -> s.getDeploymentId() != null && s.getDeploymentId().equals(deploymentId))
                         .findFirst()
                 );
             })