Test concurrent inference on default endpoint

davidkyle · Nov 15, 2024 · 4605650 · 4605650
1 parent 7d98561
commit 4605650
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 1 deletion.
diff --git a/...ice-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/DefaultEndPointsIT.java b/...ice-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/DefaultEndPointsIT.java
@@ -8,6 +8,9 @@
 package org.elasticsearch.xpack.inference;
 
 import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.ResponseListener;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService;
@@ -16,9 +19,12 @@
 import org.junit.Before;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
 
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.oneOf;
@@ -110,4 +116,39 @@ private static void assertDefaultE5Config(Map<String, Object> modelConfig) {
             Matchers.is(Map.of("enabled", true, "min_number_of_allocations", 0, "max_number_of_allocations", 32))
         );
     }
+
+    public void testMultipleInferencesTiggeringDownloadAndDeploy() throws InterruptedException {
+        assumeTrue("Default config requires a feature flag", DefaultElserFeatureFlag.isEnabled());
+
+        int numParallelRequests = 4;
+        var latch = new CountDownLatch(numParallelRequests);
+        var errors = new ArrayList<Exception>();
+
+        var listener = new ResponseListener() {
+            @Override
+            public void onSuccess(Response response) {
+                latch.countDown();
+            }
+
+            @Override
+            public void onFailure(Exception exception) {
+                errors.add(exception);
+                latch.countDown();
+            }
+        };
+
+        var inputs = List.of("Hello World", "Goodnight moon");
+        var queryParams = Map.of("timeout", "120s");
+        for (int i = 0; i < numParallelRequests; i++) {
+            var request = createInferenceRequest(
+                Strings.format("_inference/%s", ElasticsearchInternalService.DEFAULT_ELSER_ID),
+                inputs,
+                queryParams
+            );
+            client().performRequestAsync(request, listener);
+        }
+
+        latch.await();
+        assertThat(errors.toString(), errors, empty());
+    }
 }
diff --git a/...-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java b/...-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java
@@ -373,12 +373,17 @@ protected Map<String, Object> infer(String modelId, TaskType taskType, List<Stri
         return inferInternal(endpoint, input, queryParameters);
     }
 
-    private Map<String, Object> inferInternal(String endpoint, List<String> input, Map<String, String> queryParameters) throws IOException {
+    protected Request createInferenceRequest(String endpoint, List<String> input, Map<String, String> queryParameters) {
         var request = new Request("POST", endpoint);
         request.setJsonEntity(jsonBody(input));
         if (queryParameters.isEmpty() == false) {
             request.addParameters(queryParameters);
         }
+        return request;
+    }
+
+    private Map<String, Object> inferInternal(String endpoint, List<String> input, Map<String, String> queryParameters) throws IOException {
+        var request = createInferenceRequest(endpoint, input, queryParameters);
         var response = client().performRequest(request);
         assertOkOrCreated(response);
         return entityAsMap(response);