fix test failures

add feature flag

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

Author: jngz-es

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

@@ -362,6 +362,9 @@ public <T> T createPayload(String action, Map<String, String> parameters) {
     }
 
     private boolean neededStreamParameterInPayload(Map<String, String> parameters) {
+        if (parameters == null) {
+            return false;
+        }
         boolean isStream = parameters.containsKey("stream");
         if (!isStream) {
             return false;

common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java

@@ -342,4 +342,8 @@ private MLCommonsSettings() {}
     /** This setting sets the remote metadata service name */
     public static final Setting<String> REMOTE_METADATA_SERVICE_NAME = Setting
         .simpleString("plugins.ml_commons." + REMOTE_METADATA_SERVICE_NAME_KEY, Setting.Property.NodeScope, Setting.Property.Final);
+
+    /** This setting is to enable/disable streaming feature. */
+    public static final Setting<Boolean> ML_COMMONS_STREAM_ENABLED = Setting
+        .boolSetting("plugins.ml_commons.stream_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
 }

common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java

@@ -16,6 +16,7 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_REMOTE_INFERENCE_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STREAM_ENABLED;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -40,6 +41,7 @@ public class MLFeatureEnabledSetting {
 
     // This is to identify if this node is in multi-tenancy or not.
     private volatile Boolean isMultiTenancyEnabled;
+    private volatile Boolean isStreamEnabled;
 
     private volatile Boolean isMcpServerEnabled;
 
@@ -55,6 +57,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         isBatchInferenceEnabled = ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED.get(settings);
         isMultiTenancyEnabled = ML_COMMONS_MULTI_TENANCY_ENABLED.get(settings);
         isMcpServerEnabled = ML_COMMONS_MCP_SERVER_ENABLED.get(settings);
+        isStreamEnabled = ML_COMMONS_STREAM_ENABLED.get(settings);
 
         clusterService
             .getClusterSettings()
@@ -74,6 +77,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
             .getClusterSettings()
             .addSettingsUpdateConsumer(ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED, it -> isBatchInferenceEnabled = it);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MCP_SERVER_ENABLED, it -> isMcpServerEnabled = it);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_STREAM_ENABLED, it -> isStreamEnabled = it);
     }
 
     /**
@@ -144,6 +148,13 @@ public boolean isMcpServerEnabled() {
         return isMcpServerEnabled;
     }
 
+    /** Whether the streaming feature is enabled. If disabled, APIs in ml-commons will block stream.
+     * @return whether the streaming is enabled.
+     */
+    public boolean isStreamEnabled() {
+        return isStreamEnabled;
+    }
+
     public void addListener(SettingsChangeListener listener) {
         listeners.add(listener);
     }

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

@@ -154,7 +154,11 @@ 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, streamManager, threadPool);
+        if (mlModel.getAlgorithm() == FunctionName.REMOTE) {
+            predictable.initModel(mlModel, params, encryptor, streamManager, threadPool);
+        } else {
+            predictable.initModel(mlModel, params, encryptor);
+        }
         return predictable;
     }
 

ml-algorithms/src/main/java/org/opensearch/ml/engine/arrow/RemoteModelStreamProducer.java

@@ -1,3 +1,8 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 package org.opensearch.ml.engine.arrow;
 
 import java.nio.charset.StandardCharsets;

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AbstractConnectorExecutorTest.java

@@ -21,6 +21,10 @@ public class AbstractConnectorExecutorTest {
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+        when(mockConnector.getAccessKey()).thenReturn("access_key");
+        when(mockConnector.getSecretKey()).thenReturn("secret_key");
+        when(mockConnector.getSessionToken()).thenReturn("session_token");
+        when(mockConnector.getRegion()).thenReturn("us-east-1-test");
         executor = new AwsConnectorExecutor(mockConnector);
         connectorClientConfig = new ConnectorClientConfig();
     }