Skip to content

Add ML gRPC support - #2070

Open
nathaliellenaa wants to merge 1 commit into
opensearch-project:mainfrom
nathaliellenaa:ml_grpc
Open

Add ML gRPC support#2070
nathaliellenaa wants to merge 1 commit into
opensearch-project:mainfrom
nathaliellenaa:ml_grpc

Conversation

@nathaliellenaa

Copy link
Copy Markdown
Contributor

Description

This PR adds gRPC ML streaming support to the java-client-grpc module, exposing the ML Commons server-streaming APIs MLService/PredictModelStream and MLService/ExecuteAgentStream as GrpcTransport.predictModelStream() and GrpcTransport.executeAgentStream().

Sample usage for model prediction:

var transport = GrpcTransport.builder("localhost", 9400)
    .jsonpMapper(new JacksonJsonpMapper())
    .basicAuth("admin", "admin")
    .build();

var chunks = transport.predictModelStream(
    MlPredictStreamRequest.builder()
        .modelId(modelId)
        .parameters(p -> p
            .addMessage("system", "You are a helpful assistant.")
            .addMessage("user", "Hello world!")
            .llmInterface("openai/v1/chat/completions"))
        .build()
);

var answer = new StringBuilder();
while (chunks.hasNext()) {
    var chunk = chunks.next();
    if (chunk.content() != null) {
        answer.append(chunk.content());   // token-by-token
    }
}

Agent execution follows the same shape:

var chunks = transport.executeAgentStream(
    MlExecuteAgentStreamRequest.builder()
        .agentId(agentId)
        .parameters(p -> p.question("List indices in my cluster"))
        .build()
);

How it works

MlStreamRequestConverter builds the protobuf request from the typed builder; the blocking stub returns an Iterator<PredictResponse>; MlStreamResponseConverter converts each protobuf chunk into MlStreamChunk on demand as the caller advances the iterator. Nothing is buffered — one protobuf message in, one MlStreamChunk out.

Two behaviors are worth calling out:

  • The RPC is lazy. Calling predictModelStream() does not issue the request; the stub is invoked on the first hasNext()/next(). This keeps request construction free of side effects and matches how callers naturally write the loop.

  • MlStreamChunk mirrors the REST inference_results shape, so a chunk deserializes to the same structure users already know from POST /_plugins/_ml/models/<id>/_predict/stream. content() and isLast() are convenience shortcuts into the output named response. They select by name, not by position, because agent chunks carry three outputs — memory_id and parent_interaction_id (each holding only a result) precede response (which holds dataAsMap):

{"inference_results":[{"output":[
  {"name":"memory_id","result":"LvU1iJkBCzHrriq5hXbN"},
  {"name":"parent_interaction_id","result":"L_U1iJkBCzHrriq5hXbs"},
  {"name":"response","dataAsMap":{"content":"...","is_last":false}}
]}]}

Error Handling

Iterator cannot throw checked exceptions, and gRPC streaming reports failures during iteration rather than at call time. Errors are converted through the existing GrpcStatusConverter and surfaced as:

gRPC status Java exception
NOT_FOUND, INVALID_ARGUMENT, other status-bearing errors OpenSearchException with status() (e.g. 404)
UNIMPLEMENTED UnsupportedOperationException — MLService not registered, or blocked by a managed-service allowlist
UNAVAILABLE, transport failures UncheckedIOException wrapping TransportException

TransportException extends IOException, so it is wrapped rather than dropped; the cause is preserved. Chunks already delivered before an error remain valid — the iterator does not discard them.

Request Parameters

MlParameters supports messages, inputs, question, and _llm_interface. Which one applies is determined by the connector's request_body template, not by the protocol: the OpenAI connector interpolates ${parameters.messages}, the Bedrock converse connector interpolates ${parameters.inputs}, and agents use question. All are exposed so both provider families work; the request body is omitted entirely when no parameters are set.

Issues Resolved

List any issues this PR will resolve, e.g. Closes [...].

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.

@nathaliellenaa
nathaliellenaa force-pushed the ml_grpc branch 2 times, most recently from 603d04f to 4a160f8 Compare July 29, 2026 17:05
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
}
}

// ─── ML Streaming APIs ───────────────────────────────────────────────────────

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@nathaliellenaa thank you for the pull request.

A few things here:

  • GrpcTransport is a low level generic transport implementation (see please OpenSearchTransport what the expected primitives should be exposed)
  • XxxTransport classes are not supposed to be used in isolation, they are passed to OpenSearchClient which provides targeted methods (bulk, mget, ...) or/and clients (OpenSearchClusterClient, OpenSearchIndexClient, ...), in fact we already have one for ML as well - OpenSearchMlClient

The transport itself should not expose anything specific to particular API (like ml fe). The elephant in the room here is that OpenSearchTransport does not support streaming (but underlying clients do https://github.com/opensearch-project/OpenSearch/blob/main/client/rest/src/main/java/org/opensearch/client/RestClient.java#L324), so we have to close this gap first (let me pick up shortly). Once we have the streaming implementation, we could looking into bringing it to OpenSearchMlClient.

Hope it makes sense, thanks

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @reta, it makes sense that GrpcTransport shouldn't be exposing ML-specific methods. I went that route because of the gap you identified: Transport#performRequest returns a single ResponseT, so there was no way to express a server-streaming RPC returning a sequence of chunks through the existing primitives.
Happy to wait for the streaming implementation and then rebase this on top of it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants