semantic text bulk inference integration test #17
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,198 @@ | ||
| /* | ||
| * 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.action.filter; | ||
|
|
||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.action.admin.indices.refresh.RefreshRequest; | ||
| import org.elasticsearch.action.bulk.BulkItemResponse; | ||
| import org.elasticsearch.action.bulk.BulkRequestBuilder; | ||
| import org.elasticsearch.action.bulk.BulkResponse; | ||
| import org.elasticsearch.action.index.IndexRequestBuilder; | ||
| import org.elasticsearch.action.search.SearchRequest; | ||
| import org.elasticsearch.action.search.SearchResponse; | ||
| import org.elasticsearch.action.update.UpdateRequestBuilder; | ||
| import org.elasticsearch.cluster.metadata.IndexMetadata; | ||
| import org.elasticsearch.common.settings.Settings; | ||
| import org.elasticsearch.inference.InferenceServiceExtension; | ||
| import org.elasticsearch.inference.Model; | ||
| import org.elasticsearch.inference.SimilarityMeasure; | ||
| import org.elasticsearch.plugins.Plugin; | ||
| import org.elasticsearch.search.builder.SearchSourceBuilder; | ||
| import org.elasticsearch.test.ESIntegTestCase; | ||
| import org.elasticsearch.xpack.inference.InferencePlugin; | ||
| import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension; | ||
| import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension; | ||
| import org.elasticsearch.xpack.inference.registry.ModelRegistry; | ||
| import org.junit.Before; | ||
|
|
||
| import java.util.Arrays; | ||
| import java.util.Collection; | ||
| import java.util.Collections; | ||
| import java.util.HashMap; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.concurrent.CountDownLatch; | ||
| import java.util.concurrent.atomic.AtomicReference; | ||
| import java.util.function.Consumer; | ||
|
|
||
| import static org.hamcrest.CoreMatchers.is; | ||
| import static org.hamcrest.Matchers.equalTo; | ||
| import static org.hamcrest.Matchers.nullValue; | ||
|
|
||
| public class ShardBulkInferenceActionFilterIT extends ESIntegTestCase { | ||
|
|
||
| public static final String INDEX_NAME = "test-index"; | ||
|
|
||
| @Before | ||
| public void setup() throws Exception { | ||
| storeSparseModel(); | ||
| storeDenseModel(); | ||
| } | ||
|
|
||
| @Override | ||
| protected Collection<Class<? extends Plugin>> nodePlugins() { | ||
| return Arrays.asList(TestInferencePlugin.class); | ||
| } | ||
|
|
||
| public void testBulkOperations() throws Exception { | ||
| Map<String, Integer> shardsSettings = Collections.singletonMap(IndexMetadata.SETTING_NUMBER_OF_SHARDS, randomIntBetween(1, 10)); | ||
| indicesAdmin().prepareCreate(INDEX_NAME).setMapping(""" | ||
| { | ||
| "properties": { | ||
| "sparse_field": { | ||
| "type": "semantic_text", | ||
| "inference_id": "test_service" | ||
| }, | ||
| "dense_field": { | ||
| "type": "semantic_text", | ||
| "inference_id": "text_embedding_test_service" | ||
| } | ||
| } | ||
| } | ||
| """).setSettings(shardsSettings).get(); | ||
|
|
||
| int totalBulkReqs = randomIntBetween(2, 100); | ||
| long totalDocs = 0; | ||
| for (int bulkReqs = 0; bulkReqs < totalBulkReqs; bulkReqs++) { | ||
| BulkRequestBuilder bulkReqBuilder = client().prepareBulk(); | ||
| int totalBulkSize = randomIntBetween(1, 100); | ||
| for (int bulkSize = 0; bulkSize < totalBulkSize; bulkSize++) { | ||
| String id = Long.toString(totalDocs); | ||
| boolean isIndexRequest = randomBoolean(); | ||
| Map<String, Object> source = new HashMap<>(); | ||
| source.put("sparse_field", isIndexRequest && rarely() ? null : randomAlphaOfLengthBetween(0, 1000)); | ||
| source.put("dense_field", isIndexRequest && rarely() ? null : randomAlphaOfLengthBetween(0, 1000)); | ||
| if (isIndexRequest) { | ||
| bulkReqBuilder.add(new IndexRequestBuilder(client()).setIndex(INDEX_NAME).setId(id).setSource(source)); | ||
| totalDocs++; | ||
| } else { | ||
| boolean isUpsert = randomBoolean(); | ||
| UpdateRequestBuilder request = new UpdateRequestBuilder(client()).setIndex(INDEX_NAME).setDoc(source); | ||
| if (isUpsert || totalDocs == 0) { | ||
| request.setDocAsUpsert(true); | ||
| totalDocs++; | ||
| } else { | ||
| // Update already existing document | ||
| id = Long.toString(randomLongBetween(0, totalDocs - 1)); | ||
| } | ||
| request.setId(id); | ||
| bulkReqBuilder.add(request); | ||
| } | ||
| } | ||
| BulkResponse bulkResponse = bulkReqBuilder.get(); | ||
| if (bulkResponse.hasFailures()) { | ||
| // Get more details in case something fails | ||
| for (BulkItemResponse bulkItemResponse : bulkResponse.getItems()) { | ||
| if (bulkItemResponse.isFailed()) { | ||
| fail( | ||
| bulkItemResponse.getFailure().getCause(), | ||
| "Failed to index document %s: %s", | ||
| bulkItemResponse.getId(), | ||
| bulkItemResponse.getFailureMessage() | ||
| ); | ||
| } | ||
| } | ||
| } | ||
| assertFalse(bulkResponse.hasFailures()); | ||
| } | ||
|
|
||
| client().admin().indices().refresh(new RefreshRequest(INDEX_NAME)).get(); | ||
|
|
||
| SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().size(0).trackTotalHits(true); | ||
| SearchResponse searchResponse = client().search(new SearchRequest(INDEX_NAME).source(sourceBuilder)).get(); | ||
| assertThat(searchResponse.getHits().getTotalHits().value, equalTo(totalDocs)); | ||
| searchResponse.decRef(); | ||
| } | ||
|
|
||
| private void storeSparseModel() throws Exception { | ||
| Model model = new TestSparseInferenceServiceExtension.TestSparseModel( | ||
| TestSparseInferenceServiceExtension.TestInferenceService.NAME, | ||
| new TestSparseInferenceServiceExtension.TestServiceSettings( | ||
| TestSparseInferenceServiceExtension.TestInferenceService.NAME, | ||
| null, | ||
| false | ||
| ) | ||
| ); | ||
| storeModel(model); | ||
| } | ||
|
|
||
| private void storeDenseModel() throws Exception { | ||
| Model model = new TestDenseInferenceServiceExtension.TestDenseModel( | ||
| TestDenseInferenceServiceExtension.TestInferenceService.NAME, | ||
| new TestDenseInferenceServiceExtension.TestServiceSettings( | ||
| TestDenseInferenceServiceExtension.TestInferenceService.NAME, | ||
| randomIntBetween(1, 100), | ||
| // dot product means that we need normalized vectors; it's not worth doing that in this test | ||
| randomValueOtherThan(SimilarityMeasure.DOT_PRODUCT, () -> randomFrom(SimilarityMeasure.values())) | ||
| ) | ||
| ); | ||
|
|
||
| storeModel(model); | ||
| } | ||
|
|
||
| private void storeModel(Model model) throws Exception { | ||
| ModelRegistry modelRegistry = new ModelRegistry(client()); | ||
|
|
||
| AtomicReference<Boolean> storeModelHolder = new AtomicReference<>(); | ||
| AtomicReference<Exception> exceptionHolder = new AtomicReference<>(); | ||
|
|
||
| blockingCall(listener -> modelRegistry.storeModel(model, listener), storeModelHolder, exceptionHolder); | ||
|
|
||
| assertThat(storeModelHolder.get(), is(true)); | ||
| assertThat(exceptionHolder.get(), is(nullValue())); | ||
| } | ||
|
|
||
| private <T> void blockingCall(Consumer<ActionListener<T>> function, AtomicReference<T> response, AtomicReference<Exception> error) | ||
| throws InterruptedException { | ||
| CountDownLatch latch = new CountDownLatch(1); | ||
| ActionListener<T> listener = ActionListener.wrap(r -> { | ||
| response.set(r); | ||
| latch.countDown(); | ||
| }, e -> { | ||
| error.set(e); | ||
| latch.countDown(); | ||
| }); | ||
|
|
||
| function.accept(listener); | ||
| latch.await(); | ||
| } | ||
|
|
||
| public static class TestInferencePlugin extends InferencePlugin { | ||
| public TestInferencePlugin(Settings settings) { | ||
| super(settings); | ||
| } | ||
|
|
||
| @Override | ||
| public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() { | ||
| return List.of( | ||
| TestSparseInferenceServiceExtension.TestInferenceService::new, | ||
| TestDenseInferenceServiceExtension.TestInferenceService::new | ||
| ); | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Moved to specific methods on the different dense and sparse vector mock services