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

@@ -350,4 +350,8 @@ private MLCommonsSettings() {}
     // Feature flag for enabling telemetry static metric collection job -- MLStatsJobProcessor
     public static final Setting<Boolean> ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED = Setting
         .boolSetting("plugins.ml_commons.metrics_static_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+
+    /** 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

@@ -19,6 +19,7 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_REMOTE_INFERENCE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STREAM_ENABLED;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -43,6 +44,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;
 
@@ -66,6 +68,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         isRagSearchPipelineEnabled = ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED.get(settings);
         isMetricCollectionEnabled = ML_COMMONS_METRIC_COLLECTION_ENABLED.get(settings);
         isStaticMetricCollectionEnabled = ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED.get(settings);
+        isStreamEnabled = ML_COMMONS_STREAM_ENABLED.get(settings);
 
         clusterService
             .getClusterSettings()
@@ -94,6 +97,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         clusterService
             .getClusterSettings()
             .addSettingsUpdateConsumer(ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED, it -> isStaticMetricCollectionEnabled = it);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_STREAM_ENABLED, it -> isStreamEnabled = it);
     }
 
     /**
@@ -164,6 +168,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);
     }

common/src/test/java/org/opensearch/ml/common/settings/MLFeatureEnabledSettingTests.java

@@ -44,6 +44,7 @@ public void setUp() {
                     MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED,
                     MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED,
                     MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED
+                    MLCommonsSettings.ML_COMMONS_STREAM_ENABLED
                 )
         );
         when(mockClusterService.getClusterSettings()).thenReturn(mockClusterSettings);
@@ -65,6 +66,7 @@ public void testDefaults_allFeaturesEnabled() {
             .put("plugins.ml_commons.rag_pipeline_feature_enabled", true)
             .put("plugins.ml_commons.metrics_collection_enabled", true)
             .put("plugins.ml_commons.metrics_static_collection_enabled", true)
+            .put("plugins.ml_commons.stream_enabled", true)
             .build();
 
         MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
@@ -81,6 +83,7 @@ public void testDefaults_allFeaturesEnabled() {
         assertTrue(setting.isRagSearchPipelineEnabled());
         assertTrue(setting.isMetricCollectionEnabled());
         assertTrue(setting.isStaticMetricCollectionEnabled());
+        assertTrue(setting.isStreamEnabled());
     }
 
     @Test
@@ -99,6 +102,7 @@ public void testDefaults_someFeaturesDisabled() {
             .put("plugins.ml_commons.rag_pipeline_feature_enabled", false)
             .put("plugins.ml_commons.metrics_collection_enabled", false)
             .put("plugins.ml_commons.metrics_static_collection_enabled", false)
+            .put("plugins.ml_commons.stream_enabled", false)
             .build();
 
         MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
@@ -115,6 +119,7 @@ public void testDefaults_someFeaturesDisabled() {
         assertFalse(setting.isRagSearchPipelineEnabled());
         assertFalse(setting.isMetricCollectionEnabled());
         assertFalse(setting.isStaticMetricCollectionEnabled());
+        assertFalse(setting.isStreamEnabled());
     }
 
     @Test

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();
     }