[ML] Adding asynchronous start up logic for the inference API internals

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSender.java

@@ -101,15 +101,39 @@ private HttpRequestSender(
     }
 
     /**
-     * Start various internal services. This is required before sending requests.
+     * Star various internal services asynchronously. This is required before sending requests.
      */
-    public void start() {
+    @Override
+    public void startAsynchronously(ActionListener<Void> listener) {
         if (started.compareAndSet(false, true)) {
+            threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> startInternal(listener));
+        } else {
+            listener.onResponse(null);
+        }
+    }
+
+    private void startInternal(ActionListener<Void> listener) {
+        try {
             // The manager must be started before the executor service. That way we guarantee that the http client
             // is ready prior to the service attempting to use the http client to send a request
             manager.start();
             threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(service::start);
             waitForStartToComplete();
+            listener.onResponse(null);
+        } catch (Exception ex) {
+            listener.onFailure(ex);
+        }
+    }
+
+    /**
+     * Start various internal services. This is required before sending requests.
+     *
+     * NOTE: This method blocks until the startup is complete.
+     */
+    @Override
+    public void startSynchronously() {
+        if (started.compareAndSet(false, true)) {
+            startInternal(ActionListener.noop());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/Sender.java

@@ -18,7 +18,9 @@
 import java.io.Closeable;
 
 public interface Sender extends Closeable {
-    void start();
+    void startSynchronously();
+
+    void startAsynchronously(ActionListener<Void> listener);
 
     void send(
         RequestManager requestCreator,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.core.IOUtils;
@@ -73,10 +74,11 @@ public void infer(
         @Nullable TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
-        timeout = ServiceUtils.resolveInferenceTimeout(timeout, inputType, clusterService);
-        init();
-        var inferenceInput = createInput(this, model, input, inputType, query, returnDocuments, topN, stream);
-        doInfer(model, inferenceInput, taskSettings, timeout, listener);
+        SubscribableListener.newForked(this::init).<InferenceServiceResults>andThen((inferListener) -> {
+            var resolvedInferenceTimeout = ServiceUtils.resolveInferenceTimeout(timeout, inputType, clusterService);
+            var inferenceInput = createInput(this, model, input, inputType, query, returnDocuments, topN, stream);
+            doInfer(model, inferenceInput, taskSettings, resolvedInferenceTimeout, inferListener);
+        }).addListener(listener);
     }
 
     private static InferenceInputs createInput(
@@ -121,8 +123,9 @@ public void unifiedCompletionInfer(
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
-        init();
-        doUnifiedCompletionInfer(model, new UnifiedChatInput(request, true), timeout, listener);
+        SubscribableListener.newForked(this::init).<InferenceServiceResults>andThen((completionInferListener) -> {
+            doUnifiedCompletionInfer(model, new UnifiedChatInput(request, true), timeout, completionInferListener);
+        }).addListener(listener);
     }
 
     @Override
@@ -135,16 +138,16 @@ public void chunkedInfer(
         TimeValue timeout,
         ActionListener<List<ChunkedInference>> listener
     ) {
-        init();
-
-        ValidationException validationException = new ValidationException();
-        validateInputType(inputType, model, validationException);
-        if (validationException.validationErrors().isEmpty() == false) {
-            throw validationException;
-        }
+        SubscribableListener.newForked(this::init).<List<ChunkedInference>>andThen((chunkedInferListener) -> {
+            ValidationException validationException = new ValidationException();
+            validateInputType(inputType, model, validationException);
+            if (validationException.validationErrors().isEmpty() == false) {
+                throw validationException;
+            }
 
-        // a non-null query is not supported and is dropped by all providers
-        doChunkedInfer(model, input, taskSettings, inputType, timeout, listener);
+            // a non-null query is not supported and is dropped by all providers
+            doChunkedInfer(model, input, taskSettings, inputType, timeout, chunkedInferListener);
+        }).addListener(listener);
     }
 
     protected abstract void doInfer(
@@ -176,8 +179,9 @@ protected abstract void doChunkedInfer(
     );
 
     public void start(Model model, ActionListener<Boolean> listener) {
-        init();
-        doStart(model, listener);
+        SubscribableListener.newForked(this::init)
+            .<Boolean>andThen((doStartListener) -> doStart(model, doStartListener))
+            .addListener(listener);
     }
 
     @Override
@@ -189,8 +193,8 @@ protected void doStart(Model model, ActionListener<Boolean> listener) {
         listener.onResponse(true);
     }
 
-    private void init() {
-        sender.start();
+    private void init(ActionListener<Void> listener) {
+        sender.startAsynchronously(listener);
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/client/AmazonBedrockRequestSender.java

@@ -75,7 +75,7 @@ public Factory(
 
         public Sender createSender() {
             // ensure this is started
-            bedrockRequestSender.start();
+            bedrockRequestSender.startSynchronously();
             return bedrockRequestSender;
         }
     }
@@ -97,8 +97,17 @@ protected AmazonBedrockRequestSender(
         this.startCompleted = Objects.requireNonNull(startCompleted);
     }
 
+    /**
+     * TODO implement this functionality to ensure that we don't block node bootups
+     * See: https://github.com/elastic/ml-team/issues/1701
+     */
     @Override
-    public void start() {
+    public void startAsynchronously(ActionListener<Void> listener) {
+        throw new UnsupportedOperationException("not implemented");
+    }
+
+    @Override
+    public void startSynchronously() {
         if (started.compareAndSet(false, true)) {
             // The manager must be started before the executor service. That way we guarantee that the http client
             // is ready prior to the service attempting to use the http client to send a request

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandler.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
@@ -82,37 +83,33 @@ public void getAuthorization(ActionListener<ElasticInferenceServiceAuthorization
                 return;
             }
 
-            // ensure that the sender is initialized
-            sender.start();
-
-            ActionListener<InferenceServiceResults> newListener = ActionListener.wrap(results -> {
-                if (results instanceof ElasticInferenceServiceAuthorizationResponseEntity authResponseEntity) {
-                    logger.debug(() -> Strings.format("Received authorization information from gateway %s", authResponseEntity));
-                    listener.onResponse(ElasticInferenceServiceAuthorizationModel.of(authResponseEntity));
-                } else {
-                    var errorMessage = Strings.format(
-                        "%s Received an invalid response type from the Elastic Inference Service: %s",
-                        FAILED_TO_RETRIEVE_MESSAGE,
-                        results.getClass().getSimpleName()
-                    );
-
-                    logger.warn(errorMessage);
-                    listener.onFailure(new ElasticsearchException(errorMessage));
-                }
-                requestCompleteLatch.countDown();
-            }, e -> {
+            var handleFailuresListener = listener.delegateResponse((authModelListener, e) -> {
                 // unwrap because it's likely a retry exception
                 var exception = ExceptionsHelper.unwrapCause(e);
 
                 logger.warn(Strings.format(FAILED_TO_RETRIEVE_MESSAGE + " Encountered an exception: %s", exception), exception);
-                listener.onFailure(e);
-                requestCompleteLatch.countDown();
+                authModelListener.onFailure(e);
             });
 
-            var requestMetadata = extractRequestMetadataFromThreadContext(threadPool.getThreadContext());
-            var request = new ElasticInferenceServiceAuthorizationRequest(baseUrl, getCurrentTraceInfo(), requestMetadata);
+            SubscribableListener.newForked(sender::startAsynchronously).<InferenceServiceResults>andThen((authListener) -> {
+                var requestMetadata = extractRequestMetadataFromThreadContext(threadPool.getThreadContext());
+                var request = new ElasticInferenceServiceAuthorizationRequest(baseUrl, getCurrentTraceInfo(), requestMetadata);
+                sender.sendWithoutQueuing(logger, request, AUTH_RESPONSE_HANDLER, DEFAULT_AUTH_TIMEOUT, authListener);
+            }).andThenApply(authResult -> {
+                if (authResult instanceof ElasticInferenceServiceAuthorizationResponseEntity authResponseEntity) {
+                    logger.debug(() -> Strings.format("Received authorization information from gateway %s", authResponseEntity));
+                    return ElasticInferenceServiceAuthorizationModel.of(authResponseEntity);
+                }
+
+                var errorMessage = Strings.format(
+                    "%s Received an invalid response type from the Elastic Inference Service: %s",
+                    FAILED_TO_RETRIEVE_MESSAGE,
+                    authResult.getClass().getSimpleName()
+                );
 
-            sender.sendWithoutQueuing(logger, request, AUTH_RESPONSE_HANDLER, DEFAULT_AUTH_TIMEOUT, newListener);
+                logger.warn(errorMessage);
+                throw new ElasticsearchException(errorMessage);
+            }).addListener(ActionListener.runAfter(handleFailuresListener, requestCompleteLatch::countDown));
         } catch (Exception e) {
             logger.warn(Strings.format("Retrieving the authorization information encountered an exception: %s", e));
             requestCompleteLatch.countDown();

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderTests.java

@@ -105,17 +105,17 @@ public void testCreateSender_CanCallStartMultipleTimes() throws Exception {
         var senderFactory = new HttpRequestSender.Factory(createWithEmptySettings(threadPool), clientManager, mockClusterServiceEmpty());
 
         try (var sender = createSender(senderFactory)) {
-            sender.start();
-            sender.start();
-            sender.start();
+            sender.startSynchronously();
+            sender.startSynchronously();
+            sender.startSynchronously();
         }
     }
 
     public void testCreateSender_SendsRequestAndReceivesResponse() throws Exception {
         var senderFactory = new HttpRequestSender.Factory(createWithEmptySettings(threadPool), clientManager, mockClusterServiceEmpty());
 
         try (var sender = createSender(senderFactory)) {
-            sender.start();
+            sender.startSynchronously();
 
             String responseJson = """
                 {
@@ -167,7 +167,7 @@ public void testSendWithoutQueuing_SendsRequestAndReceivesResponse() throws Exce
         var senderFactory = createSenderFactory(clientManager, threadRef);
 
         try (var sender = createSender(senderFactory)) {
-            sender.start();
+            sender.startSynchronously();
 
             String responseJson = """
                 {
@@ -240,7 +240,7 @@ public void testHttpRequestSender_Throws_WhenATimeoutOccurs() throws Exception {
 
         try (var sender = senderFactory.createSender()) {
             assertThat(sender, instanceOf(HttpRequestSender.class));
-            sender.start();
+            sender.startSynchronously();
 
             PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
             sender.send(RequestManagerTests.createMock(), new EmbeddingsInput(List.of(), null), TimeValue.timeValueNanos(1), listener);
@@ -263,7 +263,7 @@ public void testHttpRequestSenderWithTimeout_Throws_WhenATimeoutOccurs() throws
         );
 
         try (var sender = senderFactory.createSender()) {
-            sender.start();
+            sender.startSynchronously();
 
             PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
             sender.send(RequestManagerTests.createMock(), new EmbeddingsInput(List.of(), null), TimeValue.timeValueNanos(1), listener);
@@ -286,7 +286,7 @@ public void testSendWithoutQueuingWithTimeout_Throws_WhenATimeoutOccurs() throws
         );
 
         try (var sender = senderFactory.createSender()) {
-            sender.start();
+            sender.startSynchronously();
 
             PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
             sender.sendWithoutQueuing(

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ChatCompletionActionTests.java

@@ -108,7 +108,7 @@ public void testExecute_ThrowsException_WhenInputIsGreaterThanOne() throws IOExc
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
 
         try (var sender = createSender(senderFactory)) {
-            sender.start();
+            sender.startSynchronously();
 
             webServer.enqueue(new MockResponse().setResponseCode(200).setBody(getResponseJson()));
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/SenderServiceTests.java

@@ -43,6 +43,8 @@
 import static org.elasticsearch.xpack.inference.Utils.mockClusterService;
 import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty;
 import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -65,7 +67,7 @@ public void shutdown() throws IOException {
     }
 
     public void testStart_InitializesTheSender() throws IOException {
-        var sender = mock(Sender.class);
+        var sender = createMockSender();
 
         var factory = mock(HttpRequestSender.Factory.class);
         when(factory.createSender()).thenReturn(sender);
@@ -75,7 +77,7 @@ public void testStart_InitializesTheSender() throws IOException {
             service.start(mock(Model.class), listener);
 
             listener.actionGet(TIMEOUT);
-            verify(sender, times(1)).start();
+            verify(sender, times(1)).startAsynchronously(any());
             verify(factory, times(1)).createSender();
         }
 
@@ -85,7 +87,7 @@ public void testStart_InitializesTheSender() throws IOException {
     }
 
     public void testStart_CallingStartTwiceKeepsSameSenderReference() throws IOException {
-        var sender = mock(Sender.class);
+        var sender = createMockSender();
 
         var factory = mock(HttpRequestSender.Factory.class);
         when(factory.createSender()).thenReturn(sender);
@@ -95,11 +97,13 @@ public void testStart_CallingStartTwiceKeepsSameSenderReference() throws IOExcep
             service.start(mock(Model.class), listener);
             listener.actionGet(TIMEOUT);
 
-            service.start(mock(Model.class), listener);
-            listener.actionGet(TIMEOUT);
+
+            PlainActionFuture<Boolean> listener2 = new PlainActionFuture<>();
+            service.start(mock(Model.class), listener2);
+            listener2.actionGet(TIMEOUT);
 
             verify(factory, times(1)).createSender();
-            verify(sender, times(2)).start();
+            verify(sender, times(2)).startAsynchronously(any());
         }
 
         verify(sender, times(1)).close();
@@ -108,7 +112,8 @@ public void testStart_CallingStartTwiceKeepsSameSenderReference() throws IOExcep
     }
 
     public void test_nullTimeoutUsesClusterSetting() throws IOException {
-        var sender = mock(Sender.class);
+        var sender = createMockSender();
+
         var factory = mock(HttpRequestSender.Factory.class);
         when(factory.createSender()).thenReturn(sender);
 
@@ -147,7 +152,7 @@ protected void doInfer(
     }
 
     public void test_providedTimeoutPropagateProperly() throws IOException {
-        var sender = mock(Sender.class);
+        var sender = createMockSender();
         var factory = mock(HttpRequestSender.Factory.class);
         when(factory.createSender()).thenReturn(sender);
 
@@ -185,6 +190,18 @@ protected void doInfer(
         }
     }
 
+    public static Sender createMockSender() {
+        var sender = mock(Sender.class);
+        doAnswer(invocationOnMock -> {
+            ActionListener<Void> listener = invocationOnMock.getArgument(0);
+            listener.onResponse(null);
+            return Void.TYPE;
+        }).when(sender).startAsynchronously(any());
+
+        return sender;
+    }
+
+
     private static class TestSenderService extends SenderService {
         TestSenderService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents, ClusterService clusterService) {
             super(factory, serviceComponents, clusterService);