Bulk inference refactoring.

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/BulkInferenceOperation.java

@@ -10,16 +10,18 @@
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefBuilder;
 import org.elasticsearch.compute.data.Block;
-import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.inference.results.ChatCompletionResults;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceOutputBuilder;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
 
 import java.util.List;
 import java.util.NoSuchElementException;
@@ -36,22 +38,27 @@ public String describe() {
 
         @Override
         public Operator get(DriverContext driverContext) {
-            return new CompletionOperator(driverContext, inferenceRunner, inferenceId, promptEvaluatorFactory.get(driverContext));
+            return new CompletionOperator(
+                driverContext,
+                inferenceRunner,
+                inferenceRunner.threadPool(),
+                inferenceId,
+                promptEvaluatorFactory.get(driverContext)
+            );
         }
     }
 
     private final ExpressionEvaluator promptEvaluator;
-    private final BlockFactory blockFactory;
 
     public CompletionOperator(
         DriverContext driverContext,
         InferenceRunner inferenceRunner,
+        ThreadPool threadPool,
         String inferenceId,
         ExpressionEvaluator promptEvaluator
     ) {
-        super(driverContext, inferenceRunner, inferenceId);
+        super(driverContext, inferenceRunner, threadPool, inferenceId);
         this.promptEvaluator = promptEvaluator;
-        this.blockFactory = driverContext.blockFactory();
     }
 
     @Override
@@ -65,8 +72,8 @@ public String toString() {
     }
 
     @Override
-    protected RequestIterator requests(Page inputPage) {
-        return new InferenceOperator.RequestIterator() {
+    protected BulkInferenceRequestIterator requests(Page inputPage) {
+        return new BulkInferenceRequestIterator() {
             private final BytesRefBlock promptBlock = (BytesRefBlock) promptEvaluator.eval(inputPage);
             private BytesRef readBuffer = new BytesRef();
             private int currentPos = 0;
@@ -105,9 +112,9 @@ public void close() {
     }
 
     @Override
-    protected OutputBuilder<ChatCompletionResults> outputBuilder(Page inputPage) {
-        return new InferenceOperator.OutputBuilder<>() {
-            private final BytesRefBlock.Builder outputBlockBuilder = blockFactory.newBytesRefBlockBuilder(inputPage.getPositionCount());
+    protected BulkInferenceOutputBuilder<ChatCompletionResults, Page> outputBuilder(Page inputPage) {
+        return new BulkInferenceOutputBuilder<>() {
+            private final BytesRefBlock.Builder outputBlockBuilder = blockFactory().newBytesRefBlockBuilder(inputPage.getPositionCount());
             private final BytesRefBuilder bytesRefBuilder = new BytesRefBuilder();
 
             @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/CompletionOperator.java

@@ -7,18 +7,28 @@
 
 package org.elasticsearch.xpack.esql.inference;
 
+import org.elasticsearch.core.TimeValue;
+
 import java.util.concurrent.ExecutorService;
 
 public class InferenceExecutionContext {
     private static final int DEFAULT_MAX_CONCURRENT_REQUESTS = 10;
+    private static final TimeValue DEFAULT_INFERENCE_EXECUTION_TIMEOUT = TimeValue.timeValueSeconds(10);
     private final InferenceRunner inferenceRunner;
     private final ExecutorService executorService;
     private final int maxConcurrentRequests;
+    private final TimeValue inferenceExecutionTimeout;
 
-    private InferenceExecutionContext(InferenceRunner inferenceRunner, ExecutorService executorService, int maxConcurrentRequests) {
+    private InferenceExecutionContext(
+        InferenceRunner inferenceRunner,
+        ExecutorService executorService,
+        int maxConcurrentRequests,
+        TimeValue inferenceExecutionTimeout
+    ) {
         this.inferenceRunner = inferenceRunner;
         this.executorService = executorService;
         this.maxConcurrentRequests = maxConcurrentRequests;
+        this.inferenceExecutionTimeout = inferenceExecutionTimeout;
     }
 
     public InferenceRunner inferenceRunner() {
@@ -33,23 +43,33 @@ public int maxConcurrentRequests() {
         return maxConcurrentRequests;
     }
 
+    public TimeValue inferenceExecutionTimeout() {
+        return inferenceExecutionTimeout;
+    }
+
     public static class Builder {
         private final InferenceRunner inferenceRunner;
         private final ExecutorService executorService;
         private int maxConcurrentRequests = DEFAULT_MAX_CONCURRENT_REQUESTS;
+        private TimeValue inferenceExecutionTimeout = DEFAULT_INFERENCE_EXECUTION_TIMEOUT;
 
         Builder(InferenceRunner inferenceRunner, ExecutorService executorService) {
             this.inferenceRunner = inferenceRunner;
             this.executorService = executorService;
         }
 
         public InferenceExecutionContext build() {
-            return new InferenceExecutionContext(inferenceRunner, executorService, maxConcurrentRequests);
+            return new InferenceExecutionContext(inferenceRunner, executorService, maxConcurrentRequests, inferenceExecutionTimeout);
         }
 
         public Builder setMaxConcurrentRequests(int maxConcurrentRequests) {
             this.maxConcurrentRequests = maxConcurrentRequests;
             return this;
         }
+
+        public Builder setInferenceExecutionTimeout(TimeValue inferenceExecutionTimeout) {
+            this.inferenceExecutionTimeout = inferenceExecutionTimeout;
+            return this;
+        }
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceExecutionContext.java

@@ -8,32 +8,38 @@
 package org.elasticsearch.xpack.esql.inference;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AsyncOperator;
 import org.elasticsearch.compute.operator.DriverContext;
-import org.elasticsearch.core.CheckedConsumer;
-import org.elasticsearch.core.Releasable;
-import org.elasticsearch.core.Releasables;
 import org.elasticsearch.inference.InferenceServiceResults;
-import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-
-import java.util.Iterator;
-
-import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceExecutionConfig;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceExecutor;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceOutputBuilder;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
 
 public abstract class InferenceOperator<InferenceResult extends InferenceServiceResults> extends AsyncOperator<Page> {
 
     // Move to a setting.
     private static final int MAX_INFERENCE_WORKER = 10;
-    private final InferenceRunner inferenceRunner;
     private final String inferenceId;
+    private final BlockFactory blockFactory;
 
-    public InferenceOperator(DriverContext driverContext, InferenceRunner inferenceRunner, String inferenceId) {
-        super(driverContext, inferenceRunner.threadContext(), MAX_INFERENCE_WORKER);
-        this.inferenceRunner = inferenceRunner;
+    private final BulkInferenceExecutor<InferenceResult, Page> bulkInferenceExecutor;
+
+    @SuppressWarnings("this-escape")
+    public InferenceOperator(DriverContext driverContext, InferenceRunner inferenceRunner, ThreadPool threadPool, String inferenceId) {
+        super(driverContext, threadPool.getThreadContext(), MAX_INFERENCE_WORKER);
+        this.blockFactory = driverContext.blockFactory();
+        this.bulkInferenceExecutor = new BulkInferenceExecutor<>(inferenceRunner, threadPool, bulkExecutionConfig());
         this.inferenceId = inferenceId;
     }
 
+    protected BlockFactory blockFactory() {
+        return blockFactory;
+    }
+
     protected String inferenceId() {
         return inferenceId;
     }
@@ -50,52 +56,14 @@ public Page getOutput() {
 
     @Override
     protected void performAsync(Page input, ActionListener<Page> listener) {
-        final RequestIterator requests = requests(input);
-        final OutputBuilder<InferenceResult> outputBuilder = outputBuilder(input);
-
-        new BulkInferenceOperation(requests, outputBuilder).execute(
-            inferenceExecutionContext(),
-            listener.delegateFailureIgnoreResponseAndWrap(l -> {
-                l.onResponse(outputBuilder.buildOutput());
-                Releasables.closeExpectNoException(requests, outputBuilder);
-            })
-        );
+        bulkInferenceExecutor.execute(requests(input), outputBuilder(input), listener);
     }
 
-    protected InferenceExecutionContext inferenceExecutionContext() {
-        return inferenceRunner.executionContextBuilder().build();
+    protected BulkInferenceExecutionConfig bulkExecutionConfig() {
+        return BulkInferenceExecutionConfig.DEFAULT;
     }
 
-    protected abstract RequestIterator requests(Page input);
-
-    protected abstract OutputBuilder<InferenceResult> outputBuilder(Page input);
-
-    public abstract static class OutputBuilder<InferenceResults extends InferenceServiceResults>
-        implements
-            CheckedConsumer<InferenceAction.Response, Exception>,
-            Releasable {
-        protected abstract Class<InferenceResults> inferenceResultsClass();
-
-        public abstract Page buildOutput();
-
-        public abstract void onInferenceResults(InferenceResults results);
-
-        @Override
-        public void accept(InferenceAction.Response response) throws Exception {
-            InferenceServiceResults results = response.getResults();
-            if (inferenceResultsClass().isInstance(response.getResults()) == false) {
-                throw new IllegalStateException(
-                    format(
-                        "Inference result has wrong type. Got [{}] while expecting [{}]",
-                        results.getClass().getName(),
-                        inferenceResultsClass().getName()
-                    )
-                );
-            }
-
-            onInferenceResults(inferenceResultsClass().cast(results));
-        }
-    }
+    protected abstract BulkInferenceRequestIterator requests(Page input);
 
-    public interface RequestIterator extends Iterator<InferenceAction.Request>, Releasable {}
+    protected abstract BulkInferenceOutputBuilder<InferenceResult, Page> outputBuilder(Page input);
 }