Skip to content

Commit 33b534a

Browse files
brianf-awsbrianf-aws
authored
[Backport 2.18] Adds preset contentRegistry for IngestProcessors (opensearch-project#3281) (opensearch-project#3318)
* Adds preset contentRegistry for IngestProcessors (opensearch-project#3281) * add preset xContentRegistry to ingestProcessors for custom parametized local models Curently local models that use the parameters map within the payload to create a request can not create objects to be used for local model prediction. This requires a opensearch core change because it needs the contentRegistry,however given there is not much dependency on the registry (currently) we can give it the preset registry given in the MachineLearningPlugin class vai the getNamedXContent() class Signed-off-by: Brian Flores <iflorbri@amazon.com> * Adds UT for proving models depend on xContentRegistry for prediction Signed-off-by: Brian Flores <iflorbri@amazon.com> * apply spotless Signed-off-by: Brian Flores <iflorbri@amazon.com> * Adds IT for Asymmetric Embedding scenario with MLInferenceIngestProcessor We needed to make sure that a IT existed so that the preset content registry on the processor could work with parametized local models. By providing an IT that uses the asymetric embedding model its proven that the content registry is needed to create the embeddings. In this specific test case I used a ingest pipeline to convert passage embeddings, by simulating the pipeline to save test time. Signed-off-by: Brian Flores <iflorbri@amazon.com> --------- Signed-off-by: Brian Flores <iflorbri@amazon.com> (cherry picked from commit df1b1ef) * Adds check for Java 11 The test passes on JAVA 21 but for some reason Java 11 has trouble passing in a message so it returns null when getMessage occurs so when Mockito tries to invoke getMessage it will get get a second NPE. and thus is why the exception is caught. This code change is to get around the gap that Java 11 has. Signed-off-by: Brian Flores <iflorbri@amazon.com> * spotless Signed-off-by: Brian Flores <iflorbri@amazon.com> --------- Signed-off-by: Brian Flores <iflorbri@amazon.com>
1 parent 6e21f20 commit 33b534a

File tree

4 files changed

+185
-3
lines changed

4 files changed

+185
-3
lines changed

‎plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -1063,10 +1063,12 @@ public void loadExtensions(ExtensionLoader loader) {
10631063
@Override
10641064
public Map<String, org.opensearch.ingest.Processor.Factory> getProcessors(org.opensearch.ingest.Processor.Parameters parameters) {
10651065
Map<String, org.opensearch.ingest.Processor.Factory> processors = new HashMap<>();
1066+
NamedXContentRegistry contentRegistry = new NamedXContentRegistry(getNamedXContent());
1067+
10661068
processors
10671069
.put(
10681070
MLInferenceIngestProcessor.TYPE,
1069-
new MLInferenceIngestProcessor.Factory(parameters.scriptService, parameters.client, xContentRegistry)
1071+
new MLInferenceIngestProcessor.Factory(parameters.scriptService, parameters.client, contentRegistry)
10701072
);
10711073
return Collections.unmodifiableMap(processors);
10721074
}

‎plugin/src/test/java/org/opensearch/ml/processor/MLInferenceIngestProcessorTests.java

+66-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,14 @@
66
package org.opensearch.ml.processor;
77

88
import static org.mockito.ArgumentMatchers.any;
9-
import static org.mockito.Mockito.*;
9+
import static org.mockito.Mockito.argThat;
10+
import static org.mockito.Mockito.doAnswer;
11+
import static org.mockito.Mockito.eq;
12+
import static org.mockito.Mockito.isNull;
13+
import static org.mockito.Mockito.mock;
14+
import static org.mockito.Mockito.times;
15+
import static org.mockito.Mockito.verify;
16+
import static org.mockito.Mockito.when;
1017
import static org.opensearch.ml.common.utils.StringUtils.toJson;
1118
import static org.opensearch.ml.processor.MLInferenceIngestProcessor.DEFAULT_OUTPUT_FIELD_NAME;
1219

@@ -31,6 +38,7 @@
3138
import org.opensearch.core.action.ActionListener;
3239
import org.opensearch.core.xcontent.NamedXContentRegistry;
3340
import org.opensearch.ingest.IngestDocument;
41+
import org.opensearch.ml.common.FunctionName;
3442
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
3543
import org.opensearch.ml.common.input.MLInput;
3644
import org.opensearch.ml.common.output.model.MLResultDataType;
@@ -138,6 +146,63 @@ public void testExecute_Exception() throws Exception {
138146

139147
}
140148

149+
/**
150+
* Models that use the parameters field need to have a valid NamedXContentRegistry object to create valid MLInputs. For example
151+
* <pre>
152+
* PUT /_plugins/_ml/_predict/text_embedding/model_id
153+
* {
154+
* "parameters": {
155+
* "content_type" : "query"
156+
* },
157+
* "text_docs" : ["what day is it today?"],
158+
* "target_response" : ["sentence_embedding"]
159+
* }
160+
* </pre>
161+
* These types of models like Local Asymmetric embedding models use the parameters field.
162+
* And as such we need to test that having the contentRegistry throws an exception as it can not
163+
* properly create a valid MLInput to perform prediction
164+
*
165+
* @implNote If you check the stack trace of the test you will see it tells you that it's a direct consequence of xContentRegistry being null
166+
*/
167+
public void testExecute_xContentRegistryNullWithLocalModel_throwsException() throws Exception {
168+
// Set the registry to null and reset after exiting the test
169+
xContentRegistry = null;
170+
171+
String localModelInput =
172+
"{ \"text_docs\": [\"What day is it today?\"],\"target_response\": [\"sentence_embedding\"], \"parameters\": { \"contentType\" : \"query\"} }";
173+
174+
MLInferenceIngestProcessor processor = createMLInferenceProcessor(
175+
"local_model_id",
176+
null,
177+
null,
178+
null,
179+
false,
180+
FunctionName.TEXT_EMBEDDING.toString(),
181+
false,
182+
false,
183+
false,
184+
localModelInput
185+
);
186+
try {
187+
String npeMessage =
188+
"Cannot invoke \"org.opensearch.ml.common.input.MLInput.setAlgorithm(org.opensearch.ml.common.FunctionName)\" because \"mlInput\" is null";
189+
190+
processor.execute(ingestDocument, handler);
191+
verify(handler)
192+
.accept(
193+
isNull(),
194+
argThat(exception -> exception instanceof NullPointerException && exception.getMessage().equals(npeMessage))
195+
);
196+
} catch (Exception e) {
197+
// Java 11 doesn't pass the message correctly resulting in an anonymous NPE thus this getMessage() results in null
198+
if (e.getMessage() != null) {
199+
assertEquals("this catch block should not get executed.", e.getMessage());
200+
}
201+
}
202+
// reset to mocked object
203+
xContentRegistry = mock(NamedXContentRegistry.class);
204+
}
205+
141206
/**
142207
* test nested object document with array of Map<String,String>
143208
*/

‎plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceIngestProcessorIT.java

+115-1
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.ml.rest;
77

88
import static org.opensearch.ml.common.MLTask.MODEL_ID_FIELD;
9+
import static org.opensearch.ml.utils.TestData.SENTENCE_TRANSFORMER_MODEL_HASH_VALUE;
910
import static org.opensearch.ml.utils.TestData.SENTENCE_TRANSFORMER_MODEL_URL;
1011
import static org.opensearch.ml.utils.TestHelper.makeRequest;
1112

@@ -25,6 +26,7 @@
2526
import org.opensearch.ml.common.transport.register.MLRegisterModelInput;
2627
import org.opensearch.ml.utils.TestHelper;
2728

29+
import com.jayway.jsonpath.DocumentContext;
2830
import com.jayway.jsonpath.JsonPath;
2931

3032
public class RestMLInferenceIngestProcessorIT extends MLCommonsRestTestCase {
@@ -431,6 +433,110 @@ public void testMLInferenceProcessorLocalModelObjectField() throws Exception {
431433
Assert.assertEquals(0.49191704, (Double) embedding2.get(0), 0.005);
432434
}
433435

436+
public void testMLInferenceIngestProcessor_simulatesIngestPipelineSuccessfully_withAsymmetricEmbeddingModelUsingPassageContentType()
437+
throws Exception {
438+
String taskId = registerModel(TestHelper.toJsonString(registerAsymmetricEmbeddingModelInput()));
439+
waitForTask(taskId, MLTaskState.COMPLETED);
440+
getTask(client(), taskId, response -> {
441+
assertNotNull(response.get(MODEL_ID_FIELD));
442+
this.localModelId = (String) response.get(MODEL_ID_FIELD);
443+
try {
444+
String deployTaskID = deployModel(this.localModelId);
445+
waitForTask(deployTaskID, MLTaskState.COMPLETED);
446+
447+
getModel(client(), this.localModelId, model -> { assertEquals("DEPLOYED", model.get("model_state")); });
448+
} catch (IOException | InterruptedException e) {
449+
throw new RuntimeException(e);
450+
}
451+
});
452+
453+
String asymmetricPipelineName = "asymmetric_embedding_pipeline";
454+
String createPipelineRequestBody = "{\n"
455+
+ " \"description\": \"ingest PASSAGE text and generate a embedding using an asymmetric model\",\n"
456+
+ " \"processors\": [\n"
457+
+ " {\n"
458+
+ " \"ml_inference\": {\n"
459+
+ "\n"
460+
+ " \"model_input\": \"{\\\"text_docs\\\":[\\\"${input_map.text_docs}\\\"],\\\"target_response\\\":[\\\"sentence_embedding\\\"],\\\"parameters\\\":{\\\"content_type\\\":\\\"passage\\\"}}\",\n"
461+
+ " \"function_name\": \"text_embedding\",\n"
462+
+ " \"model_id\": \""
463+
+ this.localModelId
464+
+ "\",\n"
465+
+ " \"input_map\": [\n"
466+
+ " {\n"
467+
+ " \"text_docs\": \"description\"\n"
468+
+ " }\n"
469+
+ " ],\n"
470+
+ " \"output_map\": [\n"
471+
+ " {\n"
472+
+ "\n"
473+
+ " "
474+
+ " \"fact_embedding\": \"$.inference_results[0].output[0].data\"\n"
475+
+ " }\n"
476+
+ " ]\n"
477+
+ " }\n"
478+
+ " }\n"
479+
+ " ]\n"
480+
+ "}";
481+
482+
createPipelineProcessor(createPipelineRequestBody, asymmetricPipelineName);
483+
String sampleDocuments = "{\n"
484+
+ " \"docs\": [\n"
485+
+ " {\n"
486+
+ " \"_index\": \"my-index\",\n"
487+
+ " \"_id\": \"1\",\n"
488+
+ " \"_source\": {\n"
489+
+ " \"title\": \"Central Park\",\n"
490+
+ " \"description\": \"A large public park in the heart of New York City, offering a wide range of recreational activities.\"\n"
491+
+ " }\n"
492+
+ " },\n"
493+
+ " {\n"
494+
+ " \"_index\": \"my-index\",\n"
495+
+ " \"_id\": \"2\",\n"
496+
+ " \"_source\": {\n"
497+
+ " \"title\": \"Empire State Building\",\n"
498+
+ " \"description\": \"An iconic skyscraper in New York City offering breathtaking views from its observation deck.\"\n"
499+
+ " }\n"
500+
+ " }\n"
501+
+ " ]\n"
502+
+ "}\n";
503+
504+
Map simulateResponseDocuments = simulateIngestPipeline(asymmetricPipelineName, sampleDocuments);
505+
506+
DocumentContext documents = JsonPath.parse(simulateResponseDocuments);
507+
508+
List centralParkFactEmbedding = documents.read("docs.[0].*._source.fact_embedding.*");
509+
assertEquals(768, centralParkFactEmbedding.size());
510+
Assert.assertEquals(0.5137818, (Double) centralParkFactEmbedding.get(0), 0.005);
511+
512+
List empireStateBuildingFactEmbedding = documents.read("docs.[1].*._source.fact_embedding.*");
513+
assertEquals(768, empireStateBuildingFactEmbedding.size());
514+
Assert.assertEquals(0.4493039, (Double) empireStateBuildingFactEmbedding.get(0), 0.005);
515+
}
516+
517+
private MLRegisterModelInput registerAsymmetricEmbeddingModelInput() {
518+
MLModelConfig modelConfig = TextEmbeddingModelConfig
519+
.builder()
520+
.modelType("bert")
521+
.frameworkType(TextEmbeddingModelConfig.FrameworkType.SENTENCE_TRANSFORMERS)
522+
.embeddingDimension(768)
523+
.queryPrefix("query >>")
524+
.passagePrefix("passage >> ")
525+
.build();
526+
527+
return MLRegisterModelInput
528+
.builder()
529+
.modelName("test_model_name")
530+
.version("1.0.0")
531+
.functionName(FunctionName.TEXT_EMBEDDING)
532+
.modelFormat(MLModelFormat.TORCH_SCRIPT)
533+
.modelConfig(modelConfig)
534+
.url(SENTENCE_TRANSFORMER_MODEL_URL)
535+
.deployModel(false)
536+
.hashValue(SENTENCE_TRANSFORMER_MODEL_HASH_VALUE)
537+
.build();
538+
}
539+
434540
// TODO: add tests for other local model types such as sparse/cross encoders
435541
public void testMLInferenceProcessorLocalModelNestedField() throws Exception {
436542

@@ -550,6 +656,14 @@ protected void createPipelineProcessor(String requestBody, final String pipeline
550656

551657
}
552658

659+
protected Map simulateIngestPipeline(String pipelineName, String sampleDocuments) throws IOException {
660+
Response ingestionResponse = TestHelper
661+
.makeRequest(client(), "POST", "/_ingest/pipeline/" + pipelineName + "/_simulate", null, sampleDocuments, null);
662+
assertEquals(200, ingestionResponse.getStatusLine().getStatusCode());
663+
664+
return parseResponseToMap(ingestionResponse);
665+
}
666+
553667
protected void createIndex(String indexName, String requestBody) throws Exception {
554668
Response response = makeRequest(client(), "PUT", indexName, null, requestBody, null);
555669
assertEquals(200, response.getStatusLine().getStatusCode());
@@ -585,7 +699,7 @@ protected MLRegisterModelInput registerModelInput() throws IOException, Interrup
585699
.modelConfig(modelConfig)
586700
.url(SENTENCE_TRANSFORMER_MODEL_URL)
587701
.deployModel(false)
588-
.hashValue("e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021")
702+
.hashValue(SENTENCE_TRANSFORMER_MODEL_HASH_VALUE)
589703
.build();
590704
}
591705

‎plugin/src/test/java/org/opensearch/ml/utils/TestData.java

+1
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ public class TestData {
3030
"https://github.com/opensearch-project/ml-commons/blob/2.x/ml-algorithms/src/test/resources/org/opensearch/ml/engine/algorithms/text_embedding/all-MiniLM-L6-v2_torchscript_huggingface.zip?raw=true";
3131
public static final String SENTENCE_TRANSFORMER_MODEL_URL =
3232
"https://github.com/opensearch-project/ml-commons/blob/2.x/ml-algorithms/src/test/resources/org/opensearch/ml/engine/algorithms/text_embedding/traced_small_model.zip?raw=true";
33+
public static final String SENTENCE_TRANSFORMER_MODEL_HASH_VALUE = "e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021";
3334
public static final String TIME_FIELD = "timestamp";
3435
public static final String HUGGINGFACE_TRANSFORMER_MODEL_HASH_VALUE =
3536
"e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021";

0 commit comments

Comments
 (0)