Add ML gRPC support - #2070
Conversation
603d04f to
4a160f8
Compare
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
| } | ||
| } | ||
|
|
||
| // ─── ML Streaming APIs ─────────────────────────────────────────────────────── |
There was a problem hiding this comment.
@nathaliellenaa thank you for the pull request.
A few things here:
GrpcTransportis a low level generic transport implementation (see pleaseOpenSearchTransportwhat the expected primitives should be exposed)XxxTransportclasses are not supposed to be used in isolation, they are passed toOpenSearchClientwhich 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
There was a problem hiding this comment.
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.
Description
This PR adds gRPC ML streaming support to the
java-client-grpcmodule, exposing the ML Commons server-streaming APIsMLService/PredictModelStreamandMLService/ExecuteAgentStreamasGrpcTransport.predictModelStream()andGrpcTransport.executeAgentStream().Sample usage for model prediction:
Agent execution follows the same shape:
How it works
MlStreamRequestConverterbuilds the protobuf request from the typed builder; the blocking stub returns anIterator<PredictResponse>;MlStreamResponseConverterconverts each protobuf chunk intoMlStreamChunkon demand as the caller advances the iterator. Nothing is buffered — one protobuf message in, oneMlStreamChunkout.Two behaviors are worth calling out:
The RPC is lazy. Calling
predictModelStream()does not issue the request; the stub is invoked on the firsthasNext()/next(). This keeps request construction free of side effects and matches how callers naturally write the loop.MlStreamChunkmirrors the RESTinference_resultsshape, so a chunk deserializes to the same structure users already know fromPOST /_plugins/_ml/models/<id>/_predict/stream.content()andisLast()are convenience shortcuts into the output namedresponse. They select by name, not by position, because agent chunks carry three outputs —memory_idandparent_interaction_id(each holding only aresult) precederesponse(which holdsdataAsMap):{"inference_results":[{"output":[ {"name":"memory_id","result":"LvU1iJkBCzHrriq5hXbN"}, {"name":"parent_interaction_id","result":"L_U1iJkBCzHrriq5hXbs"}, {"name":"response","dataAsMap":{"content":"...","is_last":false}} ]}]}Error Handling
Iteratorcannot throw checked exceptions, and gRPC streaming reports failures during iteration rather than at call time. Errors are converted through the existingGrpcStatusConverterand surfaced as:NOT_FOUND,INVALID_ARGUMENT, other status-bearing errorsOpenSearchExceptionwithstatus()(e.g. 404)UNIMPLEMENTEDUnsupportedOperationException— MLService not registered, or blocked by a managed-service allowlistUNAVAILABLE, transport failuresUncheckedIOExceptionwrappingTransportExceptionTransportException 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
MlParameterssupportsmessages,inputs,question, and_llm_interface. Which one applies is determined by the connector'srequest_bodytemplate, not by the protocol: the OpenAI connector interpolates${parameters.messages}, the Bedrock converse connector interpolates${parameters.inputs}, and agents usequestion. 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.