Configurable Inference timeout during Query time

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

@@ -145,6 +145,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -498,15 +499,10 @@ public static ExecutorBuilder<?> inferenceUtilityExecutor(Settings settings) {
 
     @Override
     public List<Setting<?>> getSettings() {
-        return List.copyOf(getInferenceSettingsInternal());
+        return List.copyOf(new ArrayList<>(getInferenceSettings()));
     }
 
-    // only used in tests
     public static Set<Setting<?>> getInferenceSettings() {
-        return Set.copyOf(getInferenceSettingsInternal());
-    }
-
-    private static Set<Setting<?>> getInferenceSettingsInternal() {
         Set<Setting<?>> settings = new HashSet<>();
         settings.addAll(HttpSettings.getSettingsDefinitions());
         settings.addAll(HttpClientManager.getSettingsDefinitions());
@@ -518,7 +514,7 @@ private static Set<Setting<?>> getInferenceSettingsInternal() {
         settings.add(INDICES_INFERENCE_BATCH_SIZE);
         settings.add(INFERENCE_QUERY_TIMEOUT);
         settings.addAll(ElasticInferenceServiceSettings.getSettingsDefinitions());
-        return settings;
+        return Collections.unmodifiableSet(settings);
     }
 
     @Override

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

@@ -22,7 +22,7 @@
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.inference.InferencePlugin;
 import org.elasticsearch.xpack.inference.services.settings.ApiKeySecrets;
@@ -1113,10 +1113,8 @@ public static TimeValue resolveInferenceTimeout(@Nullable TimeValue timeout, Inp
         if (timeout == null) {
             if (inputType == InputType.SEARCH || inputType == InputType.INTERNAL_SEARCH) {
                 return clusterService.getClusterSettings().get(InferencePlugin.INFERENCE_QUERY_TIMEOUT);
-            } else if (inputType == InputType.INGEST || inputType == InputType.INTERNAL_INGEST) {
-                return InferModelAction.Request.DEFAULT_TIMEOUT_FOR_INGEST;
             } else {
-                return InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API;
+                return InferenceAction.Request.DEFAULT_TIMEOUT;
             }
         }
         return timeout;

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

@@ -16,7 +16,7 @@
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.inference.InputType;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.inference.InferencePlugin;
 
@@ -1317,62 +1317,41 @@ public void testResolveInferenceTimeout_WithProvidedTimeout_ReturnsProvidedTimeo
         var clusterService = mockClusterService(Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), "10s").build());
         var providedTimeout = TimeValue.timeValueSeconds(45);
 
-        InputType[] inputTypes = {
-            InputType.INGEST,
-            InputType.INTERNAL_INGEST,
-            InputType.SEARCH,
-            InputType.INTERNAL_SEARCH,
-            InputType.CLASSIFICATION,
-            InputType.CLUSTERING,
-            InputType.UNSPECIFIED };
-
-        for (InputType inputType : inputTypes) {
+        for (InputType inputType : InputType.values()) {
             var result = ServiceUtils.resolveInferenceTimeout(providedTimeout, inputType, clusterService);
             assertEquals("Input type " + inputType + " should return provided timeout", providedTimeout, result);
         }
     }
 
-    public void testResolveInferenceTimeout_WithNullTimeoutAndSearchInputType_ReturnsClusterSetting() {
+    public void testResolveInferenceTimeout_WithNullTimeout_ReturnsExpectedTimeoutByInputType() {
         var configuredTimeout = TimeValue.timeValueSeconds(10);
         var clusterService = mockClusterService(
             Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), configuredTimeout).build()
         );
 
-        {
-            var result = ServiceUtils.resolveInferenceTimeout(null, InputType.SEARCH, clusterService);
-            assertEquals(configuredTimeout, result);
-        }
-        {
-            var result = ServiceUtils.resolveInferenceTimeout(null, InputType.INTERNAL_SEARCH, clusterService);
-            assertEquals(configuredTimeout, result);
-        }
-    }
-
-    public void testResolveInferenceTimeout_WithNullTimeoutAndIngestInputType_ReturnsMaxValue() {
-        var clusterService = mockClusterService(Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), "10s").build());
-
-        {
-            var result = ServiceUtils.resolveInferenceTimeout(null, InputType.INGEST, clusterService);
-            assertEquals(InferModelAction.Request.DEFAULT_TIMEOUT_FOR_INGEST, result);
-        }
-        {
-            var result = ServiceUtils.resolveInferenceTimeout(null, InputType.INTERNAL_INGEST, clusterService);
-            assertEquals(InferModelAction.Request.DEFAULT_TIMEOUT_FOR_INGEST, result);
-        }
-    }
-
-    public void testResolveInferenceTimeout_WithNullTimeoutAndOtherInputTypes_ReturnsDefaultTimeout() {
-        var clusterService = mockClusterService(Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), "10s").build());
+        Map<InputType, TimeValue> expectedTimeouts = Map.of(
+            InputType.SEARCH,
+            configuredTimeout,
+            InputType.INTERNAL_SEARCH,
+            configuredTimeout,
+            InputType.INGEST,
+            InferenceAction.Request.DEFAULT_TIMEOUT,
+            InputType.INTERNAL_INGEST,
+            InferenceAction.Request.DEFAULT_TIMEOUT,
+            InputType.CLASSIFICATION,
+            InferenceAction.Request.DEFAULT_TIMEOUT,
+            InputType.CLUSTERING,
+            InferenceAction.Request.DEFAULT_TIMEOUT,
+            InputType.UNSPECIFIED,
+            InferenceAction.Request.DEFAULT_TIMEOUT
+        );
 
-        InputType[] otherTypes = { InputType.CLASSIFICATION, InputType.CLUSTERING, InputType.UNSPECIFIED };
+        for (Map.Entry<InputType, TimeValue> entry : expectedTimeouts.entrySet()) {
+            InputType inputType = entry.getKey();
+            TimeValue expectedTimeout = entry.getValue();
 
-        for (InputType inputType : otherTypes) {
             var result = ServiceUtils.resolveInferenceTimeout(null, inputType, clusterService);
-            assertEquals(
-                "Input type " + inputType + " should return DEFAULT_TIMEOUT",
-                InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API,
-                result
-            );
+            assertEquals("Input type " + inputType + " should return expected timeout", expectedTimeout, result);
         }
     }
 }

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

@@ -1996,15 +1996,11 @@ private Client mockClientForStart(Consumer<ActionListener<CreateTrainedModelAssi
 
     @SuppressWarnings("unchecked")
     public void test_nullTimeoutUsesClusterSetting() throws InterruptedException {
-        var mlTrainedModelResults = new ArrayList<InferenceResults>();
-        mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
-        var response = new InferModelAction.Response(mlTrainedModelResults, "foo", true);
-
         Client client = mock(Client.class);
         when(client.threadPool()).thenReturn(threadPool);
         doAnswer(invocationOnMock -> {
             var listener = (ActionListener<InferModelAction.Response>) invocationOnMock.getArguments()[2];
-            listener.onResponse(response);
+            listener.onResponse(null);
             return null;
         }).when(client).execute(same(InferModelAction.INSTANCE), any(InferModelAction.Request.class), any(ActionListener.class));
 
@@ -2030,19 +2026,12 @@ public void test_nullTimeoutUsesClusterSetting() throws InterruptedException {
             null
         );
 
-        var gotResults = new AtomicBoolean();
-        var resultsListener = ActionListener.<InferenceServiceResults>wrap(serviceResponse -> {
-            assertThat(serviceResponse, instanceOf(TextEmbeddingFloatResults.class));
-            gotResults.set(true);
-        }, ESTestCase::fail);
-
         var latch = new CountDownLatch(1);
-        var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+        var latchedListener = new LatchedActionListener<>(ActionListener.<InferenceServiceResults>noop(), latch);
 
         service.infer(model, null, null, null, List.of("test input"), false, Map.of(), InputType.SEARCH, null, latchedListener);
 
-        latch.await();
-        assertTrue("Listener not called", gotResults.get());
+        assertTrue(latch.await(30, TimeUnit.SECONDS));
 
         ArgumentCaptor<InferModelAction.Request> requestCaptor = ArgumentCaptor.forClass(InferModelAction.Request.class);
         verify(client).execute(same(InferModelAction.INSTANCE), requestCaptor.capture(), any(ActionListener.class));
@@ -2051,15 +2040,11 @@ public void test_nullTimeoutUsesClusterSetting() throws InterruptedException {
 
     @SuppressWarnings("unchecked")
     public void test_providedTimeoutPropagateProperly() throws InterruptedException {
-        var mlTrainedModelResults = new ArrayList<InferenceResults>();
-        mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
-        var response = new InferModelAction.Response(mlTrainedModelResults, "foo", true);
-
         Client client = mock(Client.class);
         when(client.threadPool()).thenReturn(threadPool);
         doAnswer(invocationOnMock -> {
             var listener = (ActionListener<InferModelAction.Response>) invocationOnMock.getArguments()[2];
-            listener.onResponse(response);
+            listener.onResponse(null);
             return null;
         }).when(client).execute(same(InferModelAction.INSTANCE), any(InferModelAction.Request.class), any(ActionListener.class));
 
@@ -2085,19 +2070,12 @@ public void test_providedTimeoutPropagateProperly() throws InterruptedException
             null
         );
 
-        var gotResults = new AtomicBoolean();
-        var resultsListener = ActionListener.<InferenceServiceResults>wrap(serviceResponse -> {
-            assertThat(serviceResponse, instanceOf(TextEmbeddingFloatResults.class));
-            gotResults.set(true);
-        }, ESTestCase::fail);
-
         var latch = new CountDownLatch(1);
-        var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+        var latchedListener = new LatchedActionListener<>(ActionListener.<InferenceServiceResults>noop(), latch);
 
         service.infer(model, null, null, null, List.of("test input"), false, Map.of(), InputType.SEARCH, providedTimeout, latchedListener);
 
-        latch.await();
-        assertTrue("Listener not called", gotResults.get());
+        assertTrue(latch.await(30, TimeUnit.SECONDS));
 
         ArgumentCaptor<InferModelAction.Request> requestCaptor = ArgumentCaptor.forClass(InferModelAction.Request.class);
         verify(client).execute(same(InferModelAction.INSTANCE), requestCaptor.capture(), any(ActionListener.class));

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

@@ -12,12 +12,14 @@
 
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.LatchedActionListener;
 import org.elasticsearch.common.settings.SecureString;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.ChunkInferenceInput;
+import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.InputType;
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
@@ -42,6 +44,8 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
@@ -80,13 +84,14 @@ public class SageMakerServiceTests extends ESTestCase {
     private SageMakerClient client;
     private SageMakerSchemas schemas;
     private SageMakerService sageMakerService;
+    private ThreadPool threadPool;
 
     @Before
     public void init() {
         modelBuilder = mock();
         client = mock();
         schemas = mock();
-        ThreadPool threadPool = mock();
+        threadPool = mock();
         when(threadPool.executor(anyString())).thenReturn(EsExecutors.DIRECT_EXECUTOR_SERVICE);
         when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY));
         sageMakerService = new SageMakerService(modelBuilder, client, schemas, threadPool, Map::of, mockClusterServiceEmpty());
@@ -184,7 +189,7 @@ public void testInfer() {
     }
 
     @SuppressWarnings("unchecked")
-    public void test_nullTimeoutUsesClusterSetting() {
+    public void test_nullTimeoutUsesClusterSetting() throws InterruptedException {
         var model = mockModel();
         when(schemas.schemaFor(model)).thenReturn(mock());
 
@@ -193,22 +198,25 @@ public void test_nullTimeoutUsesClusterSetting() {
             Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), configuredTimeout).build()
         );
 
-        var service = new SageMakerService(modelBuilder, client, schemas, mock(ThreadPool.class), Map::of, clusterService);
+        var service = new SageMakerService(modelBuilder, client, schemas, threadPool, Map::of, clusterService);
 
         var capturedTimeout = new AtomicReference<TimeValue>();
         doAnswer(ans -> {
             capturedTimeout.set(ans.getArgument(2));
-            ((ActionListener<InvokeEndpointResponse>) ans.getArgument(3)).onResponse(InvokeEndpointResponse.builder().build());
+            ((ActionListener<InvokeEndpointResponse>) ans.getArgument(3)).onResponse(null);
             return null;
         }).when(client).invoke(any(), any(), any(), any());
 
-        service.infer(model, QUERY, null, null, INPUT, false, null, InputType.SEARCH, null, assertNoFailureListener(ignored -> {}));
+        var latch = new CountDownLatch(1);
+        var latchedListener = new LatchedActionListener<>(ActionListener.<InferenceServiceResults>noop(), latch);
+        service.infer(model, QUERY, null, null, INPUT, false, null, InputType.SEARCH, null, latchedListener);
 
+        assertTrue(latch.await(30, TimeUnit.SECONDS));
         assertEquals(configuredTimeout, capturedTimeout.get());
     }
 
     @SuppressWarnings("unchecked")
-    public void test_providedTimeoutPropagateProperly() {
+    public void test_providedTimeoutPropagateProperly() throws InterruptedException {
         var model = mockModel();
         when(schemas.schemaFor(model)).thenReturn(mock());
 
@@ -217,28 +225,20 @@ public void test_providedTimeoutPropagateProperly() {
             Settings.builder().put(InferencePlugin.INFERENCE_QUERY_TIMEOUT.getKey(), TimeValue.timeValueSeconds(15)).build()
         );
 
-        var service = new SageMakerService(modelBuilder, client, schemas, mock(ThreadPool.class), Map::of, clusterService);
+        var service = new SageMakerService(modelBuilder, client, schemas, threadPool, Map::of, clusterService);
 
         var capturedTimeout = new AtomicReference<TimeValue>();
         doAnswer(ans -> {
             capturedTimeout.set(ans.getArgument(2));
-            ((ActionListener<InvokeEndpointResponse>) ans.getArgument(3)).onResponse(InvokeEndpointResponse.builder().build());
+            ((ActionListener<InvokeEndpointResponse>) ans.getArgument(3)).onResponse(null);
             return null;
         }).when(client).invoke(any(), any(), any(), any());
 
-        service.infer(
-            model,
-            QUERY,
-            null,
-            null,
-            INPUT,
-            false,
-            null,
-            InputType.SEARCH,
-            providedTimeout,
-            assertNoFailureListener(ignored -> {})
-        );
+        var latch = new CountDownLatch(1);
+        var latchedListener = new LatchedActionListener<>(ActionListener.<InferenceServiceResults>noop(), latch);
+        service.infer(model, QUERY, null, null, INPUT, false, null, InputType.SEARCH, providedTimeout, latchedListener);
 
+        assertTrue(latch.await(30, TimeUnit.SECONDS));
         assertEquals(providedTimeout, capturedTimeout.get());
     }