Browse Source

[ML] Deploy default on chunked infer (#114158)

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
Mike Pellegrini 1 year ago
parent
commit
374f448722

+ 24 - 17
x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java

@@ -608,26 +608,33 @@ public class ElasticsearchInternalService extends BaseElasticsearchInternalServi
             return;
         }
 
-        var configUpdate = chunkingOptions != null
-            ? new TokenizationConfigUpdate(chunkingOptions.windowSize(), chunkingOptions.span())
-            : new TokenizationConfigUpdate(null, null);
+        if (model instanceof ElasticsearchInternalModel esModel) {
 
-        var request = buildInferenceRequest(
-            model.getConfigurations().getInferenceEntityId(),
-            configUpdate,
-            input,
-            inputType,
-            timeout,
-            true
-        );
+            var configUpdate = chunkingOptions != null
+                ? new TokenizationConfigUpdate(chunkingOptions.windowSize(), chunkingOptions.span())
+                : new TokenizationConfigUpdate(null, null);
+
+            var request = buildInferenceRequest(
+                model.getConfigurations().getInferenceEntityId(),
+                configUpdate,
+                input,
+                inputType,
+                timeout,
+                true
+            );
 
-        client.execute(
-            InferModelAction.INSTANCE,
-            request,
-            listener.delegateFailureAndWrap(
+            ActionListener<InferModelAction.Response> mlResultsListener = listener.delegateFailureAndWrap(
                 (l, inferenceResult) -> l.onResponse(translateToChunkedResults(inferenceResult.getInferenceResults()))
-            )
-        );
+            );
+
+            var maybeDeployListener = mlResultsListener.delegateResponse(
+                (l, exception) -> maybeStartDeployment(esModel, exception, request, mlResultsListener)
+            );
+
+            client.execute(InferModelAction.INSTANCE, request, maybeDeployListener);
+        } else {
+            listener.onFailure(notElasticsearchModelException(model));
+        }
     }
 
     private static List<ChunkedInferenceServiceResults> translateToChunkedResults(List<InferenceResults> inferenceResults) {