Update throttling mechanism

Author: afoucret

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

@@ -25,12 +25,8 @@
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 
 public abstract class InferenceOperator extends AsyncOperator<InferenceOperator.OngoingInference> {
-
-    // Move to a setting.
-    private static final int MAX_INFERENCE_WORKER = 10;
     private final String inferenceId;
     private final BlockFactory blockFactory;
-    private final BulkInferenceExecutionConfig bulkExecutionConfig;
     private final BulkInferenceExecutor bulkInferenceExecutor;
 
     public InferenceOperator(
@@ -40,9 +36,8 @@ public InferenceOperator(
         ThreadPool threadPool,
         String inferenceId
     ) {
-        super(driverContext, threadPool.getThreadContext(), MAX_INFERENCE_WORKER);
+        super(driverContext, threadPool.getThreadContext(), bulkExecutionConfig.workers());
         this.blockFactory = driverContext.blockFactory();
-        this.bulkExecutionConfig = bulkExecutionConfig;
         this.bulkInferenceExecutor = new BulkInferenceExecutor(inferenceRunner, threadPool, bulkExecutionConfig);
         this.inferenceId = inferenceId;
     }
@@ -60,6 +55,17 @@ protected void releaseFetchedOnAnyThread(OngoingInference result) {
         releasePageOnAnyThread(result.inputPage);
     }
 
+    @Override
+    protected void performAsync(Page input, ActionListener<OngoingInference> listener) {
+        try {
+            BulkInferenceRequestIterator requests = requests(input);
+            listener = ActionListener.releaseAfter(listener, requests);
+            bulkInferenceExecutor.execute(requests, listener.map(responses -> new OngoingInference(input, responses)));
+        } catch (Exception e) {
+            listener.onFailure(e);
+        }
+    }
+
     @Override
     public Page getOutput() {
         OngoingInference ongoingInference = fetchFromBuffer();
@@ -75,19 +81,6 @@ public Page getOutput() {
         }
     }
 
-    @Override
-    protected void performAsync(Page input, ActionListener<OngoingInference> listener) {
-        try {
-            bulkInferenceExecutor.execute(requests(input), listener.map(responses -> new OngoingInference(input, responses)));
-        } catch (Exception e) {
-            listener.onFailure(e);
-        }
-    }
-
-    protected BulkInferenceExecutionConfig bulkExecutionConfig() {
-        return bulkExecutionConfig;
-    }
-
     protected abstract BulkInferenceRequestIterator requests(Page input);
 
     protected abstract OutputBuilder outputBuilder(Page input);

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

@@ -7,13 +7,9 @@
 
 package org.elasticsearch.xpack.esql.inference.bulk;
 
-import org.elasticsearch.core.TimeValue;
+public record BulkInferenceExecutionConfig(int workers, int maxOutstandingRequests) {
+    public static final int DEFAULT_WORKERS = 2;
+    public static final int DEFAULT_MAX_OUTSTANDING_REQUESTS = 50;
 
-import java.util.concurrent.TimeUnit;
-
-public record BulkInferenceExecutionConfig(TimeValue inferenceTimeout, int workers) {
-    public static final TimeValue DEFAULT_INFERENCE_TIMEOUT = new TimeValue(10, TimeUnit.SECONDS);
-    public static final int DEFAULT_WORKERS = 10;
-
-    public static final BulkInferenceExecutionConfig DEFAULT = new BulkInferenceExecutionConfig(DEFAULT_INFERENCE_TIMEOUT, DEFAULT_WORKERS);
+    public static final BulkInferenceExecutionConfig DEFAULT = new BulkInferenceExecutionConfig(DEFAULT_WORKERS, DEFAULT_MAX_OUTSTANDING_REQUESTS);
 }

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

@@ -7,16 +7,12 @@
 
 package org.elasticsearch.xpack.esql.inference.bulk;
 
-import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.compute.operator.FailureCollector;
 import org.elasticsearch.index.seqno.LocalCheckpointTracker;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 
 import java.util.Map;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.elasticsearch.index.seqno.SequenceNumbers.NO_OPS_PERFORMED;
@@ -25,7 +21,6 @@ public class BulkInferenceExecutionState {
     private final LocalCheckpointTracker checkpoint = new LocalCheckpointTracker(NO_OPS_PERFORMED, NO_OPS_PERFORMED);
     private final FailureCollector failureCollector = new FailureCollector();
     private final Map<Long, InferenceAction.Response> bufferedResponses = new ConcurrentHashMap<>();
-    private final BlockingQueue<Long> processedSeqNoQueue = ConcurrentCollections.newBlockingQueue();
     private final AtomicBoolean finished = new AtomicBoolean(false);
 
     public long generateSeqNo() {
@@ -49,28 +44,11 @@ public void onInferenceResponse(long seqNo, InferenceAction.Response response) {
             bufferedResponses.put(seqNo, response);
         }
         checkpoint.markSeqNoAsProcessed(seqNo);
-        processedSeqNoQueue.offer(seqNo);
     }
 
     public void onInferenceException(long seqNo, Exception e) {
         failureCollector.unwrapAndCollect(e);
         checkpoint.markSeqNoAsProcessed(seqNo);
-        processedSeqNoQueue.offer(seqNo);
-    }
-
-    public long fetchProcessedSeqNo(int retry) throws InterruptedException, TimeoutException {
-        while (retry > 0) {
-            if (finished()) {
-                return -1;
-            }
-            retry--;
-            Long seqNo = processedSeqNoQueue.poll(1, TimeUnit.SECONDS);
-            if (seqNo != null) {
-                return seqNo;
-            }
-        }
-
-        throw new TimeoutException("timeout waiting for inference response");
     }
 
     public InferenceAction.Response fetchBufferedResponse(long seqNo) {

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

@@ -8,110 +8,118 @@
 package org.elasticsearch.xpack.esql.inference.bulk;
 
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.common.util.concurrent.ThrottledTaskRunner;
+import org.elasticsearch.common.util.concurrent.AbstractRunnable;
+import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.esql.inference.InferenceRunner;
 import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class BulkInferenceExecutor {
-    private static final String TASK_RUNNER_NAME = "bulk_inference_operation";
-    private static final int INFERENCE_RESPONSE_TIMEOUT = 30; // TODO: should be in the config.
     private final ThrottledInferenceRunner throttledInferenceRunner;
-    private final ExecutorService executorService;
 
     public BulkInferenceExecutor(InferenceRunner inferenceRunner, ThreadPool threadPool, BulkInferenceExecutionConfig bulkExecutionConfig) {
-        executorService = executorService(threadPool);
-        throttledInferenceRunner = ThrottledInferenceRunner.create(inferenceRunner, executorService, bulkExecutionConfig);
+        throttledInferenceRunner = ThrottledInferenceRunner.create(inferenceRunner, executorService(threadPool), bulkExecutionConfig);
     }
 
     public void execute(BulkInferenceRequestIterator requests, ActionListener<List<InferenceAction.Response>> listener) throws Exception {
-        final ResponseHandler responseHandler = new ResponseHandler();
-        runInferenceRequests(requests, listener.delegateFailureAndWrap(responseHandler::handleResponses));
-    }
+        if (requests.hasNext() == false) {
+            listener.onResponse(List.of());
+            return;
+        }
 
-    private void runInferenceRequests(BulkInferenceRequestIterator requests, ActionListener<BulkInferenceExecutionState> listener) {
         final BulkInferenceExecutionState bulkExecutionState = new BulkInferenceExecutionState();
-        try {
-            executorService.execute(() -> {
-                while (bulkExecutionState.finished() == false && requests.hasNext()) {
-                    InferenceAction.Request request = requests.next();
-                    long seqNo = bulkExecutionState.generateSeqNo();
-                    throttledInferenceRunner.doInference(
-                        request,
-                        ActionListener.wrap(
-                            r -> bulkExecutionState.onInferenceResponse(seqNo, r),
-                            e -> bulkExecutionState.onInferenceException(seqNo, e)
-                        )
-                    );
-                }
+        final ResponseHandler responseHandler = new ResponseHandler(bulkExecutionState, listener, requests.estimatedSize());
+
+        while (bulkExecutionState.finished() == false && requests.hasNext()) {
+            InferenceAction.Request request = requests.next();
+            long seqNo = bulkExecutionState.generateSeqNo();
+
+            if (requests.hasNext() == false) {
                 bulkExecutionState.finish();
-            });
-        } catch (RejectedExecutionException e) {
-            bulkExecutionState.addFailure(new IllegalStateException("Unable to enqueue inference requests", e));
-            bulkExecutionState.finish();
-        } finally {
-            listener.onResponse(bulkExecutionState);
+            }
+
+            throttledInferenceRunner.doInference(
+                request,
+                ActionListener.runAfter(
+                    ActionListener.wrap(
+                        r -> bulkExecutionState.onInferenceResponse(seqNo, r),
+                        e -> bulkExecutionState.onInferenceException(seqNo, e)
+                    ),
+                    responseHandler::persistsInferenceResponses
+                )
+            );
         }
     }
 
     private static class ResponseHandler {
-        private final List<InferenceAction.Response> responses = new ArrayList<>();
-
-        private void handleResponses(ActionListener<List<InferenceAction.Response>> listener, BulkInferenceExecutionState bulkExecutionState) {
-
-            try {
-                persistsInferenceResponses(bulkExecutionState);
-            } catch (InterruptedException | TimeoutException e) {
-                bulkExecutionState.addFailure(e);
-                bulkExecutionState.finish();
-            }
+        private final List<InferenceAction.Response> responses;
+        private final ActionListener<List<InferenceAction.Response>> listener;
+        private final BulkInferenceExecutionState bulkExecutionState;
+        private final AtomicBoolean responseSent = new AtomicBoolean(false);
+
+        private ResponseHandler(
+            BulkInferenceExecutionState bulkExecutionState,
+            ActionListener<List<InferenceAction.Response>> listener,
+            int estimatedSize
+        ) {
+            this.listener = listener;
+            this.bulkExecutionState = bulkExecutionState;
+            this.responses = new ArrayList<>(estimatedSize);
+        }
 
-            if (bulkExecutionState.hasFailure() == false) {
-                try {
-                    listener.onResponse(responses);
-                    return;
-                } catch (Exception e) {
-                    bulkExecutionState.addFailure(e);
+        public synchronized void persistsInferenceResponses() {
+            long persistedSeqNo = bulkExecutionState.getPersistedCheckpoint();
+
+            while (persistedSeqNo < bulkExecutionState.getProcessedCheckpoint()) {
+                persistedSeqNo++;
+                InferenceAction.Response response = bulkExecutionState.fetchBufferedResponse(persistedSeqNo);
+                assert response != null || bulkExecutionState.hasFailure();
+                if (bulkExecutionState.hasFailure() == false) {
+                    try {
+                        responses.add(response);
+                    } catch (Exception e) {
+                        bulkExecutionState.addFailure(e);
+                    }
                 }
+                bulkExecutionState.markSeqNoAsPersisted(persistedSeqNo);
             }
 
-            listener.onFailure(bulkExecutionState.getFailure());
+            sendResponseOnCompletion();
         }
 
-        private void persistsInferenceResponses(BulkInferenceExecutionState bulkExecutionState) throws TimeoutException,
-            InterruptedException {
-            while (bulkExecutionState.finished() == false && bulkExecutionState.fetchProcessedSeqNo(INFERENCE_RESPONSE_TIMEOUT) >= 0) {
-                long persistedSeqNo = bulkExecutionState.getPersistedCheckpoint();
-
-                while (persistedSeqNo < bulkExecutionState.getProcessedCheckpoint()) {
-                    persistedSeqNo++;
-                    InferenceAction.Response response = bulkExecutionState.fetchBufferedResponse(persistedSeqNo);
-                    assert response != null || bulkExecutionState.hasFailure();
-                    if (bulkExecutionState.hasFailure() == false) {
-                        try {
-                            responses.add(response);
-                        } catch (Exception e) {
-                            bulkExecutionState.addFailure(e);
-                        }
+        private void sendResponseOnCompletion() {
+            if (bulkExecutionState.finished() && responseSent.compareAndSet(false, true)) {
+                if (bulkExecutionState.hasFailure() == false) {
+                    try {
+                        listener.onResponse(responses);
+                        return;
+                    } catch (Exception e) {
+                        bulkExecutionState.addFailure(e);
                     }
-                    bulkExecutionState.markSeqNoAsPersisted(persistedSeqNo);
                 }
+
+                listener.onFailure(bulkExecutionState.getFailure());
             }
         }
     }
 
-    private static class ThrottledInferenceRunner extends ThrottledTaskRunner {
+    private static class ThrottledInferenceRunner {
         private final InferenceRunner inferenceRunner;
+        private final ExecutorService executorService;
+        private final BlockingQueue<AbstractRunnable> pendingRequests = ConcurrentCollections.newBlockingQueue();
+        private final Semaphore permits;
 
         private ThrottledInferenceRunner(InferenceRunner inferenceRunner, ExecutorService executorService, int maxRunningTasks) {
-            super(TASK_RUNNER_NAME, maxRunningTasks, executorService);
+            this.executorService = executorService;
+            this.permits = new Semaphore(maxRunningTasks);
             this.inferenceRunner = inferenceRunner;
         }
 
@@ -120,13 +128,58 @@ public static ThrottledInferenceRunner create(
             ExecutorService executorService,
             BulkInferenceExecutionConfig bulkExecutionConfig
         ) {
-            return new ThrottledInferenceRunner(inferenceRunner, executorService, bulkExecutionConfig.workers());
+            return new ThrottledInferenceRunner(inferenceRunner, executorService, bulkExecutionConfig.maxOutstandingRequests());
         }
 
         public void doInference(InferenceAction.Request request, ActionListener<InferenceAction.Response> listener) {
-            this.enqueueTask(listener.delegateFailureAndWrap((l, releasable) -> {
-                inferenceRunner.doInference(request, ActionListener.releaseAfter(l, releasable));
-            }));
+            enqueueTask(request, listener);
+            executePendingRequests();
+        }
+
+        private void executePendingRequests() {
+            while (permits.tryAcquire()) {
+                AbstractRunnable task = pendingRequests.poll();
+
+                if (task == null) {
+                    permits.release();
+                    return;
+                }
+
+                try {
+                    executorService.execute(task);
+                } catch (Exception e){
+                    task.onFailure(e);
+                    permits.release();
+                }
+            }
+        }
+
+        private void enqueueTask(InferenceAction.Request request, ActionListener<InferenceAction.Response> listener) {
+            try {
+                pendingRequests.add(createTask(request, listener));
+                executePendingRequests();
+            } catch (Exception e) {
+                listener.onFailure(new IllegalStateException("An error occurred while adding the inference request to the queue", e));
+            }
+        }
+
+        private AbstractRunnable createTask(InferenceAction.Request request, ActionListener<InferenceAction.Response> listener) {
+            final ActionListener<InferenceAction.Response> completionListener = ActionListener.runAfter(listener, () -> {
+                permits.release();
+                executePendingRequests();
+            });
+
+            return new AbstractRunnable() {
+                @Override
+                protected void doRun() {
+                    inferenceRunner.doInference(request, completionListener);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    completionListener.onFailure(e);
+                }
+            };
         }
     }
 

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

@@ -7,10 +7,12 @@
 
 package org.elasticsearch.xpack.esql.inference.bulk;
 
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 
 import java.util.Iterator;
 
-public interface BulkInferenceRequestIterator extends Iterator<InferenceAction.Request> {
+public interface BulkInferenceRequestIterator extends Iterator<InferenceAction.Request>, Releasable {
+    int estimatedSize();
 
 }