[RFC] Vertex AI support in the ml_inference processor

[RFC] Vertex AI support in the ml_inference processor

Problem Statement

The ml_inference processor (data-prepper-plugins/ml-inference-processor) drives offline
batch inference by calling the ml-commons Batch Predict API and staging batch input/output in
object storage. It supports two AI services — SageMaker and Bedrock — and is wired to
AWS end to end:

1. Service coverage is AWS-only. ServiceName is {sagemaker, bedrock} and
   MLBatchJobCreatorFactory switches on exactly those two. There is no path for
   GCP Vertex AI, even though ml-commons now ships a google_cloud connector with Vertex
   AI batch support (PR #4921 / RFC #4915).
2. Batch staging is S3-only. SageMakerBatchJobCreator builds an S3 manifest, uses an
   S3Client, and templates s3:// input/output URIs. There is no GCS-staging creator.
3. Auth to OpenSearch is SigV4-only. MlCommonRequester signs every request to the
   ml-commons REST API with AWS SigV4 (Aws4Signer, signingName = "es", region from the
   mandatory aws: block). A self-managed or GCP-hosted OpenSearch cluster has no SigV4
   identity, so the processor cannot authenticate to it at all.

The result: a pipeline running on GCP cannot use ml_inference for batch vectorization, even
once ml-commons and the GCS plugins (companion RFC #7067) are in place. This RFC closes the
processor-side gap.

Current state

• configuration/ServiceName.java — enum {SAGEMAKER, BEDROCK}.
• common/MLBatchJobCreatorFactory.java — switch over the two AWS services;
  SageMakerBatchJobCreator / BedrockBatchJobCreator are S3-native.
• util/MlCommonRequester.java — SigV4 signing hard-wired via Aws4Signer.
• MLProcessorConfig.java — aws block is @NotNull; aws_sigv4 boolean; output_path
  is an s3:// URI.

Proposed Solution

Three additive changes, none altering existing AWS behavior:

1. New vertexai service

Add VERTEXAI("vertexai") to ServiceName and a VertexAiBatchJobCreator registered in
MLBatchJobCreatorFactory. It mirrors SageMakerBatchJobCreator but stages JSONL to GCS
and templates gs:// input/output URIs in the ml-commons _batch_predict request, matching
the Vertex AI batch blueprint that PR #4921 added.

The creator writes batch input to GCS and reads the batch-job output location using the
shared GCP auth module proposed in RFC #7067 (GcpCredentialsSupplier + google-cloud-storage
client). This RFC depends on that module rather than introducing its own GCS client.

2. Pluggable auth to OpenSearch

Generalize MlCommonRequester so authentication to the ml-commons REST endpoint is not
hard-wired to SigV4. Introduce an auth abstraction with (initially) two implementations:

• aws_sigv4 — the existing behavior, unchanged and still the default when an aws:
  block is present, so current pipelines are unaffected.
• basic — HTTP Basic auth (username/password), the common case for self-managed and
  non-AWS-hosted OpenSearch (including GCP deployments). Credentials sourced via Data
  Prepper's existing secret-reference mechanism (${{...}}), e.g. the GCP Secret Manager
  provider (RFC #7010) or environment.

This decouples how the processor talks to OpenSearch from which AI service runs the batch
job — they are independent axes today only because everything happened to be AWS.

3. Config changes

• service_name: vertexai becomes valid.
• The aws: block becomes conditionally required: mandatory for sagemaker/bedrock
  and for auth_mode: aws_sigv4, not required otherwise. (Today it is unconditionally
  @NotNull.)
• A new opensearch_auth block selects the to-OpenSearch auth mode.
• output_path accepts a gs:// URI for the Vertex path.

Example (GCP-native pipeline):

processor:
  - ml_inference:
      host: "https://my-opensearch:9200"
      action_type: "batch_predict"
      service_name: "vertexai"
      model_id: "<vertex batch model id registered in ml-commons>"
      output_path: "gs://my-bucket/batch-output/"
      input_key: key
      opensearch_auth:
        mode: basic
        username: "${{gcp_secrets:os-creds:username}}"
        password: "${{gcp_secrets:os-creds:password}}"
      gcp:
        project_id: my-project
        # auth_mode: adc   # or service-account key fields
      ml_when: /bucket == "offlinebatch"

Existing AWS pipeline (unchanged):

processor:
  - ml_inference:
      host: "<opensearch url>"
      aws_sigv4: true
      action_type: "batch_predict"
      service_name: "bedrock"
      model_id: "<model id>"
      output_path: "s3://my-bucket/output/"
      aws:
        region: "us-east-1"
        sts_role_arn: "<arn>"

Scope

In scope

• vertexai ServiceName + VertexAiBatchJobCreator (GCS-staged).
• Pluggable to-OpenSearch auth: aws_sigv4 (existing) + basic (new).
• Conditional aws: block requirement; new opensearch_auth and gcp: blocks.
• Reuse of the shared GCP auth module (RFC #7067) and secret providers (RFC #7010).

Out of scope

• The GCS source/sink plugins themselves (RFC #7067).
• The ml-commons google_cloud connector and Vertex batch terminal-state handling
  (ml-commons PR #4921 — separate repo).
• Real-time (non-batch) ml_inference invocation.

Dependencies

• ml-commons PR #4921 / RFC #4915 — the google_cloud connector and Vertex AI batch
  submit/status/cancel this processor's _batch_predict calls rely on.
• RFC #7067 — GCS source/sink and the shared gcp-plugin auth module
  (GcpCredentialsSupplier) this processor reuses for GCS staging.
• RFC #7010 — GCP Secret Manager provider, a natural source for the basic auth
  credentials and GCP service-account key.

Alternatives Considered

1. Only add vertexai, leave SigV4-only auth. Rejected — a GCP-hosted OpenSearch cluster
   cannot be reached with SigV4, so the Vertex service would be unusable in exactly the
   deployments that need it. The two changes must land together to be useful.
2. A generic HTTP connector in the processor instead of per-service creators. Rejected —
   the existing design is per-service (SageMakerBatchJobCreator, BedrockBatchJobCreator);
   a Vertex creator matching that pattern is lower risk than reworking the abstraction.
3. Token/bearer auth to OpenSearch instead of basic. Basic covers the majority of
   self-managed clusters; a bearer/token mode can be added later behind the same
   opensearch_auth.mode switch if demand appears.

Open Questions

1. Auth block shape. Introduce a unified opensearch_auth: { mode: aws_sigv4 | basic }
   and deprecate the standalone aws_sigv4: true boolean, or keep the boolean for
   back-compat and only add the new block? Proposing: keep the boolean working, treat it as
   sugar for opensearch_auth.mode: aws_sigv4.
2. Sequencing vs. #7067. This RFC's VertexAiBatchJobCreator needs the shared GCP auth
   module from #7067. Land #7067's gcp-plugin module first, then this?
3. Bedrock/SageMaker with basic auth. The new basic mode is independent of service — a
   user could run Bedrock batch against a self-managed OpenSearch. Confirm we want to allow
   any (service × auth-mode) combination rather than restricting pairs.

Request for Comments

Feedback sought on: the opensearch_auth block design and back-compat approach, the
conditional aws: requirement, and the dependency sequencing across #7067 / #7010 / this RFC.

Start by reading configuration/ServiceName.java, common/MLBatchJobCreatorFactory.java, util/MlCommonRequester.java, and MLProcessorConfig.java to trace service selection, authentication, and validation. Review the dependencies on RFC #7067, RFC #7010, and ml-commons PR #4921 before deciding the configuration shape. Done means Vertex AI GCS staging, basic and existing SigV4 authentication, conditional AWS configuration, and the proposed Vertex configuration are supported without changing existing AWS behavior.