Configurable Inference timeout during Query time

docs/changelog/131551.yaml

@@ -0,0 +1,5 @@
+pr: 131551
+summary: Added support to configure query timeout for inference
+area: Inference
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/inference/InferenceService.java

@@ -108,7 +108,9 @@ default boolean hideFromConfigurationApi() {
      * @param stream          Stream inference results
      * @param taskSettings    Settings in the request to override the model's defaults
      * @param inputType       For search, ingest etc
-     * @param timeout         The timeout for the request
+     * @param timeout         The timeout for the request. Callers should normally pass in a timeout.
+     *                        Passing in null is specifically for query builders who do not have access to the cluster settings
+     *                        to determine the appropriate timeout value set by the user within semantic_text.
      * @param listener        Inference result listener
      */
     void infer(
@@ -120,7 +122,7 @@ void infer(
         boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
-        TimeValue timeout,
+        @Nullable TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     );
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilder.java

@@ -31,7 +31,6 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.action.CoordinatedInferenceAction;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
 import org.elasticsearch.xpack.core.ml.inference.results.WarningInferenceResults;
@@ -279,7 +278,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
             List.of(query),
             TextExpansionConfigUpdate.EMPTY_UPDATE,
             false,
-            InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API
+            null
         );
         inferRequest.setHighPriority(true);
         inferRequest.setPrefixType(TrainedModelPrefixStrings.PrefixType.SEARCH);

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/vectors/TextEmbeddingQueryVectorBuilder.java

@@ -19,7 +19,6 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.action.CoordinatedInferenceAction;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings;
 import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
@@ -116,7 +115,7 @@ public void buildVector(Client client, ActionListener<float[]> listener) {
             List.of(modelText),
             TextEmbeddingConfigUpdate.EMPTY_INSTANCE,
             false,
-            InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API
+            null
         );
 
         inferRequest.setHighPriority(true);

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilderTests.java

@@ -124,7 +124,7 @@ protected boolean canSimulateMethod(Method method, Object[] args) throws NoSuchM
     @Override
     protected Object simulateMethod(Method method, Object[] args) {
         CoordinatedInferenceAction.Request request = (CoordinatedInferenceAction.Request) args[1];
-        assertEquals(InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API, request.getInferenceTimeout());
+        assertNull(request.getInferenceTimeout());
         assertEquals(TrainedModelPrefixStrings.PrefixType.SEARCH, request.getPrefixType());
         assertEquals(CoordinatedInferenceAction.Request.RequestModelType.NLP_MODEL, request.getRequestModelType());
 

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

@@ -148,6 +148,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -180,6 +181,12 @@ public class InferencePlugin extends Plugin
         Setting.Property.NodeScope,
         Setting.Property.Dynamic
     );
+    public static final Setting<TimeValue> INFERENCE_QUERY_TIMEOUT = Setting.timeSetting(
+        "xpack.inference.query_timeout",
+        TimeValue.timeValueSeconds(TimeUnit.SECONDS.toSeconds(10)),
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
 
     public static final LicensedFeature.Momentary INFERENCE_API_FEATURE = LicensedFeature.momentary(
         "inference",
@@ -499,6 +506,7 @@ public List<Setting<?>> getSettings() {
         settings.addAll(RequestExecutorServiceSettings.getSettingsDefinitions());
         settings.add(SKIP_VALIDATE_AND_START);
         settings.add(INDICES_INFERENCE_BATCH_SIZE);
+        settings.add(INFERENCE_QUERY_TIMEOUT);
         settings.addAll(ElasticInferenceServiceSettings.getSettingsDefinitions());
 
         return settings;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java

@@ -31,7 +31,6 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
 import org.elasticsearch.xpack.core.ml.inference.results.ErrorInferenceResults;
 import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
@@ -237,7 +236,7 @@ private SemanticQueryBuilder doRewriteGetInferenceResults(QueryRewriteContext qu
                 List.of(query),
                 Map.of(),
                 InputType.INTERNAL_SEARCH,
-                InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API,
+                null,
                 false
             );
 

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services;
+
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.InputType;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.inference.InferencePlugin;
+
+public class InferenceTimeoutUtils {
+    public static TimeValue resolveInferenceTimeout(@Nullable TimeValue timeout, InputType inputType, ClusterService clusterService) {
+        if (timeout == null) {
+            if (inputType == InputType.SEARCH || inputType == InputType.INTERNAL_SEARCH) {
+                return clusterService.getClusterSettings().get(InferencePlugin.INFERENCE_QUERY_TIMEOUT);
+            } else {
+                return InferenceAction.Request.DEFAULT_TIMEOUT;
+            }
+        }
+        return timeout;
+    }
+}

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

@@ -70,9 +70,10 @@ public void infer(
         boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
-        TimeValue timeout,
+        @Nullable TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
+        timeout = InferenceTimeoutUtils.resolveInferenceTimeout(timeout, inputType, clusterService);
         init();
         var chunkInferenceInput = input.stream().map(i -> new ChunkInferenceInput(i, null)).toList();
         var inferenceInput = createInput(this, model, chunkInferenceInput, inputType, query, returnDocuments, topN, stream);

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

@@ -299,6 +299,10 @@ private void preferredVariantFromPlatformArchitecture(ActionListener<PreferredMo
         );
     }
 
+    protected ClusterService getClusterService() {
+        return clusterService;
+    }
+
     boolean isClusterInElasticCloud() {
         // Use the ml lazy node count as a heuristic to determine if in Elastic cloud.
         // A value > 0 means scaling should be available for ml nodes

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

@@ -58,6 +58,7 @@
 import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
+import org.elasticsearch.xpack.inference.services.InferenceTimeoutUtils;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
 
 import java.util.ArrayList;
@@ -610,9 +611,10 @@ public void infer(
         boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
-        TimeValue timeout,
+        @Nullable TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
+        timeout = InferenceTimeoutUtils.resolveInferenceTimeout(timeout, inputType, getClusterService());
         if (model instanceof ElasticsearchInternalModel esModel) {
             var taskType = model.getConfigurations().getTaskType();
             if (TaskType.TEXT_EMBEDDING.equals(taskType)) {

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

@@ -31,6 +31,7 @@
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
+import org.elasticsearch.xpack.inference.services.InferenceTimeoutUtils;
 import org.elasticsearch.xpack.inference.services.sagemaker.model.SageMakerModel;
 import org.elasticsearch.xpack.inference.services.sagemaker.model.SageMakerModelBuilder;
 import org.elasticsearch.xpack.inference.services.sagemaker.schema.SageMakerSchemas;
@@ -160,7 +161,7 @@ public void infer(
             listener.onFailure(createInvalidModelException(model));
             return;
         }
-
+        timeout = InferenceTimeoutUtils.resolveInferenceTimeout(timeout, inputType, clusterService);
         var inferenceRequest = new SageMakerInferenceRequest(query, returnDocuments, topN, input, stream, inputType);
 
         try {
@@ -173,7 +174,7 @@ public void infer(
                 client.invokeStream(
                     regionAndSecrets,
                     request,
-                    timeout != null ? timeout : DEFAULT_TIMEOUT,
+                    timeout,
                     ActionListener.wrap(
                         response -> listener.onResponse(schema.streamResponse(sageMakerModel, response)),
                         e -> listener.onFailure(schema.error(sageMakerModel, e))
@@ -185,7 +186,7 @@ public void infer(
                 client.invoke(
                     regionAndSecrets,
                     request,
-                    timeout != null ? timeout : DEFAULT_TIMEOUT,
+                    timeout,
                     ActionListener.wrap(
                         response -> listener.onResponse(schema.response(sageMakerModel, response, threadPool.getThreadContext())),
                         e -> listener.onFailure(schema.error(sageMakerModel, e))

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

@@ -81,6 +81,8 @@ public static ClusterService mockClusterService(Settings settings) {
             ElasticInferenceServiceSettings.getSettingsDefinitions()
         ).flatMap(Collection::stream).collect(Collectors.toSet());
 
+        registeredSettings.add(InferencePlugin.INFERENCE_QUERY_TIMEOUT);
+
         var cSettings = new ClusterSettings(settings, registeredSettings);
         when(clusterService.getClusterSettings()).thenReturn(cSettings);
 

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services;
+
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.InputType;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.inference.InferencePlugin;
+import org.junit.Before;
+
+import static org.elasticsearch.xpack.inference.Utils.mockClusterService;
+
+public class InferenceTimeoutUtilsTests extends ESTestCase {
+
+    private ClusterService clusterService;
+    private static final TimeValue configuredTimeout = TimeValue.timeValueSeconds(10);
+
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+        var settings = Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), configuredTimeout).build();
+        clusterService = mockClusterService(settings);
+    }
+
+    public void testResolveInferenceTimeout_WithProvidedTimeout_ReturnsProvidedTimeout() {
+        var providedTimeout = TimeValue.timeValueSeconds(45);
+        {
+            var result = InferenceTimeoutUtils.resolveInferenceTimeout(providedTimeout, InputType.SEARCH, clusterService);
+            assertEquals(providedTimeout, result);
+        }
+        {
+            var result = InferenceTimeoutUtils.resolveInferenceTimeout(providedTimeout, InputType.INTERNAL_SEARCH, clusterService);
+            assertEquals(providedTimeout, result);
+        }
+        {
+            var result = InferenceTimeoutUtils.resolveInferenceTimeout(providedTimeout, InputType.INGEST, clusterService);
+            assertEquals(providedTimeout, result);
+        }
+    }
+
+    public void testResolveInferenceTimeout_WithNullTimeoutAndSearchInputType_ReturnsClusterSetting() {
+        {
+            var result = InferenceTimeoutUtils.resolveInferenceTimeout(null, InputType.SEARCH, clusterService);
+            assertEquals(configuredTimeout, result);
+        }
+        {
+            var result = InferenceTimeoutUtils.resolveInferenceTimeout(null, InputType.INTERNAL_SEARCH, clusterService);
+            assertEquals(configuredTimeout, result);
+        }
+    }
+
+    public void testResolveInferenceTimeout_WithNullTimeoutAndIngestInputType_ReturnsDefaultTimeout() {
+        var result = InferenceTimeoutUtils.resolveInferenceTimeout(null, InputType.INGEST, clusterService);
+        assertEquals(InferenceAction.Request.DEFAULT_TIMEOUT, result);
+    }
+}

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

@@ -12,6 +12,8 @@
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.ValidationException;
+import org.elasticsearch.common.settings.ClusterSettings;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.inference.InferenceServiceConfiguration;
@@ -21,6 +23,7 @@
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.InferencePlugin;
 import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
@@ -34,7 +37,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool;
 import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty;
@@ -103,7 +108,49 @@ public void testStart_CallingStartTwiceKeepsSameSenderReference() throws IOExcep
         verifyNoMoreInteractions(sender);
     }
 
-    private static final class TestSenderService extends SenderService {
+    public void test_nullTimeoutUsesClusterSetting() throws IOException {
+        var sender = mock(Sender.class);
+        var factory = mock(HttpRequestSender.Factory.class);
+        when(factory.createSender()).thenReturn(sender);
+
+        var configuredTimeout = TimeValue.timeValueSeconds(30);
+        var clusterSettings = new ClusterSettings(
+            Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), configuredTimeout).build(),
+            Set.of(InferencePlugin.INFERENCE_QUERY_TIMEOUT)
+        );
+        var clusterService = mock(ClusterService.class);
+        when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
+
+        var capturedTimeout = new AtomicReference<TimeValue>();
+        var testService = new TestSenderService(factory, createWithEmptySettings(threadPool), clusterService) {
+            // Override doInfer to capture the timeout value and return a mock response
+            @Override
+            protected void doInfer(
+                Model model,
+                InferenceInputs inputs,
+                Map<String, Object> taskSettings,
+                TimeValue timeout,
+                ActionListener<InferenceServiceResults> listener
+            ) {
+                capturedTimeout.set(timeout);
+                listener.onResponse(mock(InferenceServiceResults.class));
+            }
+        };
+
+        try (testService) {
+            var model = mock(Model.class);
+            when(model.getTaskType()).thenReturn(TaskType.TEXT_EMBEDDING);
+
+            PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
+
+            testService.infer(model, null, null, null, List.of("test input"), false, Map.of(), InputType.SEARCH, null, listener);
+
+            listener.actionGet(TIMEOUT);
+            assertEquals(configuredTimeout, capturedTimeout.get());
+        }
+    }
+
+    private static class TestSenderService extends SenderService {
         TestSenderService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents, ClusterService clusterService) {
             super(factory, serviceComponents, clusterService);
         }