From 6edda9bebd81f5fd21af5b0ec5728ac8bb571e57 Mon Sep 17 00:00:00 2001 From: b4sjoo Date: Wed, 18 Sep 2024 22:08:50 +0000 Subject: [PATCH] Fix IT Signed-off-by: b4sjoo --- .../org/opensearch/ml/model/MLModelManager.java | 13 +------------ .../opensearch/ml/plugin/MachineLearningPlugin.java | 3 +-- .../opensearch/ml/model/MLModelManagerTests.java | 7 +------ 3 files changed, 3 insertions(+), 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 index eea7dcc8de..45f80f5113 100644 --- 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 */ public synchronized void deployControllerWithDeployedModel(String modelId, ActionListener 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 index 843c670779..1e15c8a4e5 100644 --- 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 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 index baec2f0026..835440040b 100644 --- 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 ) );