Fix IT

Signed-off-by: b4sjoo <[email protected]>
b4sjoo · Sep 18, 2024 · 6edda9b · 6edda9b
1 parent 4acea8d
commit 6edda9b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 20 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -47,7 +47,6 @@
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_MAX_REGISTER_MODEL_TASKS_PER_NODE;
 import static org.opensearch.ml.stats.ActionName.REGISTER;
 import static org.opensearch.ml.stats.MLActionLevelStat.ML_ACTION_REQUEST_COUNT;
-import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
 import static org.opensearch.ml.utils.MLExceptionUtils.logException;
 import static org.opensearch.ml.utils.MLNodeUtils.checkOpenCircuitBreaker;
 import static org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
@@ -132,7 +131,6 @@
 import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.engine.utils.FileUtils;
 import org.opensearch.ml.profile.MLModelProfile;
-import org.opensearch.ml.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.stats.ActionName;
 import org.opensearch.ml.stats.MLActionLevelStat;
 import org.opensearch.ml.stats.MLNodeLevelStat;
@@ -180,7 +178,6 @@ public class MLModelManager {
     private final MLTaskManager mlTaskManager;
     private final MLEngine mlEngine;
     private final DiscoveryNodeHelper nodeHelper;
-    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     private volatile Integer maxModelPerNode;
     private volatile Integer maxRegisterTasksPerNode;
@@ -211,8 +208,7 @@ public MLModelManager(
         MLTaskManager mlTaskManager,
         MLModelCacheHelper modelCacheHelper,
         MLEngine mlEngine,
-        DiscoveryNodeHelper nodeHelper,
-        MLFeatureEnabledSetting mlFeatureEnabledSetting
+        DiscoveryNodeHelper nodeHelper
     ) {
         this.client = client;
         this.sdkClient = sdkClient;
@@ -228,7 +224,6 @@ public MLModelManager(
         this.mlTaskManager = mlTaskManager;
         this.mlEngine = mlEngine;
         this.nodeHelper = nodeHelper;
-        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
 
         this.maxModelPerNode = ML_COMMONS_MAX_MODELS_PER_NODE.get(settings);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MAX_MODELS_PER_NODE, it -> maxModelPerNode = it);
@@ -1334,9 +1329,6 @@ public synchronized void updateModelCache(String modelId, ActionListener<String>
      */
     public synchronized void deployControllerWithDeployedModel(String modelId, ActionListener<String> listener) {
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-            if (!mlFeatureEnabledSetting.isControllerEnabled()) {
-                throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
-            }
             if (!modelCacheHelper.isModelDeployed(modelId)) {
                 throw new OpenSearchStatusException(
                     "The model of this model controller has not deployed yet, please deploy the model first.",
@@ -1506,9 +1498,6 @@ private synchronized void deployControllerWithDeployingModel(
      * @param mlModel ml model
      */
     public void deployControllerWithDeployingModel(MLModel mlModel, Integer eligibleNodeCount) {
-        if (!mlFeatureEnabledSetting.isControllerEnabled()) {
-            throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
-        }
         if (mlModel.getModelState() != MLModelState.DEPLOYING) {
             throw new OpenSearchStatusException(
                 "This method should only be called when model is in DEPLOYING state, but the model is in state: " + mlModel.getModelState(),

diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -511,8 +511,7 @@ public Collection<Object> createComponents(
             mlTaskManager,
             modelCacheHelper,
             mlEngine,
-            nodeHelper,
-            mlFeatureEnabledSetting
+            nodeHelper
         );
         mlInputDatasetHandler = new MLInputDatasetHandler(client);
         modelAccessControlHelper = new ModelAccessControlHelper(clusterService, settings);

diff --git a/plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java b/plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java
@@ -122,7 +122,6 @@
 import org.opensearch.ml.engine.encryptor.EncryptorImpl;
 import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.sdkclient.SdkClientFactory;
-import org.opensearch.ml.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.stats.ActionName;
 import org.opensearch.ml.stats.MLActionLevelStat;
 import org.opensearch.ml.stats.MLNodeLevelStat;
@@ -209,8 +208,6 @@ public class MLModelManagerTests extends OpenSearchTestCase {
 
     @Mock
     private MLTask pretrainedMLTask;
-    @Mock
-    MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     @Before
     public void setup() throws URISyntaxException, IOException {
@@ -287,7 +284,6 @@ public void setup() throws URISyntaxException, IOException {
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
         when(threadPool.executor(any())).thenReturn(testThreadPool.executor(GENERAL_THREAD_POOL));
-        when(mlFeatureEnabledSetting.isControllerEnabled()).thenReturn(true);
 
         modelManager = spy(
             new MLModelManager(
@@ -305,8 +301,7 @@ public void setup() throws URISyntaxException, IOException {
                 mlTaskManager,
                 modelCacheHelper,
                 mlEngine,
-                nodeHelper,
-                mlFeatureEnabledSetting
+                nodeHelper
             )
         );