Skip to content

Commit a72181a

Browse files
b4sjoob4sjoo
authored
Add feature enable setting for controller index (opensearch-project#2652)
Signed-off-by: b4sjoo <sicheng.song@outlook.com>
1 parent 0d26931 commit a72181a

22 files changed

+269
-29
lines changed

‎plugin/src/main/java/org/opensearch/ml/action/controller/CreateControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import static org.opensearch.ml.common.FunctionName.REMOTE;
1010
import static org.opensearch.ml.common.FunctionName.TEXT_EMBEDDING;
1111
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
12+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1213

1314
import java.util.ArrayList;
1415
import java.util.Arrays;
@@ -51,6 +52,7 @@
5152
import org.opensearch.ml.helper.ModelAccessControlHelper;
5253
import org.opensearch.ml.model.MLModelCacheHelper;
5354
import org.opensearch.ml.model.MLModelManager;
55+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
5456
import org.opensearch.ml.utils.RestActionUtils;
5557
import org.opensearch.tasks.Task;
5658
import org.opensearch.transport.TransportService;
@@ -68,6 +70,7 @@ public class CreateControllerTransportAction extends HandledTransportAction<Acti
6870
ClusterService clusterService;
6971
MLModelCacheHelper mlModelCacheHelper;
7072
ModelAccessControlHelper modelAccessControlHelper;
73+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
7174

7275
@Inject
7376
public CreateControllerTransportAction(
@@ -78,7 +81,8 @@ public CreateControllerTransportAction(
7881
ClusterService clusterService,
7982
ModelAccessControlHelper modelAccessControlHelper,
8083
MLModelCacheHelper mlModelCacheHelper,
81-
MLModelManager mlModelManager
84+
MLModelManager mlModelManager,
85+
MLFeatureEnabledSetting mlFeatureEnabledSetting
8286
) {
8387
super(MLCreateControllerAction.NAME, transportService, actionFilters, MLCreateControllerRequest::new);
8488
this.mlIndicesHandler = mlIndicesHandler;
@@ -87,6 +91,7 @@ public CreateControllerTransportAction(
8791
this.clusterService = clusterService;
8892
this.mlModelCacheHelper = mlModelCacheHelper;
8993
this.modelAccessControlHelper = modelAccessControlHelper;
94+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
9095
}
9196

9297
@Override
@@ -98,6 +103,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLCrea
98103
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
99104

100105
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
106+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
107+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
108+
}
101109
ActionListener<MLCreateControllerResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
102110
mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
103111
FunctionName functionName = mlModel.getAlgorithm();

‎plugin/src/main/java/org/opensearch/ml/action/controller/DeleteControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
99
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1010
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
11+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1112

1213
import java.util.ArrayList;
1314
import java.util.Arrays;
@@ -41,6 +42,7 @@
4142
import org.opensearch.ml.helper.ModelAccessControlHelper;
4243
import org.opensearch.ml.model.MLModelCacheHelper;
4344
import org.opensearch.ml.model.MLModelManager;
45+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
4446
import org.opensearch.ml.utils.RestActionUtils;
4547
import org.opensearch.tasks.Task;
4648
import org.opensearch.transport.TransportService;
@@ -58,6 +60,7 @@ public class DeleteControllerTransportAction extends HandledTransportAction<Acti
5860
MLModelManager mlModelManager;
5961
MLModelCacheHelper mlModelCacheHelper;
6062
ModelAccessControlHelper modelAccessControlHelper;
63+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
6164

6265
@Inject
6366
public DeleteControllerTransportAction(
@@ -68,7 +71,8 @@ public DeleteControllerTransportAction(
6871
ClusterService clusterService,
6972
MLModelManager mlModelManager,
7073
MLModelCacheHelper mlModelCacheHelper,
71-
ModelAccessControlHelper modelAccessControlHelper
74+
ModelAccessControlHelper modelAccessControlHelper,
75+
MLFeatureEnabledSetting mlFeatureEnabledSetting
7276
) {
7377
super(MLControllerDeleteAction.NAME, transportService, actionFilters, MLControllerDeleteRequest::new);
7478
this.client = client;
@@ -77,6 +81,7 @@ public DeleteControllerTransportAction(
7781
this.mlModelManager = mlModelManager;
7882
this.mlModelCacheHelper = mlModelCacheHelper;
7983
this.modelAccessControlHelper = modelAccessControlHelper;
84+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
8085
}
8186

8287
@Override
@@ -86,6 +91,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
8691
User user = RestActionUtils.getUserContext(client);
8792
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
8893
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
94+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
95+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
96+
}
8997
ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
9098
mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
9199
Boolean isHidden = mlModel.getIsHidden();

‎plugin/src/main/java/org/opensearch/ml/action/controller/GetControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
99
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1010
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
11+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1112
import static org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
1213
import static org.opensearch.ml.utils.RestActionUtils.getFetchSourceContext;
1314

@@ -33,6 +34,7 @@
3334
import org.opensearch.ml.common.transport.controller.MLControllerGetResponse;
3435
import org.opensearch.ml.helper.ModelAccessControlHelper;
3536
import org.opensearch.ml.model.MLModelManager;
37+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
3638
import org.opensearch.ml.utils.RestActionUtils;
3739
import org.opensearch.search.fetch.subphase.FetchSourceContext;
3840
import org.opensearch.tasks.Task;
@@ -50,6 +52,7 @@ public class GetControllerTransportAction extends HandledTransportAction<ActionR
5052
ClusterService clusterService;
5153
MLModelManager mlModelManager;
5254
ModelAccessControlHelper modelAccessControlHelper;
55+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
5356

5457
@Inject
5558
public GetControllerTransportAction(
@@ -59,14 +62,16 @@ public GetControllerTransportAction(
5962
NamedXContentRegistry xContentRegistry,
6063
ClusterService clusterService,
6164
MLModelManager mlModelManager,
62-
ModelAccessControlHelper modelAccessControlHelper
65+
ModelAccessControlHelper modelAccessControlHelper,
66+
MLFeatureEnabledSetting mlFeatureEnabledSetting
6367
) {
6468
super(MLControllerGetAction.NAME, transportService, actionFilters, MLControllerGetRequest::new);
6569
this.client = client;
6670
this.xContentRegistry = xContentRegistry;
6771
this.clusterService = clusterService;
6872
this.mlModelManager = mlModelManager;
6973
this.modelAccessControlHelper = modelAccessControlHelper;
74+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
7075
}
7176

7277
@Override
@@ -79,6 +84,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLCont
7984
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
8085

8186
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
87+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
88+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
89+
}
8290
ActionListener<MLControllerGetResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
8391
client.get(getRequest, ActionListener.wrap(r -> {
8492
if (r != null && r.isExists()) {

‎plugin/src/main/java/org/opensearch/ml/action/controller/UpdateControllerTransportAction.java

+9-1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import static org.opensearch.ml.common.FunctionName.REMOTE;
1010
import static org.opensearch.ml.common.FunctionName.TEXT_EMBEDDING;
1111
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
12+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1213

1314
import java.util.ArrayList;
1415
import java.util.Arrays;
@@ -46,6 +47,7 @@
4647
import org.opensearch.ml.helper.ModelAccessControlHelper;
4748
import org.opensearch.ml.model.MLModelCacheHelper;
4849
import org.opensearch.ml.model.MLModelManager;
50+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
4951
import org.opensearch.ml.utils.RestActionUtils;
5052
import org.opensearch.tasks.Task;
5153
import org.opensearch.transport.TransportService;
@@ -62,6 +64,7 @@ public class UpdateControllerTransportAction extends HandledTransportAction<Acti
6264
MLModelCacheHelper mlModelCacheHelper;
6365
ClusterService clusterService;
6466
ModelAccessControlHelper modelAccessControlHelper;
67+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
6568

6669
@Inject
6770
public UpdateControllerTransportAction(
@@ -71,14 +74,16 @@ public UpdateControllerTransportAction(
7174
ClusterService clusterService,
7275
ModelAccessControlHelper modelAccessControlHelper,
7376
MLModelCacheHelper mlModelCacheHelper,
74-
MLModelManager mlModelManager
77+
MLModelManager mlModelManager,
78+
MLFeatureEnabledSetting mlFeatureEnabledSetting
7579
) {
7680
super(MLUpdateControllerAction.NAME, transportService, actionFilters, MLUpdateControllerRequest::new);
7781
this.client = client;
7882
this.mlModelManager = mlModelManager;
7983
this.clusterService = clusterService;
8084
this.mlModelCacheHelper = mlModelCacheHelper;
8185
this.modelAccessControlHelper = modelAccessControlHelper;
86+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
8287
}
8388

8489
@Override
@@ -90,6 +95,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
9095
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
9196

9297
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
98+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
99+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
100+
}
93101
ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
94102
mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
95103
FunctionName functionName = mlModel.getAlgorithm();

‎plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java

+7
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_MAX_REGISTER_MODEL_TASKS_PER_NODE;
4646
import static org.opensearch.ml.stats.ActionName.REGISTER;
4747
import static org.opensearch.ml.stats.MLActionLevelStat.ML_ACTION_REQUEST_COUNT;
48+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
4849
import static org.opensearch.ml.utils.MLExceptionUtils.logException;
4950
import static org.opensearch.ml.utils.MLNodeUtils.checkOpenCircuitBreaker;
5051
import static org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
@@ -1254,6 +1255,9 @@ public synchronized void updateModelCache(String modelId, ActionListener<String>
12541255
*/
12551256
public synchronized void deployControllerWithDeployedModel(String modelId, ActionListener<String> listener) {
12561257
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
1258+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
1259+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
1260+
}
12571261
if (!modelCacheHelper.isModelDeployed(modelId)) {
12581262
throw new OpenSearchStatusException(
12591263
"The model of this model controller has not deployed yet, please deploy the model first.",
@@ -1423,6 +1427,9 @@ private synchronized void deployControllerWithDeployingModel(
14231427
* @param mlModel ml model
14241428
*/
14251429
public void deployControllerWithDeployingModel(MLModel mlModel, Integer eligibleNodeCount) {
1430+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
1431+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
1432+
}
14261433
if (mlModel.getModelState() != MLModelState.DEPLOYING) {
14271434
throw new OpenSearchStatusException(
14281435
"This method should only be called when model is in DEPLOYING state, but the model is in state: " + mlModel.getModelState(),

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

+6-5
Original file line numberDiff line numberDiff line change
@@ -755,10 +755,10 @@ public List<RestHandler> getRestHandlers(
755755
RestMemorySearchInteractionsAction restSearchInteractionsAction = new RestMemorySearchInteractionsAction();
756756
RestMemoryGetConversationAction restGetConversationAction = new RestMemoryGetConversationAction();
757757
RestMemoryGetInteractionAction restGetInteractionAction = new RestMemoryGetInteractionAction();
758-
RestMLCreateControllerAction restMLCreateControllerAction = new RestMLCreateControllerAction();
759-
RestMLGetControllerAction restMLGetControllerAction = new RestMLGetControllerAction();
760-
RestMLUpdateControllerAction restMLUpdateControllerAction = new RestMLUpdateControllerAction();
761-
RestMLDeleteControllerAction restMLDeleteControllerAction = new RestMLDeleteControllerAction();
758+
RestMLCreateControllerAction restMLCreateControllerAction = new RestMLCreateControllerAction(mlFeatureEnabledSetting);
759+
RestMLGetControllerAction restMLGetControllerAction = new RestMLGetControllerAction(mlFeatureEnabledSetting);
760+
RestMLUpdateControllerAction restMLUpdateControllerAction = new RestMLUpdateControllerAction(mlFeatureEnabledSetting);
761+
RestMLDeleteControllerAction restMLDeleteControllerAction = new RestMLDeleteControllerAction(mlFeatureEnabledSetting);
762762
RestMLGetAgentAction restMLGetAgentAction = new RestMLGetAgentAction(mlFeatureEnabledSetting);
763763
RestMLDeleteAgentAction restMLDeleteAgentAction = new RestMLDeleteAgentAction(mlFeatureEnabledSetting);
764764
RestMemoryUpdateConversationAction restMemoryUpdateConversationAction = new RestMemoryUpdateConversationAction();
@@ -969,7 +969,8 @@ public List<Setting<?>> getSettings() {
969969
MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_COMPLETED_REGEX,
970970
MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_CANCELLED_REGEX,
971971
MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_CANCELLING_REGEX,
972-
MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_EXPIRED_REGEX
972+
MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_EXPIRED_REGEX,
973+
MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED
973974
);
974975
return settings;
975976
}

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLCreateControllerAction.java

+10-1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
99
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
10+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1011
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
1112
import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
1213

@@ -20,6 +21,7 @@
2021
import org.opensearch.ml.common.controller.MLController;
2122
import org.opensearch.ml.common.transport.controller.MLCreateControllerAction;
2223
import org.opensearch.ml.common.transport.controller.MLCreateControllerRequest;
24+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
2325
import org.opensearch.rest.BaseRestHandler;
2426
import org.opensearch.rest.RestRequest;
2527
import org.opensearch.rest.action.RestToXContentListener;
@@ -29,11 +31,14 @@
2931
public class RestMLCreateControllerAction extends BaseRestHandler {
3032

3133
public final static String ML_CREATE_CONTROLLER_ACTION = "ml_create_controller_action";
34+
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
3235

3336
/**
3437
* Constructor
3538
*/
36-
public RestMLCreateControllerAction() {}
39+
public RestMLCreateControllerAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
40+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
41+
}
3742

3843
@Override
3944
public String getName() {
@@ -61,6 +66,10 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
6166
* @return MLCreateControllerRequest
6267
*/
6368
private MLCreateControllerRequest getRequest(RestRequest request) throws IOException {
69+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
70+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
71+
}
72+
6473
if (!request.hasContent()) {
6574
throw new OpenSearchParseException("Create model controller request has empty body");
6675
}

‎plugin/src/main/java/org/opensearch/ml/rest/RestMLDeleteControllerAction.java

+10-1
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.ml.rest;
77

88
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
9+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
910
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
1011

1112
import java.io.IOException;
@@ -15,6 +16,7 @@
1516
import org.opensearch.client.node.NodeClient;
1617
import org.opensearch.ml.common.transport.controller.MLControllerDeleteAction;
1718
import org.opensearch.ml.common.transport.controller.MLControllerDeleteRequest;
19+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
1820
import org.opensearch.rest.BaseRestHandler;
1921
import org.opensearch.rest.RestRequest;
2022
import org.opensearch.rest.action.RestToXContentListener;
@@ -25,9 +27,13 @@
2527
* This class consists of the REST handler to delete ML Model.
2628
*/
2729
public class RestMLDeleteControllerAction extends BaseRestHandler {
30+
2831
private static final String ML_DELETE_CONTROLLER_ACTION = "ml_delete_controller_action";
32+
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
2933

30-
public void RestMLDeleteControllerAction() {}
34+
public RestMLDeleteControllerAction(MLFeatureEnabledSetting mlFeatureEnabledSetting) {
35+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
36+
}
3137

3238
@Override
3339
public String getName() {
@@ -42,6 +48,9 @@ public List<Route> routes() {
4248

4349
@Override
4450
protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
51+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
52+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
53+
}
4554
String modelId = request.param(PARAMETER_MODEL_ID);
4655

4756
MLControllerDeleteRequest mlControllerDeleteRequest = new MLControllerDeleteRequest(modelId);

0 commit comments

Comments
 (0)