Skip to content

Commit

Permalink
Addressing martin and hemmin comments
Browse files Browse the repository at this point in the history
Signed-off-by: Varun Jain <[email protected]>
  • Loading branch information
vibrantvarun committed Jan 25, 2024
1 parent 7bccb42 commit 98050a1
Show file tree
Hide file tree
Showing 14 changed files with 274 additions and 333 deletions.
2 changes: 2 additions & 0 deletions qa/restart-upgrade/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ task testAgainstOldCluster(type: StandaloneRestIntegTestTask) {
systemProperty 'tests.skip_delete_model_index', 'true'
systemProperty 'tests.plugin_bwc_version', neural_search_bwc_version

//Excluding MultiModalSearchIT, HybridSearchIT, NeuralSparseSearchIT tests from neural search version 2.9 and 2.10 because these features were released in 2.11 version.
if (neural_search_bwc_version.startsWith("2.9") || neural_search_bwc_version.startsWith("2.10")){
filter {
excludeTestsMatching "org.opensearch.neuralsearch.bwc.MultiModalSearchIT.*"
Expand All @@ -61,6 +62,7 @@ task testAgainstNewCluster(type: StandaloneRestIntegTestTask) {
systemProperty 'tests.is_old_cluster', 'false'
systemProperty 'tests.plugin_bwc_version', neural_search_bwc_version

//Excluding MultiModalSearchIT, HybridSearchIT, NeuralSparseSearchIT tests from neural search version 2.9 and 2.10 because these features were released in 2.11 version.
if (neural_search_bwc_version.startsWith("2.9") || neural_search_bwc_version.startsWith("2.10")){
filter {
excludeTestsMatching "org.opensearch.neuralsearch.bwc.MultiModalSearchIT.*"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,18 @@
*/
package org.opensearch.neuralsearch.bwc;

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Locale;
import java.util.Optional;
import org.junit.Before;
import org.opensearch.common.settings.Settings;
import org.opensearch.neuralsearch.BaseNeuralSearchIT;
import static org.opensearch.neuralsearch.TestUtils.NEURAL_SEARCH_BWC_PREFIX;
import static org.opensearch.neuralsearch.TestUtils.CLIENT_TIMEOUT_VALUE;
import static org.opensearch.neuralsearch.TestUtils.RESTART_UPGRADE_OLD_CLUSTER;
import static org.opensearch.neuralsearch.TestUtils.BWC_VERSION;
import static org.opensearch.neuralsearch.TestUtils.NEURAL_SEARCH_BWC_PREFIX;
import static org.opensearch.neuralsearch.TestUtils.generateModelId;
import org.opensearch.test.rest.OpenSearchRestTestCase;

public abstract class AbstractRestartUpgradeRestTestCase extends BaseNeuralSearchIT {
Expand Down Expand Up @@ -57,4 +60,43 @@ protected static final boolean isRunningAgainstOldCluster() {
protected final Optional<String> getBWCVersion() {
return Optional.ofNullable(System.getProperty(BWC_VERSION, null));
}

protected String uploadTextEmbeddingModel() throws Exception {
String requestBody = Files.readString(Path.of(classLoader.getResource("processor/UploadModelRequestBody.json").toURI()));
return registerModelGroupAndGetModelId(requestBody);
}

protected String registerModelGroupAndGetModelId(final String requestBody) throws Exception {
String modelGroupRegisterRequestBody = Files.readString(
Path.of(classLoader.getResource("processor/CreateModelGroupRequestBody.json").toURI())
);
String modelGroupId = registerModelGroup(String.format(LOCALE, modelGroupRegisterRequestBody, generateModelId()));
return uploadModel(String.format(LOCALE, requestBody, modelGroupId));
}

protected void createPipelineProcessor(final String modelId, final String pipelineName) throws Exception {
String requestBody = Files.readString(Path.of(classLoader.getResource("processor/PipelineConfiguration.json").toURI()));
createPipelineProcessor(requestBody, pipelineName, modelId);
}

protected String uploadSparseEncodingModel() throws Exception {
String requestBody = Files.readString(
Path.of(classLoader.getResource("processor/UploadSparseEncodingModelRequestBody.json").toURI())
);
return registerModelGroupAndGetModelId(requestBody);
}

protected void createPipelineForTextImageProcessor(final String modelId, final String pipelineName) throws Exception {
String requestBody = Files.readString(
Path.of(classLoader.getResource("processor/PipelineForTextImageProcessorConfiguration.json").toURI())
);
createPipelineProcessor(requestBody, pipelineName, modelId);
}

protected void createPipelineForSparseEncodingProcessor(final String modelId, final String pipelineName) throws Exception {
String requestBody = Files.readString(
Path.of(classLoader.getResource("processor/PipelineForSparseEncodingProcessorConfiguration.json").toURI())
);
createPipelineProcessor(requestBody, pipelineName, modelId);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,14 @@
*/
package org.opensearch.neuralsearch.bwc;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import org.opensearch.index.query.MatchQueryBuilder;
import static org.opensearch.neuralsearch.TestUtils.getModelId;
import static org.opensearch.neuralsearch.TestUtils.generateModelId;
import static org.opensearch.neuralsearch.TestUtils.NODES_BWC_CLUSTER;
import static org.opensearch.neuralsearch.TestUtils.PARAM_NAME_WEIGHTS;
import static org.opensearch.neuralsearch.TestUtils.TEXT_EMBEDDING_PROCESSOR;
Expand Down Expand Up @@ -48,28 +48,18 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr
Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())),
PIPELINE_NAME
);
addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT_1, null, null);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_2, null, null);
addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_3, null, null);
addDocument(getIndexNameForTest(), "3", TEST_FIELD, TEXT_4, null, null);
addDocument(getIndexNameForTest(), "4", TEST_FIELD, TEXT_5, null, null);
createSearchPipeline(
SEARCH_PIPELINE_NAME,
DEFAULT_NORMALIZATION_METHOD,
DEFAULT_COMBINATION_METHOD,
Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f }))
);
addDocuments(getIndexNameForTest(), true);
createSearchPipeline(SEARCH_PIPELINE_NAME);
} else {
Map<String, Object> pipeline = getIngestionPipeline(PIPELINE_NAME);
assertNotNull(pipeline);
String modelId = getModelId(pipeline, TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "5", TEST_FIELD, TEXT_6, null, null);
validateTestIndex(modelId, getIndexNameForTest(), SEARCH_PIPELINE_NAME);
deleteSearchPipeline(SEARCH_PIPELINE_NAME);
deletePipeline(PIPELINE_NAME);
deleteModel(modelId);
deleteIndex(getIndexNameForTest());
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocuments(getIndexNameForTest(), false);
validateTestIndex(modelId, getIndexNameForTest(), SEARCH_PIPELINE_NAME);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME);
}
}
}

Expand All @@ -87,50 +77,43 @@ public void testNormalizationProcessor_whenIndexWithSingleShard_E2EFlow() throws
Files.readString(Path.of(classLoader.getResource("processor/IndexMappingSingleShard.json").toURI())),
PIPELINE1_NAME
);
addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT_1, null, null);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_2, null, null);
addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_3, null, null);
addDocument(getIndexNameForTest(), "3", TEST_FIELD, TEXT_4, null, null);
addDocument(getIndexNameForTest(), "4", TEST_FIELD, TEXT_5, null, null);
createSearchPipeline(
SEARCH_PIPELINE1_NAME,
DEFAULT_NORMALIZATION_METHOD,
DEFAULT_COMBINATION_METHOD,
Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f }))
);
addDocuments(getIndexNameForTest(), true);
createSearchPipeline(SEARCH_PIPELINE1_NAME);
} else {
Map<String, Object> pipeline = getIngestionPipeline(PIPELINE1_NAME);
assertNotNull(pipeline);
String modelId = getModelId(pipeline, TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "5", TEST_FIELD, TEXT_6, null, null);
validateTestIndex(modelId, getIndexNameForTest(), SEARCH_PIPELINE1_NAME);
deleteSearchPipeline(SEARCH_PIPELINE1_NAME);
deletePipeline(PIPELINE1_NAME);
deleteModel(modelId);
deleteIndex(getIndexNameForTest());
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE1_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocuments(getIndexNameForTest(), false);
validateTestIndex(modelId, getIndexNameForTest(), SEARCH_PIPELINE1_NAME);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE1_NAME, modelId, SEARCH_PIPELINE1_NAME);
}
}
}

private String uploadTextEmbeddingModel() throws Exception {
String requestBody = Files.readString(Path.of(classLoader.getResource("processor/UploadModelRequestBody.json").toURI()));
return registerModelGroupAndGetModelId(requestBody);
private void addDocuments(final String indexName, boolean isRunningAgainstOldCluster) throws IOException {
if (isRunningAgainstOldCluster) {
addDocument(indexName, "0", TEST_FIELD, TEXT_1, null, null);
addDocument(indexName, "1", TEST_FIELD, TEXT_2, null, null);
addDocument(indexName, "2", TEST_FIELD, TEXT_3, null, null);
addDocument(indexName, "3", TEST_FIELD, TEXT_4, null, null);
addDocument(indexName, "4", TEST_FIELD, TEXT_5, null, null);
} else {
addDocument(indexName, "5", TEST_FIELD, TEXT_6, null, null);
}
}

private String registerModelGroupAndGetModelId(String requestBody) throws Exception {
String modelGroupRegisterRequestBody = Files.readString(
Path.of(classLoader.getResource("processor/CreateModelGroupRequestBody.json").toURI())
private void createSearchPipeline(final String pipelineName) {
createSearchPipeline(
pipelineName,
DEFAULT_NORMALIZATION_METHOD,
DEFAULT_COMBINATION_METHOD,
Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f }))
);
String modelGroupId = registerModelGroup(String.format(LOCALE, modelGroupRegisterRequestBody, generateModelId()));
return uploadModel(String.format(LOCALE, requestBody, modelGroupId));
}

private void createPipelineProcessor(String modelId, String pipelineName) throws Exception {
String requestBody = Files.readString(Path.of(classLoader.getResource("processor/PipelineConfiguration.json").toURI()));
createPipelineProcessor(requestBody, pipelineName, modelId);
}

private void validateTestIndex(String modelId, String index, String searchPipeline) throws Exception {
private void validateTestIndex(final String modelId, final String index, final String searchPipeline) throws Exception {
int docCount = getDocCount(index);
assertEquals(6, docCount);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId);
Expand All @@ -144,7 +127,7 @@ private void validateTestIndex(String modelId, String index, String searchPipeli
}
}

public HybridQueryBuilder getQueryBuilder(String modelId) {
public HybridQueryBuilder getQueryBuilder(final String modelId) {
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
neuralQueryBuilder.fieldName("passage_embedding");
neuralQueryBuilder.modelId(modelId);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
import static org.opensearch.neuralsearch.TestUtils.NODES_BWC_CLUSTER;
import static org.opensearch.neuralsearch.TestUtils.TEXT_IMAGE_EMBEDDING_PROCESSOR;
import static org.opensearch.neuralsearch.TestUtils.getModelId;
import static org.opensearch.neuralsearch.TestUtils.generateModelId;
import org.opensearch.neuralsearch.query.NeuralQueryBuilder;

public class MultiModalSearchIT extends AbstractRestartUpgradeRestTestCase {
Expand All @@ -29,54 +28,34 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
waitForClusterHealthGreen(NODES_BWC_CLUSTER);

if (isRunningAgainstOldCluster()) {
String modelId = uploadTextImageEmbeddingModel();
String modelId = uploadTextEmbeddingModel();
loadModel(modelId);
createPipelineProcessor(modelId, PIPELINE_NAME);
createPipelineForTextImageProcessor(modelId, PIPELINE_NAME);
createIndexWithConfiguration(
getIndexNameForTest(),
Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())),
PIPELINE_NAME
);
addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT);
} else {
Map<String, Object> pipeline = getIngestionPipeline(PIPELINE_NAME);
assertNotNull(pipeline);
String modelId = getModelId(pipeline, TEXT_IMAGE_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
validateTestIndex(modelId);
deletePipeline(PIPELINE_NAME);
deleteModel(modelId);
deleteIndex(getIndexNameForTest());
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
validateTestIndex(modelId);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
}
}
}

private void validateTestIndex(String modelId) throws Exception {
private void validateTestIndex(final String modelId) throws Exception {
int docCount = getDocCount(getIndexNameForTest());
assertEquals(2, docCount);
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder("passage_embedding", TEXT, TEST_IMAGE_TEXT, modelId, 1, null, null);
Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
assertNotNull(response);
}

private String uploadTextImageEmbeddingModel() throws Exception {
String requestBody = Files.readString(Path.of(classLoader.getResource("processor/UploadModelRequestBody.json").toURI()));
return registerModelGroupAndGetModelId(requestBody);
}

private String registerModelGroupAndGetModelId(String requestBody) throws Exception {
String modelGroupRegisterRequestBody = Files.readString(
Path.of(classLoader.getResource("processor/CreateModelGroupRequestBody.json").toURI())
);
String modelGroupId = registerModelGroup(String.format(LOCALE, modelGroupRegisterRequestBody, generateModelId()));
return uploadModel(String.format(LOCALE, requestBody, modelGroupId));
}

protected void createPipelineProcessor(String modelId, String pipelineName) throws Exception {
String requestBody = Files.readString(
Path.of(classLoader.getResource("processor/PipelineForTextImageProcessorConfiguration.json").toURI())
);
createPipelineProcessor(requestBody, pipelineName, modelId);
}

}
Loading

0 comments on commit 98050a1

Please sign in to comment.