add ml_inference processor for offline batch inference #5507

Zhangxunmt · 2025-03-06T20:19:23Z

Description

Adding a new ml_inference processor to interact with ml-commons plugin in OpenSearch for ML related applications.

Some examples that work well:

ml-batch-job-sagemaker-pipeline:
  source:
    s3:
      codec:
        ndjson:
      compression: none
      aws:
        region: "us-east-1"
      default_bucket_owner: <your aws account>
      scan:
        scheduling:
          interval: PT6M
        buckets:
          - bucket:
              name: "<your bucket>"
              data_selection: metadata_only
              filter:
                include_prefix:
                  - <your input prefix>
          - bucket:
              name: "offlinebatch"
              data_selection: data_only
              filter:
                include_prefix: 
                  - <your output prefix>

  buffer:
    bounded_blocking:
      buffer_size: 2048 # max number of records the buffer accepts
      batch_size: 512 # max number of records the buffer drains after each read

  processor:
    - ml_inference:
        host: "<your AOS endpoint url>"
        aws_sigv4: true
        action_type: "batch_predict"
        service_name: "sagemaker|bedrock"
        model_id: "<your model id used in vector search>"
        output_path: "s3://offlinebatch/sagemaker/output"
        aws:
          region: "us-east-1"
        ml_when: /bucket == "offlinebatch"
    - copy_values:
        entries:
          - to_key: chapter
            from_key: /content/0
          - to_key: title
            from_key: /content/1
          - to_key: chapter_embedding
            from_key: /SageMakerOutput/0
          - to_key: title_embedding
            from_key: /SageMakerOutput/1
    - delete_entries:
        with_keys: [content, SageMakerOutput]
  
  route:
      - ml-ingest-route: "/chapter != null and /title != null"

  sink:
    - opensearch:
        hosts: ["<your AOS endpoint url>"]
        aws_sigv4: true
        index: "test-nlp-index"
        routes: [ml-ingest-route]
        username: "<you username>"
        password: "<your password>"

ml-batch-job-bedrock-pipeline:
  source:
    s3:
      codec:
        ndjson:
      compression: none
      aws:
        region: "us-east-1"
      default_bucket_owner: <your aws account>
      scan:
        scheduling:
          interval: PT2M
        buckets:
          - bucket:
              name: "<your bucket>"
              data_selection: metadata_only
              filter:
                include_prefix:
                  - bedrock-multisource/my_batch
                exclude_suffix:
                  - .out
          - bucket:
              name: "<your bucket>"
              data_selection: data_only
              filter:
                include_prefix:
                  - bedrock-multisource/output-multisource/
                exclude_suffix:
                  - manifest.json.out

  buffer:
    bounded_blocking:
      buffer_size: 2048 # max number of records the buffer accepts
      batch_size: 512 # max number of records the buffer drains after each read

  processor:
    - ml_inference:
        host: "<your AOS endpoint url>"
        aws_sigv4: true
        action_type: "batch_predict"
        service_name: "bedrock"
        model_id: "<your model id used in vector search>"
        output_path: "s3://offlinebatch/bedrock-multisource/output-multisource/"
        aws:
          region: "us-east-1"
        ml_when: /bucket == "offlinebatch"
    - copy_values:
        entries:
          - to_key: chapter
            from_key: /modelInput/inputText
          - to_key: chapter_embedding
            from_key: /modelOutput/embedding
    - delete_entries:
        with_keys: [modelInput, modelOutput, recordId, s3]

  route:
      - ml-ingest-route: "/chapter != null and /chapter_embedding != null"

  sink:
    - opensearch:
        hosts: ["<your AOS endpoint url>"]
        aws_sigv4: true
        index: "my-nlp-index-bedrock"
        routes: [ml-ingest-route]

Issues Resolved

#5470
#5433
#5509

Check List

New functionality includes testing.
New functionality has a documentation issue. Please link to it in this PR.
- New functionality has javadoc added
Commits are signed with a real name per the DCO

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

data-prepper-plugins/ml-processor/build.gradle

.../ml-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

settings.gradle

...ns/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectHandler.java

.../ml-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

...in/java/org/opensearch/dataprepper/plugins/ml/processor/common/SageMakerBatchJobCreator.java

dlvenable · 2025-03-12T21:38:37Z

...in/java/org/opensearch/dataprepper/plugins/ml/processor/common/SageMakerBatchJobCreator.java

+    private String findCommonPrefix(Collection<Record<Event>> records) {
+        List<String> keys = new ArrayList<>();
+        for (Record<Event> record : records) {
+            keys.add(record.getData().getJsonNode().get("key").asText());


This seems to have a very rigid expectation on having a key named key. If you need this, it should be configurable by users.

Also, prefer:

record.getData().get(eventKey, String.class)

You can get an EventKey in the constructor. See #4636 for an example of how you can use this.

added the "input_key" to allow CX to configure the key.

Thanks for using the EventKey. I'm still unsure what input_key is doing though. What does this represent exactly? Are these the keys for obtaining the bucket name and the S3 key?

We should have more concrete names for the configuration. For example: bucket_key makes sense. But, that would also tend toward key_key which is confusing!

The input_key here represents the S3 file uri. Currently the S3 scan in meta data mode sends data like this
{"bucket":"offlinebatch","length":158893,"time":1736279088.000000000,"key":"bedrock-multisource/my_batch.jsonl"}
So the input_key would be "key" in this S3 Scan mode, and that's why it's listed as "input_key: key" in the pipeline yaml file. I think it makes sense to change to "object_key"? Basically that defines the key name in the json record which has the S3 uri to be processed.

Given that the S3 Scan uses the "key" as the name of the S3Uri from the metadata scan, I use "key" as the default field to read S3Uris if this input_key is not provided in the pipeline. Otherwise, it uses the EventKey to read it from the input_key. Please let me know if you'd want a different name instead of input_key.

dlvenable · 2025-03-12T21:46:38Z

@Zhangxunmt , Thank you for this great processor!

We will also need some unit tests. I'm ok accepting this PR without them as long as we have the @Experimental annotation.

Zhangxunmt · 2025-03-12T23:13:03Z

@Zhangxunmt , Thank you for this great processor!

We will also need some unit tests. I'm ok accepting this PR without them as long as we have the @Experimental annotation.

dlvenable Thanks David for the comments. Looks like there're no major concerns. I will add the remaining UTs soon and the @experimental annotation.

...-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/util/RetryUtil.java

...nce-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

...ocessor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessorConfig.java

...nce-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

data-prepper-plugins/ml-inference-processor/build.gradle

dlvenable · 2025-03-27T21:49:42Z

.../ml-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

+            mlBatchJobCreator.createMLBatchJob(recordsToMlCommons);
+            numberOfMLProcessorSuccessCounter.increment();
+        } catch (Exception e) {
+            LOG.error(NOISY, e.getMessage(), e);


This is still unresolved.

dlvenable · 2025-03-27T21:49:52Z

.../ml-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

+        try {
+            mlBatchJobCreator.createMLBatchJob(recordsToMlCommons);
+            numberOfMLProcessorSuccessCounter.increment();
+        } catch (Exception e) {


This is still unresolved.

...nce-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java

dlvenable · 2025-03-27T21:57:28Z

...in/java/org/opensearch/dataprepper/plugins/ml/processor/common/SageMakerBatchJobCreator.java

+    private String findCommonPrefix(Collection<Record<Event>> records) {
+        List<String> keys = new ArrayList<>();
+        for (Record<Event> record : records) {
+            keys.add(record.getData().getJsonNode().get("key").asText());


Thanks for using the EventKey. I'm still unsure what input_key is doing though. What does this represent exactly? Are these the keys for obtaining the bucket name and the S3 key?

We should have more concrete names for the configuration. For example: bucket_key makes sense. But, that would also tend toward key_key which is confusing!

Signed-off-by: Xun Zhang <[email protected]>

…uation Signed-off-by: Xun Zhang <[email protected]>

…put_key is null Signed-off-by: Xun Zhang <[email protected]>

Zhangxunmt · 2025-03-31T21:51:46Z

@dlvenable Please review the latest commit for the updates to requested changes, after a rebase to the main. The Gradle Builds somehow fail due to unrelated tests.

Zhangxunmt requested review from sb2k16, chenqi0805, engechas, san81, srikanthjg, graytaylor0, dinujoh, kkondaka, KarstenSchnitter, dlvenable and oeyh as code owners March 6, 2025 20:19

Zhangxunmt force-pushed the main branch 3 times, most recently from 86e4f34 to a59fb6a Compare March 10, 2025 22:43

dlvenable requested changes Mar 12, 2025

View reviewed changes

Zhangxunmt force-pushed the main branch 5 times, most recently from 0b384c9 to 9ce8a77 Compare March 24, 2025 22:26

Zhangxunmt force-pushed the main branch 3 times, most recently from 5ae7861 to c670ca5 Compare March 25, 2025 18:28

Zhangxunmt changed the title ~~add ml processor for offline batch inference~~ add ml-inference processor for offline batch inference Mar 25, 2025

Zhangxunmt changed the title ~~add ml-inference processor for offline batch inference~~ add ml_inference processor for offline batch inference Mar 25, 2025

Zhangxunmt force-pushed the main branch 3 times, most recently from d0aa269 to 0108602 Compare March 25, 2025 20:00

graytaylor0 reviewed Mar 25, 2025

View reviewed changes

Zhangxunmt force-pushed the main branch from c14524e to 2e96545 Compare March 25, 2025 22:47

srikanthjg reviewed Mar 27, 2025

View reviewed changes

...nce-processor/src/main/java/org/opensearch/dataprepper/plugins/ml/processor/MLProcessor.java Outdated Show resolved Hide resolved

dlvenable requested changes Mar 27, 2025

View reviewed changes

Zhangxunmt added 6 commits March 31, 2025 14:07

add ml processor for offline batch inference

b8c9d31

Signed-off-by: Xun Zhang <[email protected]>

add more UTs and address comments

1b92a90

Signed-off-by: Xun Zhang <[email protected]>

rename to ml-inference

56351cd

Signed-off-by: Xun Zhang <[email protected]>

update readme

068ed30

Signed-off-by: Xun Zhang <[email protected]>

move retryUtils into common lib and add try catch for conditions eval…

78ec7ce

…uation Signed-off-by: Xun Zhang <[email protected]>

add failure tag, catch known error types, and use default key when in…

bf4a4c6

…put_key is null Signed-off-by: Xun Zhang <[email protected]>

Zhangxunmt force-pushed the main branch from 2e96545 to bf4a4c6 Compare March 31, 2025 21:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add ml_inference processor for offline batch inference #5507

add ml_inference processor for offline batch inference #5507

Zhangxunmt commented Mar 6, 2025 •

edited

Loading

dlvenable Mar 12, 2025 •

edited

Loading

Zhangxunmt Mar 24, 2025

dlvenable Mar 27, 2025

Zhangxunmt Mar 27, 2025

Zhangxunmt Mar 31, 2025

dlvenable commented Mar 12, 2025

Zhangxunmt commented Mar 12, 2025 •

edited

Loading

dlvenable Mar 27, 2025

dlvenable Mar 27, 2025

dlvenable Mar 27, 2025

Zhangxunmt commented Mar 31, 2025 •

edited

Loading

add ml_inference processor for offline batch inference #5507

Are you sure you want to change the base?

add ml_inference processor for offline batch inference #5507

Conversation

Zhangxunmt commented Mar 6, 2025 • edited Loading

Description

Issues Resolved

Check List

dlvenable Mar 12, 2025 • edited Loading

Choose a reason for hiding this comment

Zhangxunmt Mar 24, 2025

Choose a reason for hiding this comment

dlvenable Mar 27, 2025

Choose a reason for hiding this comment

Zhangxunmt Mar 27, 2025

Choose a reason for hiding this comment

Zhangxunmt Mar 31, 2025

Choose a reason for hiding this comment

dlvenable commented Mar 12, 2025

Zhangxunmt commented Mar 12, 2025 • edited Loading

dlvenable Mar 27, 2025

Choose a reason for hiding this comment

dlvenable Mar 27, 2025

Choose a reason for hiding this comment

dlvenable Mar 27, 2025

Choose a reason for hiding this comment

Zhangxunmt commented Mar 31, 2025 • edited Loading

Zhangxunmt commented Mar 6, 2025 •

edited

Loading

dlvenable Mar 12, 2025 •

edited

Loading

Zhangxunmt commented Mar 12, 2025 •

edited

Loading

Zhangxunmt commented Mar 31, 2025 •

edited

Loading