rest stream predict api

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

Author: jngz-es

common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java

@@ -36,6 +36,9 @@
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
 
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
+
 import lombok.Builder;
 import lombok.EqualsAndHashCode;
 import lombok.NoArgsConstructor;
@@ -346,6 +349,10 @@ public <T> T createPayload(String action, Map<String, String> parameters) {
 
             if (!isJson(payload)) {
                 throw new IllegalArgumentException("Invalid payload: " + payload);
+            } else if (parameters.containsKey("stream")) {
+                JsonObject jsonObject = JsonParser.parseString(payload).getAsJsonObject();
+                jsonObject.addProperty("stream", true);
+                payload = jsonObject.toString();
             }
             return (T) payload;
         }

ml-algorithms/build.gradle

@@ -14,13 +14,22 @@ plugins {
 }
 
 repositories {
+    mavenLocal()
     mavenCentral()
 }
 
 dependencies {
     implementation project(path: ":${rootProject.name}-spi", configuration: 'shadow')
     implementation project(path: ":${rootProject.name}-common", configuration: 'shadow')
     implementation project(':opensearch-ml-memory')
+
+    implementation "org.opensearch:opensearch-arrow-spi:${opensearch_version}"
+    implementation "org.apache.arrow:arrow-vector:${versions.arrow}"
+    implementation "org.apache.arrow:arrow-format:${versions.arrow}"
+    implementation "org.apache.arrow:arrow-memory-core:${versions.arrow}"
+    runtimeOnly "org.apache.arrow:arrow-memory-netty:${versions.arrow}"
+    runtimeOnly "org.apache.arrow:arrow-memory-netty-buffer-patch:${versions.arrow}"
+
     compileOnly group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
     implementation "org.opensearch.client:opensearch-rest-client:${opensearch_version}"
     testImplementation "org.opensearch.test:framework:${opensearch_version}"
@@ -88,6 +97,8 @@ dependencies {
     testImplementation("com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}")
     testImplementation("com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}")
     testImplementation group: 'com.networknt' , name: 'json-schema-validator', version: '1.4.0'
+    api group: 'com.squareup.okhttp3', name: 'okhttp', version: '4.12.0'
+    implementation group: 'com.squareup.okhttp3', name: 'okhttp-sse', version: '4.12.0'
 }
 
 lombok {

ml-algorithms/src/main/java/org/opensearch/ml/engine/MLEngine.java

@@ -11,7 +11,9 @@
 import java.nio.file.Path;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Supplier;
 
+import org.opensearch.arrow.spi.StreamManager;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLModel;
@@ -49,11 +51,14 @@ public class MLEngine {
 
     private Encryptor encryptor;
 
-    public MLEngine(Path opensearchDataFolder, Encryptor encryptor) {
+    private Supplier<StreamManager> streamManager;
+
+    public MLEngine(Path opensearchDataFolder, Encryptor encryptor, Supplier<StreamManager> streamManager) {
         this.mlCachePath = opensearchDataFolder.resolve("ml_cache");
         this.mlModelsCachePath = mlCachePath.resolve("models_cache");
         this.mlConfigPath = mlCachePath.resolve("config");
         this.encryptor = encryptor;
+        this.streamManager = streamManager;
     }
 
     public String getPrebuiltModelMetaListPath() {
@@ -141,7 +146,7 @@ public Map<String, String> getConnectorCredential(Connector connector) {
 
     public Predictable deploy(MLModel mlModel, Map<String, Object> params) {
         Predictable predictable = MLEngineClassLoader.initInstance(mlModel.getAlgorithm(), null, MLAlgoParams.class);
-        predictable.initModel(mlModel, params, encryptor);
+        predictable.initModel(mlModel, params, encryptor, streamManager);
         return predictable;
     }
 

ml-algorithms/src/main/java/org/opensearch/ml/engine/Predictable.java

@@ -6,7 +6,9 @@
 package org.opensearch.ml.engine;
 
 import java.util.Map;
+import java.util.function.Supplier;
 
+import org.opensearch.arrow.spi.StreamManager;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.input.MLInput;
@@ -47,7 +49,13 @@ default void asyncPredict(MLInput mlInput, ActionListener<MLTaskResponse> action
      * @param params other parameters
      * @param encryptor encryptor
      */
-    void initModel(MLModel model, Map<String, Object> params, Encryptor encryptor);
+    default void initModel(MLModel model, Map<String, Object> params, Encryptor encryptor) {
+
+    };
+
+    default void initModel(MLModel model, Map<String, Object> params, Encryptor encryptor, Supplier<StreamManager> streamManager) {
+
+    };
 
     /**
      * Close resources like deployed model.

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java

@@ -17,13 +17,15 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
 
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.bulk.BackoffPolicy;
 import org.opensearch.action.support.GroupedActionListener;
 import org.opensearch.action.support.RetryableAction;
+import org.opensearch.arrow.spi.StreamManager;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.unit.TimeValue;
@@ -221,6 +223,8 @@ && getUserRateLimiterMap().get(user.getName()) != null
             }
             if (getConnectorClientConfig().getMaxRetryTimes() != 0) {
                 invokeRemoteServiceWithRetry(action, mlInput, parameters, payload, executionContext, actionListener);
+            } else if (parameters.containsKey("stream")) {
+                invokeRemoteServiceStream(action, mlInput, parameters, payload, executionContext, actionListener);
             } else {
                 invokeRemoteService(action, mlInput, parameters, payload, executionContext, actionListener);
             }
@@ -337,4 +341,19 @@ class RetryableActionExtensionArgs {
         private final ExecutionContext executionContext;
         private final String payload;
     }
+
+    void invokeRemoteServiceStream(
+        String action,
+        MLInput mlInput,
+        Map<String, String> parameters,
+        String payload,
+        ExecutionContext executionContext,
+        ActionListener<Tuple<Integer, ModelTensors>> actionListener
+    );
+
+    default void setStreamManager(Supplier<StreamManager> streamManager) {}
+
+    default Supplier<StreamManager> getStreamManager() {
+        return null;
+    };
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteModel.java

@@ -9,7 +9,9 @@
 
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
 
+import org.opensearch.arrow.spi.StreamManager;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.util.TokenBucket;
 import org.opensearch.core.action.ActionListener;
@@ -98,7 +100,7 @@ public boolean isModelReady() {
     }
 
     @Override
-    public void initModel(MLModel model, Map<String, Object> params, Encryptor encryptor) {
+    public void initModel(MLModel model, Map<String, Object> params, Encryptor encryptor, Supplier<StreamManager> streamManager) {
         try {
             Connector connector = model.getConnector().cloneConnector();
             connector
@@ -112,6 +114,7 @@ public void initModel(MLModel model, Map<String, Object> params, Encryptor encry
             this.connectorExecutor.setUserRateLimiterMap((Map<String, TokenBucket>) params.get(USER_RATE_LIMITER_MAP));
             this.connectorExecutor.setMlGuard((MLGuard) params.get(GUARDRAILS));
             this.connectorExecutor.setConnectorPrivateIpEnabled((AtomicBoolean) params.get(CONNECTOR_PRIVATE_IP_ENABLED));
+            this.connectorExecutor.setStreamManager(streamManager);
         } catch (RuntimeException e) {
             log.error("Failed to init remote model.", e);
             throw e;

plugin/build.gradle

@@ -57,6 +57,19 @@ dependencies {
     implementation project(':opensearch-ml-algorithms')
     implementation project(':opensearch-ml-search-processors')
     implementation project(':opensearch-ml-memory')
+
+    implementation "org.opensearch:opensearch-arrow-spi:${opensearchVersion}"
+    implementation "org.apache.arrow:arrow-vector:${versions.arrow}"
+    implementation "org.apache.arrow:arrow-format:${versions.arrow}"
+    implementation "org.apache.arrow:arrow-memory-core:${versions.arrow}"
+    runtimeOnly "org.apache.arrow:arrow-memory-netty:${versions.arrow}"
+    runtimeOnly "org.apache.arrow:arrow-memory-netty-buffer-patch:${versions.arrow}"
+
+    implementation "io.netty:netty-buffer:${versions.netty}"
+    implementation "io.netty:netty-common:${versions.netty}"
+    compileOnly 'org.checkerframework:checker-qual:3.44.0'
+
+
     compileOnly "com.google.guava:guava:32.1.3-jre"
 
     implementation group: 'software.amazon.awssdk', name: 'aws-core', version: "2.30.18"
@@ -424,6 +437,16 @@ configurations.all {
     resolutionStrategy.force "jakarta.json:jakarta.json-api:2.1.3"
     resolutionStrategy.force "org.opensearch:opensearch:${opensearch_version}"
     resolutionStrategy.force "org.bouncycastle:bcprov-jdk18on:1.78.1"
+    resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib-jdk8:1.9.10"
+    resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib:1.9.10"
+    resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib-jdk7:1.9.10"
+    resolutionStrategy.force "org.jetbrains.kotlin:kotlin-stdlib-common:1.9.10"
+    resolutionStrategy.force "org.checkerframework:checker-qual:3.44.0"
+    resolutionStrategy.force 'io.netty:netty-buffer:4.1.118.Final'
+    resolutionStrategy.force 'io.netty:netty-common:4.1.118.Final'
+    resolutionStrategy.force 'com.fasterxml.jackson.core:jackson-annotations:2.18.2'
+    resolutionStrategy.force 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.18.2'
+    resolutionStrategy.force 'com.google.flatbuffers:flatbuffers-java:24.3.25'
 }
 
 apply plugin: 'com.netflix.nebula.ospackage'

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

@@ -40,6 +40,7 @@
 import java.util.function.Supplier;
 
 import org.opensearch.action.ActionRequest;
+import org.opensearch.arrow.spi.StreamManager;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.cluster.service.ClusterService;
@@ -312,6 +313,7 @@
 import org.opensearch.plugins.Plugin;
 import org.opensearch.plugins.SearchPipelinePlugin;
 import org.opensearch.plugins.SearchPlugin;
+import org.opensearch.plugins.StreamManagerPlugin;
 import org.opensearch.plugins.SystemIndexPlugin;
 import org.opensearch.remote.metadata.client.SdkClient;
 import org.opensearch.remote.metadata.client.impl.SdkClientFactory;
@@ -343,7 +345,8 @@ public class MachineLearningPlugin extends Plugin
         SearchPipelinePlugin,
         ExtensiblePlugin,
         IngestPlugin,
-        SystemIndexPlugin {
+        SystemIndexPlugin,
+        StreamManagerPlugin {
     public static final String ML_THREAD_POOL_PREFIX = "thread_pool.ml_commons.";
     public static final String GENERAL_THREAD_POOL = "opensearch_ml_general";
     public static final String SDK_CLIENT_THREAD_POOL = "opensearch_ml_sdkclient";
@@ -354,6 +357,7 @@ public class MachineLearningPlugin extends Plugin
     public static final String INGEST_THREAD_POOL = "opensearch_ml_ingest";
     public static final String REGISTER_THREAD_POOL = "opensearch_ml_register";
     public static final String DEPLOY_THREAD_POOL = "opensearch_ml_deploy";
+    public static final String STREAM_PREDICT_THREAD_POOL = "opensearch_ml_predict_stream";
     public static final String ML_BASE_URI = "/_plugins/_ml";
 
     private MLStats mlStats;
@@ -398,6 +402,14 @@ public class MachineLearningPlugin extends Plugin
     private ScriptService scriptService;
     private Encryptor encryptor;
 
+    private StreamManager streamManager;
+
+    private StreamManager getStreamManagerRef() {
+        return this.streamManager;
+    }
+
+    private Supplier<StreamManager> streamManagerSupplier = () -> { return getStreamManagerRef(); };
+
     public MachineLearningPlugin(Settings settings) {
         // Handle this here as this feature is tied to Search/Query API, not to a ml-common API
         // and as such, it can't be lazy-loaded when a ml-commons API is invoked.
@@ -523,7 +535,7 @@ public Collection<Object> createComponents(
 
         encryptor = new EncryptorImpl(clusterService, client, sdkClient, mlIndicesHandler);
 
-        mlEngine = new MLEngine(dataPath, encryptor);
+        mlEngine = new MLEngine(dataPath, encryptor, streamManagerSupplier);
         nodeHelper = new DiscoveryNodeHelper(clusterService, settings);
         modelCacheHelper = new MLModelCacheHelper(clusterService, settings);
         cmHandler = new OpenSearchConversationalMemoryHandler(client, clusterService);
@@ -753,7 +765,8 @@ public Collection<Object> createComponents(
                 mlCircuitBreakerService,
                 mlModelAutoRedeployer,
                 cmHandler,
-                sdkClient
+                sdkClient,
+                streamManagerSupplier
             );
     }
 
@@ -771,6 +784,11 @@ public List<RestHandler> getRestHandlers(
         RestMLTrainingAction restMLTrainingAction = new RestMLTrainingAction();
         RestMLTrainAndPredictAction restMLTrainAndPredictAction = new RestMLTrainAndPredictAction();
         RestMLPredictionAction restMLPredictionAction = new RestMLPredictionAction(mlModelManager, mlFeatureEnabledSetting);
+        RestMLPredictionStreamingAction restMLPredictionStreamingAction = new RestMLPredictionStreamingAction(
+            mlModelManager,
+            mlFeatureEnabledSetting,
+            streamManagerSupplier
+        );
         RestMLExecuteAction restMLExecuteAction = new RestMLExecuteAction(mlFeatureEnabledSetting);
         RestMLGetModelAction restMLGetModelAction = new RestMLGetModelAction(mlFeatureEnabledSetting);
         RestMLDeleteModelAction restMLDeleteModelAction = new RestMLDeleteModelAction(mlFeatureEnabledSetting);
@@ -835,6 +853,7 @@ public List<RestHandler> getRestHandlers(
                 restMLStatsAction,
                 restMLTrainingAction,
                 restMLPredictionAction,
+                restMLPredictionStreamingAction,
                 restMLExecuteAction,
                 restMLTrainAndPredictAction,
                 restMLGetModelAction,
@@ -964,6 +983,14 @@ public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {
             ML_THREAD_POOL_PREFIX + INGEST_THREAD_POOL,
             false
         );
+        FixedExecutorBuilder streamPredictThreadPool = new FixedExecutorBuilder(
+            settings,
+            STREAM_PREDICT_THREAD_POOL,
+            OpenSearchExecutors.allocatedProcessors(settings) * 2,
+            10000,
+            ML_THREAD_POOL_PREFIX + STREAM_PREDICT_THREAD_POOL,
+            false
+        );
 
         return ImmutableList
             .of(
@@ -975,7 +1002,8 @@ public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {
                 predictThreadPool,
                 remotePredictThreadPool,
                 batchIngestThreadPool,
-                sdkClientThreadPool
+                sdkClientThreadPool,
+                streamPredictThreadPool
             );
     }
 
@@ -1174,4 +1202,10 @@ public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings sett
         systemIndexDescriptors.add(new SystemIndexDescriptor(ML_STOP_WORDS_INDEX, "ML Commons Stop Words Index"));
         return systemIndexDescriptors;
     }
+
+    @Override
+    public void onStreamManagerInitialized(Supplier<StreamManager> streamManager) {
+        this.streamManager = streamManager.get();
+    }
+
 }