waitForChe
return this;
}
+ /**
+ * Set the timeout for the {@link #setWaitForCheckpoints(Map)} request.
+ */
+ public SearchRequestBuilder setWaitForCheckpointsTimeout(final TimeValue waitForCheckpointsTimeout) {
+ request.setWaitForCheckpointsTimeout(waitForCheckpointsTimeout);
+ return this;
+ }
+
/**
* Specifies what type of requested indices to ignore and wildcard indices expressions.
*
diff --git a/server/src/main/resources/org/elasticsearch/TransportVersions.csv b/server/src/main/resources/org/elasticsearch/TransportVersions.csv
index b0ef5b780e775..26c518962c19a 100644
--- a/server/src/main/resources/org/elasticsearch/TransportVersions.csv
+++ b/server/src/main/resources/org/elasticsearch/TransportVersions.csv
@@ -131,3 +131,4 @@
8.15.1,8702002
8.15.2,8702003
8.15.3,8702003
+8.15.4,8702003
diff --git a/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv b/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv
index e3681cc975988..6cab0b513ee63 100644
--- a/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv
+++ b/server/src/main/resources/org/elasticsearch/index/IndexVersions.csv
@@ -131,3 +131,4 @@
8.15.1,8512000
8.15.2,8512000
8.15.3,8512000
+8.15.4,8512000
diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
index 6f6224d505327..25d9509ecdc7a 100644
--- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
+++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
@@ -10,14 +10,15 @@
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
-import org.elasticsearch.common.xcontent.support.XContentMapValues;
import org.elasticsearch.test.rest.ESRestTestCase;
import java.io.IOException;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.Set;
-import java.util.stream.Collectors;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.hasSize;
public class MlRestTestStateCleaner {
@@ -30,24 +31,29 @@ public MlRestTestStateCleaner(Logger logger, RestClient adminClient) {
}
public void resetFeatures() throws IOException {
- waitForMlStatsIndexToInitialize();
- deleteAllTrainedModelIngestPipelines();
+ deletePipelinesWithInferenceProcessors();
// This resets all features, not just ML, but they should have been getting reset between tests anyway so it shouldn't matter
adminClient.performRequest(new Request("POST", "/_features/_reset"));
}
@SuppressWarnings("unchecked")
- private void deleteAllTrainedModelIngestPipelines() throws IOException {
- final Request getAllTrainedModelStats = new Request("GET", "/_ml/trained_models/_stats");
- getAllTrainedModelStats.addParameter("size", "10000");
- final Response trainedModelsStatsResponse = adminClient.performRequest(getAllTrainedModelStats);
+ private void deletePipelinesWithInferenceProcessors() throws IOException {
+ final Response pipelinesResponse = adminClient.performRequest(new Request("GET", "/_ingest/pipeline"));
+ final Map pipelines = ESRestTestCase.entityAsMap(pipelinesResponse);
+
+ var pipelinesWithInferenceProcessors = new HashSet();
+ for (var entry : pipelines.entrySet()) {
+ var pipelineDef = (Map) entry.getValue(); // each top level object is a separate pipeline
+ var processors = (List>) pipelineDef.get("processors");
+ for (var processor : processors) {
+ assertThat(processor.entrySet(), hasSize(1));
+ if ("inference".equals(processor.keySet().iterator().next())) {
+ pipelinesWithInferenceProcessors.add(entry.getKey());
+ }
+ }
+ }
- final List> pipelines = (List>) XContentMapValues.extractValue(
- "trained_model_stats.ingest.pipelines",
- ESRestTestCase.entityAsMap(trainedModelsStatsResponse)
- );
- Set pipelineIds = pipelines.stream().flatMap(m -> m.keySet().stream()).collect(Collectors.toSet());
- for (String pipelineId : pipelineIds) {
+ for (String pipelineId : pipelinesWithInferenceProcessors) {
try {
adminClient.performRequest(new Request("DELETE", "/_ingest/pipeline/" + pipelineId));
} catch (Exception ex) {
@@ -55,12 +61,4 @@ private void deleteAllTrainedModelIngestPipelines() throws IOException {
}
}
}
-
- private void waitForMlStatsIndexToInitialize() throws IOException {
- ESRestTestCase.ensureHealth(adminClient, ".ml-stats-*", (request) -> {
- request.addParameter("wait_for_no_initializing_shards", "true");
- request.addParameter("level", "shards");
- request.addParameter("timeout", "30s");
- });
- }
}
diff --git a/x-pack/plugin/fleet/qa/rest/build.gradle b/x-pack/plugin/fleet/qa/rest/build.gradle
index fda9251c7ef34..dec624bc3cc56 100644
--- a/x-pack/plugin/fleet/qa/rest/build.gradle
+++ b/x-pack/plugin/fleet/qa/rest/build.gradle
@@ -1,8 +1,15 @@
-apply plugin: 'elasticsearch.legacy-yaml-rest-test'
+/*
+ * 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.
+ */
-dependencies {
- yamlRestTestImplementation(testArtifact(project(xpackModule('core'))))
-}
+import org.elasticsearch.gradle.internal.info.BuildParams
+
+apply plugin: 'elasticsearch.internal-yaml-rest-test'
+apply plugin: 'elasticsearch.yaml-rest-compat-test'
+apply plugin: 'elasticsearch.internal-test-artifact'
restResources {
restApi {
@@ -10,11 +17,17 @@ restResources {
}
}
-testClusters.configureEach {
- testDistribution = 'DEFAULT'
- setting 'xpack.security.enabled', 'true'
- setting 'xpack.license.self_generated.type', 'trial'
- extraConfigFile 'roles.yml', file('roles.yml')
- user username: 'elastic_admin', password: 'admin-password'
- user username: 'fleet_unprivileged_secrets', password: 'password', role: 'unprivileged_secrets'
+artifacts {
+ restXpackTests(new File(projectDir, "src/yamlRestTest/resources/rest-api-spec/test"))
+}
+
+tasks.named('yamlRestTest') {
+ usesDefaultDistribution()
+}
+tasks.named('yamlRestCompatTest') {
+ usesDefaultDistribution()
+}
+if (BuildParams.inFipsJvm){
+ // This test cluster is using a BASIC license and FIPS 140 mode is not supported in BASIC
+ tasks.named("yamlRestTest").configure{enabled = false }
}
diff --git a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java
index 202149abf11e1..bc49649bc1139 100644
--- a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java
+++ b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/java/org/elasticsearch/xpack/fleet/FleetRestIT.java
@@ -12,8 +12,12 @@
import org.elasticsearch.common.settings.SecureString;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.concurrent.ThreadContext;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.elasticsearch.test.cluster.local.distribution.DistributionType;
+import org.elasticsearch.test.cluster.util.resource.Resource;
import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
+import org.junit.ClassRule;
public class FleetRestIT extends ESClientYamlSuiteTestCase {
@@ -21,14 +25,30 @@ public FleetRestIT(final ClientYamlTestCandidate testCandidate) {
super(testCandidate);
}
+ @ClassRule
+ public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
+ .distribution(DistributionType.DEFAULT)
+ .setting("xpack.license.self_generated.type", "basic")
+ .setting("xpack.security.enabled", "true")
+ .rolesFile(Resource.fromClasspath("roles.yml"))
+ .user("elastic_admin", "admin-password", "superuser", true)
+ .user("fleet_unprivileged_secrets", "password", "unprivileged_secrets", true)
+ .build();
+
+ @Override
+ protected String getTestRestCluster() {
+ return cluster.getHttpAddresses();
+ }
+
@Override
protected Settings restClientSettings() {
- String authentication = basicAuthHeaderValue("elastic_admin", new SecureString("admin-password".toCharArray()));
- return Settings.builder().put(super.restClientSettings()).put(ThreadContext.PREFIX + ".Authorization", authentication).build();
+ String token = basicAuthHeaderValue("elastic_admin", new SecureString("admin-password".toCharArray()));
+ return Settings.builder().put(ThreadContext.PREFIX + ".Authorization", token).build();
}
@ParametersFactory
public static Iterable parameters() throws Exception {
return ESClientYamlSuiteTestCase.createParameters();
}
+
}
diff --git a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml
index 5610502a65d23..4c168c8feb0cd 100644
--- a/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml
+++ b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/fleet/20_wait_for_checkpoints.yml
@@ -105,6 +105,7 @@ setup:
index: "test-after-refresh"
allow_partial_search_results: false
wait_for_checkpoints: 2
+ wait_for_checkpoints_timeout: 1m
body: { query: { match_all: {} } }
---
@@ -115,7 +116,7 @@ setup:
body:
- { "allow_partial_search_results": false, wait_for_checkpoints: 1 }
- { query: { match_all: { } } }
- - { "allow_partial_search_results": false, wait_for_checkpoints: 2 }
+ - { "allow_partial_search_results": false, wait_for_checkpoints: 2, wait_for_checkpoints_timeout: 1m }
- { query: { match_all: { } } }
- match: { responses.0._shards.successful: 1 }
@@ -128,7 +129,7 @@ setup:
- {query: { match_all: {} } }
- { "index": "test-alias", "allow_partial_search_results": false, wait_for_checkpoints: 1 }
- { query: { match_all: { } } }
- - {"index": "test-refresh-disabled", "allow_partial_search_results": false, wait_for_checkpoints: 2}
+ - { "index": "test-refresh-disabled", "allow_partial_search_results": false, wait_for_checkpoints: 2, wait_for_checkpoints_timeout: 1m }
- {query: { match_all: {} } }
- match: { responses.0._shards.successful: 1 }
diff --git a/x-pack/plugin/fleet/qa/rest/roles.yml b/x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/roles.yml
similarity index 100%
rename from x-pack/plugin/fleet/qa/rest/roles.yml
rename to x-pack/plugin/fleet/qa/rest/src/yamlRestTest/resources/roles.yml
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java
index 83249266c79ab..fe83acc8574aa 100644
--- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java
+++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java
@@ -68,6 +68,7 @@
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Stream;
@@ -680,25 +681,13 @@ public void chunkedInfer(
esModel.getConfigurations().getChunkingSettings()
).batchRequestsWithListeners(listener);
- for (var batch : batchedRequests) {
- var inferenceRequest = buildInferenceRequest(
- esModel.mlNodeDeploymentId(),
- EmptyConfigUpdate.INSTANCE,
- batch.batch().inputs(),
- inputType,
- timeout
- );
-
- ActionListener mlResultsListener = batch.listener()
- .delegateFailureAndWrap(
- (l, inferenceResult) -> translateToChunkedResult(model.getTaskType(), inferenceResult.getInferenceResults(), l)
- );
-
- var maybeDeployListener = mlResultsListener.delegateResponse(
- (l, exception) -> maybeStartDeployment(esModel, exception, inferenceRequest, mlResultsListener)
- );
-
- client.execute(InferModelAction.INSTANCE, inferenceRequest, maybeDeployListener);
+ if (batchedRequests.isEmpty()) {
+ listener.onResponse(List.of());
+ } else {
+ // Avoid filling the inference queue by executing the batches in series
+ // Each batch contains up to EMBEDDING_MAX_BATCH_SIZE inference request
+ var sequentialRunner = new BatchIterator(esModel, inputType, timeout, batchedRequests);
+ sequentialRunner.run();
}
} else {
listener.onFailure(notElasticsearchModelException(model));
@@ -1018,6 +1007,82 @@ static TaskType inferenceConfigToTaskType(InferenceConfig config) {
}
}
+ /**
+ * Iterates over the batch executing a limited number requests at a time to avoid
+ * filling the ML node inference queue.
+ *
+ * First, a single request is executed, which can also trigger deploying a model
+ * if necessary. When this request is successfully executed, a callback executes
+ * N requests in parallel next. Each of these requests also has a callback that
+ * executes one more request, so that at all time N requests are in-flight. This
+ * continues until all requests are executed.
+ */
+ class BatchIterator {
+ private static final int NUM_REQUESTS_INFLIGHT = 20; // * batch size = 200
+
+ private final AtomicInteger index = new AtomicInteger();
+ private final ElasticsearchInternalModel esModel;
+ private final List requestAndListeners;
+ private final InputType inputType;
+ private final TimeValue timeout;
+
+ BatchIterator(
+ ElasticsearchInternalModel esModel,
+ InputType inputType,
+ TimeValue timeout,
+ List requestAndListeners
+ ) {
+ this.esModel = esModel;
+ this.requestAndListeners = requestAndListeners;
+ this.inputType = inputType;
+ this.timeout = timeout;
+ }
+
+ void run() {
+ // The first request may deploy the model, and upon completion runs
+ // NUM_REQUESTS_INFLIGHT in parallel.
+ inferenceExecutor.execute(() -> inferBatch(NUM_REQUESTS_INFLIGHT, true));
+ }
+
+ private void inferBatch(int runAfterCount, boolean maybeDeploy) {
+ int batchIndex = index.getAndIncrement();
+ if (batchIndex >= requestAndListeners.size()) {
+ return;
+ }
+ executeRequest(batchIndex, maybeDeploy, () -> {
+ for (int i = 0; i < runAfterCount; i++) {
+ // Subsequent requests may not deploy the model, because the first request
+ // already did so. Upon completion, it runs one more request.
+ inferenceExecutor.execute(() -> inferBatch(1, false));
+ }
+ });
+ }
+
+ private void executeRequest(int batchIndex, boolean maybeDeploy, Runnable runAfter) {
+ EmbeddingRequestChunker.BatchRequestAndListener batch = requestAndListeners.get(batchIndex);
+ var inferenceRequest = buildInferenceRequest(
+ esModel.mlNodeDeploymentId(),
+ EmptyConfigUpdate.INSTANCE,
+ batch.batch().inputs(),
+ inputType,
+ timeout
+ );
+ logger.trace("Executing batch index={}", batchIndex);
+
+ ActionListener listener = batch.listener()
+ .delegateFailureAndWrap(
+ (l, inferenceResult) -> translateToChunkedResult(esModel.getTaskType(), inferenceResult.getInferenceResults(), l)
+ );
+ if (runAfter != null) {
+ listener = ActionListener.runAfter(listener, runAfter);
+ }
+ if (maybeDeploy) {
+ listener = listener.delegateResponse((l, exception) -> maybeStartDeployment(esModel, exception, inferenceRequest, l));
+ }
+ client.execute(InferModelAction.INSTANCE, inferenceRequest, listener);
+ }
+ }
+
public static class Configuration {
public static InferenceServiceConfiguration get() {
return configuration.getOrCompute();
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java
index c1be537a6b0a7..4fdf254101d3e 100644
--- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java
+++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java
@@ -24,12 +24,25 @@
import java.util.concurrent.atomic.AtomicReference;
import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.empty;
import static org.hamcrest.Matchers.hasSize;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.startsWith;
public class EmbeddingRequestChunkerTests extends ESTestCase {
+ public void testEmptyInput() {
+ var embeddingType = randomFrom(EmbeddingRequestChunker.EmbeddingType.values());
+ var batches = new EmbeddingRequestChunker(List.of(), 100, 100, 10, embeddingType).batchRequestsWithListeners(testListener());
+ assertThat(batches, empty());
+ }
+
+ public void testBlankInput() {
+ var embeddingType = randomFrom(EmbeddingRequestChunker.EmbeddingType.values());
+ var batches = new EmbeddingRequestChunker(List.of(""), 100, 100, 10, embeddingType).batchRequestsWithListeners(testListener());
+ assertThat(batches, hasSize(1));
+ }
+
public void testShortInputsAreSingleBatch() {
String input = "one chunk";
var embeddingType = randomFrom(EmbeddingRequestChunker.EmbeddingType.values());
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java
index 89a27a921cbea..9a4d0dda82238 100644
--- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java
+++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java
@@ -12,6 +12,7 @@
import org.apache.logging.log4j.Level;
import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.LatchedActionListener;
import org.elasticsearch.action.support.PlainActionFuture;
import org.elasticsearch.client.internal.Client;
import org.elasticsearch.cluster.service.ClusterService;
@@ -65,6 +66,7 @@
import org.elasticsearch.xpack.inference.InferencePlugin;
import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsTests;
import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
+import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings;
import org.elasticsearch.xpack.inference.services.ServiceFields;
import org.junit.After;
import org.junit.Before;
@@ -72,12 +74,14 @@
import org.mockito.Mockito;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
@@ -832,16 +836,16 @@ public void testParsePersistedConfig() {
}
}
- public void testChunkInfer_E5WithNullChunkingSettings() {
+ public void testChunkInfer_E5WithNullChunkingSettings() throws InterruptedException {
testChunkInfer_e5(null);
}
- public void testChunkInfer_E5ChunkingSettingsSet() {
+ public void testChunkInfer_E5ChunkingSettingsSet() throws InterruptedException {
testChunkInfer_e5(ChunkingSettingsTests.createRandomChunkingSettings());
}
@SuppressWarnings("unchecked")
- private void testChunkInfer_e5(ChunkingSettings chunkingSettings) {
+ private void testChunkInfer_e5(ChunkingSettings chunkingSettings) throws InterruptedException {
var mlTrainedModelResults = new ArrayList();
mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
@@ -889,6 +893,9 @@ private void testChunkInfer_e5(ChunkingSettings chunkingSettings) {
gotResults.set(true);
}, ESTestCase::fail);
+ var latch = new CountDownLatch(1);
+ var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+
service.chunkedInfer(
model,
null,
@@ -897,22 +904,23 @@ private void testChunkInfer_e5(ChunkingSettings chunkingSettings) {
InputType.SEARCH,
new ChunkingOptions(null, null),
InferenceAction.Request.DEFAULT_TIMEOUT,
- ActionListener.runAfter(resultsListener, () -> terminate(threadPool))
+ latchedListener
);
+ latch.await();
assertTrue("Listener not called", gotResults.get());
}
- public void testChunkInfer_SparseWithNullChunkingSettings() {
+ public void testChunkInfer_SparseWithNullChunkingSettings() throws InterruptedException {
testChunkInfer_Sparse(null);
}
- public void testChunkInfer_SparseWithChunkingSettingsSet() {
+ public void testChunkInfer_SparseWithChunkingSettingsSet() throws InterruptedException {
testChunkInfer_Sparse(ChunkingSettingsTests.createRandomChunkingSettings());
}
@SuppressWarnings("unchecked")
- private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) {
+ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) throws InterruptedException {
var mlTrainedModelResults = new ArrayList();
mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults());
mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults());
@@ -936,6 +944,7 @@ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) {
var service = createService(client);
var gotResults = new AtomicBoolean();
+
var resultsListener = ActionListener.>wrap(chunkedResponse -> {
assertThat(chunkedResponse, hasSize(2));
assertThat(chunkedResponse.get(0), instanceOf(InferenceChunkedSparseEmbeddingResults.class));
@@ -955,6 +964,9 @@ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) {
gotResults.set(true);
}, ESTestCase::fail);
+ var latch = new CountDownLatch(1);
+ var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+
service.chunkedInfer(
model,
null,
@@ -963,22 +975,23 @@ private void testChunkInfer_Sparse(ChunkingSettings chunkingSettings) {
InputType.SEARCH,
new ChunkingOptions(null, null),
InferenceAction.Request.DEFAULT_TIMEOUT,
- ActionListener.runAfter(resultsListener, () -> terminate(threadPool))
+ latchedListener
);
+ latch.await();
assertTrue("Listener not called", gotResults.get());
}
- public void testChunkInfer_ElserWithNullChunkingSettings() {
+ public void testChunkInfer_ElserWithNullChunkingSettings() throws InterruptedException {
testChunkInfer_Elser(null);
}
- public void testChunkInfer_ElserWithChunkingSettingsSet() {
+ public void testChunkInfer_ElserWithChunkingSettingsSet() throws InterruptedException {
testChunkInfer_Elser(ChunkingSettingsTests.createRandomChunkingSettings());
}
@SuppressWarnings("unchecked")
- private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) {
+ private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) throws InterruptedException {
var mlTrainedModelResults = new ArrayList();
mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults());
mlTrainedModelResults.add(TextExpansionResultsTests.createRandomResults());
@@ -1022,6 +1035,9 @@ private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) {
gotResults.set(true);
}, ESTestCase::fail);
+ var latch = new CountDownLatch(1);
+ var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+
service.chunkedInfer(
model,
null,
@@ -1030,9 +1046,10 @@ private void testChunkInfer_Elser(ChunkingSettings chunkingSettings) {
InputType.SEARCH,
new ChunkingOptions(null, null),
InferenceAction.Request.DEFAULT_TIMEOUT,
- ActionListener.runAfter(resultsListener, () -> terminate(threadPool))
+ latchedListener
);
+ latch.await();
assertTrue("Listener not called", gotResults.get());
}
@@ -1093,7 +1110,7 @@ public void testChunkInferSetsTokenization() {
}
@SuppressWarnings("unchecked")
- public void testChunkInfer_FailsBatch() {
+ public void testChunkInfer_FailsBatch() throws InterruptedException {
var mlTrainedModelResults = new ArrayList();
mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
@@ -1129,6 +1146,9 @@ public void testChunkInfer_FailsBatch() {
gotResults.set(true);
}, ESTestCase::fail);
+ var latch = new CountDownLatch(1);
+ var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+
service.chunkedInfer(
model,
null,
@@ -1137,12 +1157,86 @@ public void testChunkInfer_FailsBatch() {
InputType.SEARCH,
new ChunkingOptions(null, null),
InferenceAction.Request.DEFAULT_TIMEOUT,
- ActionListener.runAfter(resultsListener, () -> terminate(threadPool))
+ latchedListener
);
+ latch.await();
assertTrue("Listener not called", gotResults.get());
}
+ @SuppressWarnings("unchecked")
+ public void testChunkingLargeDocument() throws InterruptedException {
+ int numBatches = randomIntBetween(3, 6);
+
+ // how many response objects to return in each batch
+ int[] numResponsesPerBatch = new int[numBatches];
+ for (int i = 0; i < numBatches - 1; i++) {
+ numResponsesPerBatch[i] = ElasticsearchInternalService.EMBEDDING_MAX_BATCH_SIZE;
+ }
+ numResponsesPerBatch[numBatches - 1] = randomIntBetween(1, ElasticsearchInternalService.EMBEDDING_MAX_BATCH_SIZE);
+ int numChunks = Arrays.stream(numResponsesPerBatch).sum();
+
+ // build a doc with enough words to make numChunks of chunks
+ int wordsPerChunk = 10;
+ int numWords = numChunks * wordsPerChunk;
+ var input = "word ".repeat(numWords);
+
+ Client client = mock(Client.class);
+ when(client.threadPool()).thenReturn(threadPool);
+
+ // mock the inference response
+ doAnswer(invocationOnMock -> {
+ var request = (InferModelAction.Request) invocationOnMock.getArguments()[1];
+ var listener = (ActionListener) invocationOnMock.getArguments()[2];
+ var mlTrainedModelResults = new ArrayList();
+ for (int i = 0; i < request.numberOfDocuments(); i++) {
+ mlTrainedModelResults.add(MlTextEmbeddingResultsTests.createRandomResults());
+ }
+ var response = new InferModelAction.Response(mlTrainedModelResults, "foo", true);
+ listener.onResponse(response);
+ return null;
+ }).when(client).execute(same(InferModelAction.INSTANCE), any(InferModelAction.Request.class), any(ActionListener.class));
+
+ var service = createService(client);
+
+ var gotResults = new AtomicBoolean();
+ var resultsListener = ActionListener.>wrap(chunkedResponse -> {
+ assertThat(chunkedResponse, hasSize(1));
+ assertThat(chunkedResponse.get(0), instanceOf(InferenceChunkedTextEmbeddingFloatResults.class));
+ var sparseResults = (InferenceChunkedTextEmbeddingFloatResults) chunkedResponse.get(0);
+ assertThat(sparseResults.chunks(), hasSize(numChunks));
+
+ gotResults.set(true);
+ }, ESTestCase::fail);
+
+ // Create model using the word boundary chunker.
+ var model = new MultilingualE5SmallModel(
+ "foo",
+ TaskType.TEXT_EMBEDDING,
+ "e5",
+ new MultilingualE5SmallInternalServiceSettings(1, 1, "cross-platform", null),
+ new WordBoundaryChunkingSettings(wordsPerChunk, 0)
+ );
+
+ var latch = new CountDownLatch(1);
+ var latchedListener = new LatchedActionListener<>(resultsListener, latch);
+
+ // For the given input we know how many requests will be made
+ service.chunkedInfer(
+ model,
+ null,
+ List.of(input),
+ Map.of(),
+ InputType.SEARCH,
+ new ChunkingOptions(null, null),
+ InferenceAction.Request.DEFAULT_TIMEOUT,
+ latchedListener
+ );
+
+ latch.await();
+ assertTrue("Listener not called with results", gotResults.get());
+ }
+
public void testParsePersistedConfig_Rerank() {
// with task settings
{
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java
index bab7bb52f928f..5a06308a3c8cc 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingCapacity.java
@@ -17,7 +17,11 @@ public static Builder builder(ByteSizeValue nodeSize, ByteSizeValue tierSize) {
}
public static Builder from(AutoscalingCapacity autoscalingCapacity) {
- return builder(autoscalingCapacity.node().memory(), autoscalingCapacity.total().memory());
+ if (autoscalingCapacity == null) {
+ return builder(null, null);
+ } else {
+ return builder(autoscalingCapacity.node().memory(), autoscalingCapacity.total().memory());
+ }
}
@Override
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java
index dfe0e557f749d..0ff6aece95ab1 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlMemoryAutoscalingDecider.java
@@ -809,7 +809,7 @@ static MlMemoryAutoscalingCapacity ensureScaleDown(
MlMemoryAutoscalingCapacity scaleDownResult,
MlMemoryAutoscalingCapacity currentCapacity
) {
- if (scaleDownResult == null || currentCapacity == null) {
+ if (scaleDownResult == null || currentCapacity == null || currentCapacity.isUndetermined()) {
return null;
}
MlMemoryAutoscalingCapacity newCapacity = MlMemoryAutoscalingCapacity.builder(