Skip to content

Commit 2f08772

Browse files
b4sjoob4sjoo
committed
Add feature enable setting for controller index
Signed-off-by: b4sjoo <sicheng.song@outlook.com>
1 parent 66d8e2b commit 2f08772

22 files changed

+280
-30
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
// TODO: Add support for multi tenancy
103111
mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {

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

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

88
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
99
import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
10+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
1011

1112
import java.util.ArrayList;
1213
import java.util.Arrays;
@@ -40,6 +41,7 @@
4041
import org.opensearch.ml.helper.ModelAccessControlHelper;
4142
import org.opensearch.ml.model.MLModelCacheHelper;
4243
import org.opensearch.ml.model.MLModelManager;
44+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
4345
import org.opensearch.ml.utils.RestActionUtils;
4446
import org.opensearch.tasks.Task;
4547
import org.opensearch.transport.TransportService;
@@ -57,6 +59,7 @@ public class DeleteControllerTransportAction extends HandledTransportAction<Acti
5759
MLModelManager mlModelManager;
5860
MLModelCacheHelper mlModelCacheHelper;
5961
ModelAccessControlHelper modelAccessControlHelper;
62+
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
6063

6164
@Inject
6265
public DeleteControllerTransportAction(
@@ -67,7 +70,8 @@ public DeleteControllerTransportAction(
6770
ClusterService clusterService,
6871
MLModelManager mlModelManager,
6972
MLModelCacheHelper mlModelCacheHelper,
70-
ModelAccessControlHelper modelAccessControlHelper
73+
ModelAccessControlHelper modelAccessControlHelper,
74+
MLFeatureEnabledSetting mlFeatureEnabledSetting
7175
) {
7276
super(MLControllerDeleteAction.NAME, transportService, actionFilters, MLControllerDeleteRequest::new);
7377
this.client = client;
@@ -76,6 +80,7 @@ public DeleteControllerTransportAction(
7680
this.mlModelManager = mlModelManager;
7781
this.mlModelCacheHelper = mlModelCacheHelper;
7882
this.modelAccessControlHelper = modelAccessControlHelper;
83+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
7984
}
8085

8186
@Override
@@ -85,6 +90,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
8590
User user = RestActionUtils.getUserContext(client);
8691
String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
8792
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
93+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
94+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
95+
}
8896
ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
8997
// TODO: Add support for multi tenancy
9098
mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {

‎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
// TODO: Add support for multi tenancy
95103
mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {

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

+12-1
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_MAX_REGISTER_MODEL_TASKS_PER_NODE;
4848
import static org.opensearch.ml.stats.ActionName.REGISTER;
4949
import static org.opensearch.ml.stats.MLActionLevelStat.ML_ACTION_REQUEST_COUNT;
50+
import static org.opensearch.ml.utils.MLExceptionUtils.CONTROLLER_DISABLED_ERR_MSG;
5051
import static org.opensearch.ml.utils.MLExceptionUtils.logException;
5152
import static org.opensearch.ml.utils.MLNodeUtils.checkOpenCircuitBreaker;
5253
import static org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
@@ -131,6 +132,7 @@
131132
import org.opensearch.ml.engine.indices.MLIndicesHandler;
132133
import org.opensearch.ml.engine.utils.FileUtils;
133134
import org.opensearch.ml.profile.MLModelProfile;
135+
import org.opensearch.ml.settings.MLFeatureEnabledSetting;
134136
import org.opensearch.ml.stats.ActionName;
135137
import org.opensearch.ml.stats.MLActionLevelStat;
136138
import org.opensearch.ml.stats.MLNodeLevelStat;
@@ -178,6 +180,7 @@ public class MLModelManager {
178180
private final MLTaskManager mlTaskManager;
179181
private final MLEngine mlEngine;
180182
private final DiscoveryNodeHelper nodeHelper;
183+
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
181184

182185
private volatile Integer maxModelPerNode;
183186
private volatile Integer maxRegisterTasksPerNode;
@@ -208,7 +211,8 @@ public MLModelManager(
208211
MLTaskManager mlTaskManager,
209212
MLModelCacheHelper modelCacheHelper,
210213
MLEngine mlEngine,
211-
DiscoveryNodeHelper nodeHelper
214+
DiscoveryNodeHelper nodeHelper,
215+
MLFeatureEnabledSetting mlFeatureEnabledSetting
212216
) {
213217
this.client = client;
214218
this.sdkClient = sdkClient;
@@ -224,6 +228,7 @@ public MLModelManager(
224228
this.mlTaskManager = mlTaskManager;
225229
this.mlEngine = mlEngine;
226230
this.nodeHelper = nodeHelper;
231+
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
227232

228233
this.maxModelPerNode = ML_COMMONS_MAX_MODELS_PER_NODE.get(settings);
229234
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MAX_MODELS_PER_NODE, it -> maxModelPerNode = it);
@@ -1329,6 +1334,9 @@ public synchronized void updateModelCache(String modelId, ActionListener<String>
13291334
*/
13301335
public synchronized void deployControllerWithDeployedModel(String modelId, ActionListener<String> listener) {
13311336
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
1337+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
1338+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
1339+
}
13321340
if (!modelCacheHelper.isModelDeployed(modelId)) {
13331341
throw new OpenSearchStatusException(
13341342
"The model of this model controller has not deployed yet, please deploy the model first.",
@@ -1498,6 +1506,9 @@ private synchronized void deployControllerWithDeployingModel(
14981506
* @param mlModel ml model
14991507
*/
15001508
public void deployControllerWithDeployingModel(MLModel mlModel, Integer eligibleNodeCount) {
1509+
if (!mlFeatureEnabledSetting.isControllerEnabled()) {
1510+
throw new IllegalStateException(CONTROLLER_DISABLED_ERR_MSG);
1511+
}
15011512
if (mlModel.getModelState() != MLModelState.DEPLOYING) {
15021513
throw new OpenSearchStatusException(
15031514
"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

+7-5
Original file line numberDiff line numberDiff line change
@@ -511,7 +511,8 @@ public Collection<Object> createComponents(
511511
mlTaskManager,
512512
modelCacheHelper,
513513
mlEngine,
514-
nodeHelper
514+
nodeHelper,
515+
mlFeatureEnabledSetting
515516
);
516517
mlInputDatasetHandler = new MLInputDatasetHandler(client);
517518
modelAccessControlHelper = new ModelAccessControlHelper(clusterService, settings);
@@ -750,10 +751,10 @@ public List<RestHandler> getRestHandlers(
750751
RestMemorySearchInteractionsAction restSearchInteractionsAction = new RestMemorySearchInteractionsAction();
751752
RestMemoryGetConversationAction restGetConversationAction = new RestMemoryGetConversationAction();
752753
RestMemoryGetInteractionAction restGetInteractionAction = new RestMemoryGetInteractionAction();
753-
RestMLCreateControllerAction restMLCreateControllerAction = new RestMLCreateControllerAction();
754-
RestMLGetControllerAction restMLGetControllerAction = new RestMLGetControllerAction();
755-
RestMLUpdateControllerAction restMLUpdateControllerAction = new RestMLUpdateControllerAction();
756-
RestMLDeleteControllerAction restMLDeleteControllerAction = new RestMLDeleteControllerAction();
754+
RestMLCreateControllerAction restMLCreateControllerAction = new RestMLCreateControllerAction(mlFeatureEnabledSetting);
755+
RestMLGetControllerAction restMLGetControllerAction = new RestMLGetControllerAction(mlFeatureEnabledSetting);
756+
RestMLUpdateControllerAction restMLUpdateControllerAction = new RestMLUpdateControllerAction(mlFeatureEnabledSetting);
757+
RestMLDeleteControllerAction restMLDeleteControllerAction = new RestMLDeleteControllerAction(mlFeatureEnabledSetting);
757758
RestMLGetAgentAction restMLGetAgentAction = new RestMLGetAgentAction(mlFeatureEnabledSetting);
758759
RestMLDeleteAgentAction restMLDeleteAgentAction = new RestMLDeleteAgentAction(mlFeatureEnabledSetting);
759760
RestMemoryUpdateConversationAction restMemoryUpdateConversationAction = new RestMemoryUpdateConversationAction();
@@ -946,6 +947,7 @@ public List<Setting<?>> getSettings() {
946947
MLCommonsSettings.ML_COMMONS_AGENT_FRAMEWORK_ENABLED,
947948
MLCommonsSettings.ML_COMMONS_MODEL_AUTO_DEPLOY_ENABLE,
948949
MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED,
950+
MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED,
949951
// Settings for SdkClient
950952
SdkClientSettings.REMOTE_METADATA_TYPE,
951953
SdkClientSettings.REMOTE_METADATA_ENDPOINT,

‎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
}

0 commit comments

Comments
 (0)