Skip to content

Commit 6bff20d

Browse files
owaiskazi19rbhavnaylwu-amzndhrubo-osb4sjoo
authored
Move allow model setting from rest to transport (opensearch-project#1961)
* Backport multiple PRs to main from 2.x (opensearch-project#1652) * fix parameter name in preprocess function; fix remote model function … (opensearch-project#1362) * fix parameter name in preprocess function; fix remote model function name Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix failed unit test Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> * throw exception when model group not found during update request (opensearch-project#1447) Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * add status code to model tensor (opensearch-project#1443) (opensearch-project#1453) Signed-off-by: Yaliang Wu <ylwu@amazon.com> * register new versions to a model group based on the name provided (opensearch-project#1452) Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * fixing metrics correlation algorithm (opensearch-project#1448) * fixing metrics correlation algorithm Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> * if model version fails to register, update model group accordingly (opensearch-project#1463) * if model version fails to register, update model group accordingly Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * Update Model API (opensearch-project#1350) * Update Model API POC Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Using GetRequest to get model Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Finalize model update API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compile Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compileTest Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add Unit Test Cases for Update Model API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Tune back test coverage thereshold Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add more unit tests on Update model API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add unit test for TransportUpdateModelAction class Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix a test error Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change exception thrown to failure response Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Move the function judgement to the outter block Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Check if model is undeployed before update model Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add more unit test for update model API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix unit test due to blocking java 11 CI workflow Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Enabling auto bumping model version during registering to a new model group and address reviewers' other concern Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Autobump new model groups' latest version when register to a new model Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change the REST API method from POST to PUT Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change the update REST API endpoint Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add a setting to control the update connector API (opensearch-project#1465) * Add a setting to control the update connector API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Enabling the update connnector setting in unit test Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Enabling the update connnector setting in corresponding unit test Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * fix update connector API (opensearch-project#1484) * fix update connector API Signed-off-by: Yaliang Wu <ylwu@amazon.com> * Performance enhacement for predict action by caching model info (opensearch-project#1472) (opensearch-project#1508) * Performance enhacement for predict action by caching model info Signed-off-by: zane-neo <zaniu@amazon.com> * Add context.restore() to avoid missing info Signed-off-by: zane-neo <zaniu@amazon.com> --------- Signed-off-by: zane-neo <zaniu@amazon.com> (cherry picked from commit a985f6e) Co-authored-by: zane-neo <zaniu@amazon.com> * fix failed ut from PR 1472 (opensearch-project#1479) (opensearch-project#1510) * fix failed ut from PR 1472 Signed-off-by: Yaliang Wu <ylwu@amazon.com> * exclude class for low coverage Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit da5d829) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * [Backport to 2.11] throw exception if remote model doesn't return 2xx status code; fix p… (opensearch-project#1477) (opensearch-project#1509) * throw exception if remote model doesn't return 2xx status code; fix p… (opensearch-project#1473) * throw exception if remote model doesn't return 2xx status code; fix predict runner Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix kmeans model deploy bug Signed-off-by: Yaliang Wu <ylwu@amazon.com> * support multiple docs for remote embedding model Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix ut Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> * fix wrong class Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 201c8a8) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * fix no worker node exception for remote embedding model (opensearch-project#1482) (opensearch-project#1511) * fix no worker node exception for remote embedding model Signed-off-by: Yaliang Wu <ylwu@amazon.com> * only add model info to cache if model cache exist Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 6f83b9f) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * fix for delete model group API throwing incorrect error when model index not created (opensearch-project#1485) (opensearch-project#1486) (opensearch-project#1512) * fix for delete model group API throwing incorrect error when model index not created Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> (cherry picked from commit 60ef0fd) Co-authored-by: Bhavana Ramaram <rbhavna@amazon.com> (cherry picked from commit 5544681) Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> * fix no worker node error on multi-node cluster (opensearch-project#1487) (opensearch-project#1513) Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit cea1cd6) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * add prefix to show the error is from remote service (opensearch-project#1499) (opensearch-project#1515) Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 3897ad1) Co-authored-by: Yaliang Wu <ylwu@amazon.com> * fix multiple docs support (opensearch-project#1516) Signed-off-by: Yaliang Wu <ylwu@amazon.com> * adding another fix issue to the release note (opensearch-project#1498) (opensearch-project#1514) Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> (cherry picked from commit 440155c) Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> * add bedrockURL to trusted connector regex list (opensearch-project#1461) Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * return parsing exception 400 for parsing errors Signed-off-by: Xun Zhang <xunzh@amazon.com> * add more ut in restupdateconnector Signed-off-by: Xun Zhang <xunzh@amazon.com> * fix format violations Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> * Fix model/connector update API to address security concern (opensearch-project#1595) * Fix model/connector update API to address appsec concern Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compile and build failure Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Improve unit test coverage Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix spotless Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Merge update connector feature flag to remote inference feature flag Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix compile Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Fix exception status Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Keep fixing exception status Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Spotless fix Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add UT on parsing exception Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * change XContentFactory to MediaTypeRegistry builder in MLRegisterModelInputTest class Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Signed-off-by: Sicheng Song <sicheng.song@outlook.com> Signed-off-by: Xun Zhang <xunzh@amazon.com> Co-authored-by: Yaliang Wu <ylwu@amazon.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> Co-authored-by: Sicheng Song <sicheng.song@outlook.com> Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Co-authored-by: zane-neo <zaniu@amazon.com> Co-authored-by: Xun Zhang <xunzh@amazon.com> * Allow model setting to transport from rest Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> * Added test Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> * Fixed integ test Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Signed-off-by: Sicheng Song <sicheng.song@outlook.com> Signed-off-by: Xun Zhang <xunzh@amazon.com> Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> Co-authored-by: Bhavana Ramaram <rbhavna@amazon.com> Co-authored-by: Yaliang Wu <ylwu@amazon.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> Co-authored-by: Sicheng Song <sicheng.song@outlook.com> Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Co-authored-by: zane-neo <zaniu@amazon.com> Co-authored-by: Xun Zhang <xunzh@amazon.com>
1 parent 0fd07f0 commit 6bff20d

File tree

5 files changed

+65
-23
lines changed

5 files changed

+65
-23
lines changed

‎plugin/src/main/java/org/opensearch/ml/action/register/TransportRegisterModelAction.java

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

88
import static org.opensearch.ml.common.MLTask.STATE_FIELD;
99
import static org.opensearch.ml.common.MLTaskState.FAILED;
10+
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
1011
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX;
1112
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_URL_REGEX;
1213
import static org.opensearch.ml.task.MLTaskManager.TASK_SEMAPHORE_TIMEOUT;
@@ -89,6 +90,7 @@ public class TransportRegisterModelAction extends HandledTransportAction<ActionR
8990
private List<String> trustedConnectorEndpointsRegex;
9091

9192
ModelAccessControlHelper modelAccessControlHelper;
93+
private volatile boolean isModelUrlAllowed;
9294

9395
ConnectorAccessControlHelper connectorAccessControlHelper;
9496
MLModelGroupManager mlModelGroupManager;
@@ -132,6 +134,9 @@ public TransportRegisterModelAction(
132134
trustedUrlRegex = ML_COMMONS_TRUSTED_URL_REGEX.get(settings);
133135
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_TRUSTED_URL_REGEX, it -> trustedUrlRegex = it);
134136

137+
isModelUrlAllowed = ML_COMMONS_ALLOW_MODEL_URL.get(settings);
138+
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_ALLOW_MODEL_URL, it -> isModelUrlAllowed = it);
139+
135140
trustedConnectorEndpointsRegex = ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX.get(settings);
136141
clusterService
137142
.getClusterSettings()
@@ -142,6 +147,11 @@ public TransportRegisterModelAction(
142147
protected void doExecute(Task task, ActionRequest request, ActionListener<MLRegisterModelResponse> listener) {
143148
MLRegisterModelRequest registerModelRequest = MLRegisterModelRequest.fromActionRequest(request);
144149
MLRegisterModelInput registerModelInput = registerModelRequest.getRegisterModelInput();
150+
if (registerModelInput.getUrl() != null && !isModelUrlAllowed) {
151+
throw new IllegalArgumentException(
152+
"To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use OpenSearch pre-trained models."
153+
);
154+
}
145155
registerModelInput.setIsHidden(RestActionUtils.isSuperAdminUser(clusterService, client));
146156
if (StringUtils.isEmpty(registerModelInput.getModelGroupId())) {
147157
mlModelGroupManager.validateUniqueModelGroupName(registerModelInput.getModelName(), ActionListener.wrap(modelGroups -> {

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

-9
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
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.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
1110
import static org.opensearch.ml.utils.MLExceptionUtils.REMOTE_INFERENCE_DISABLED_ERR_MSG;
1211
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_DEPLOY_MODEL;
1312
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
@@ -35,7 +34,6 @@
3534

3635
public class RestMLRegisterModelAction extends BaseRestHandler {
3736
private static final String ML_REGISTER_MODEL_ACTION = "ml_register_model_action";
38-
private volatile boolean isModelUrlAllowed;
3937
private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
4038

4139
/**
@@ -51,8 +49,6 @@ public RestMLRegisterModelAction(MLFeatureEnabledSetting mlFeatureEnabledSetting
5149
* @param settings settings
5250
*/
5351
public RestMLRegisterModelAction(ClusterService clusterService, Settings settings, MLFeatureEnabledSetting mlFeatureEnabledSetting) {
54-
isModelUrlAllowed = ML_COMMONS_ALLOW_MODEL_URL.get(settings);
55-
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_ALLOW_MODEL_URL, it -> isModelUrlAllowed = it);
5652
this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
5753
}
5854

@@ -103,11 +99,6 @@ MLRegisterModelRequest getRequest(RestRequest request) throws IOException {
10399
if (mlInput.getFunctionName() == FunctionName.REMOTE && !mlFeatureEnabledSetting.isRemoteInferenceEnabled()) {
104100
throw new IllegalStateException(REMOTE_INFERENCE_DISABLED_ERR_MSG);
105101
}
106-
if (mlInput.getUrl() != null && !isModelUrlAllowed) {
107-
throw new IllegalArgumentException(
108-
"To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use opensearch pre-trained models."
109-
);
110-
}
111102
return new MLRegisterModelRequest(mlInput);
112103
}
113104
}

‎plugin/src/test/java/org/opensearch/ml/action/models/SearchModelITTests.java

+8
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,14 @@
55

66
package org.opensearch.ml.action.models;
77

8+
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
9+
810
import org.junit.Before;
911
import org.junit.Rule;
1012
import org.junit.rules.ExpectedException;
1113
import org.opensearch.action.search.SearchRequest;
1214
import org.opensearch.action.search.SearchResponse;
15+
import org.opensearch.common.settings.Settings;
1316
import org.opensearch.index.query.BoolQueryBuilder;
1417
import org.opensearch.index.query.QueryBuilders;
1518
import org.opensearch.ml.action.MLCommonsIntegTestCase;
@@ -187,4 +190,9 @@ private void test_matchPhrase_search() {
187190
assertEquals(1, response.getHits().getTotalHits().value);
188191
}
189192

193+
@Override
194+
protected Settings nodeSettings(int ordinal) {
195+
return Settings.builder().put(super.nodeSettings(ordinal)).put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), true).build();
196+
}
197+
190198
}

‎plugin/src/test/java/org/opensearch/ml/action/register/TransportRegisterModelActionTests.java

+47
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import static org.mockito.Mockito.mock;
1515
import static org.mockito.Mockito.verify;
1616
import static org.mockito.Mockito.when;
17+
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
1718
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX;
1819
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_URL_REGEX;
1920
import static org.opensearch.ml.utils.TestHelper.clusterSetting;
@@ -155,12 +156,14 @@ public void setup() throws IOException {
155156
settings = Settings
156157
.builder()
157158
.put(ML_COMMONS_TRUSTED_URL_REGEX.getKey(), trustedUrlRegex)
159+
.put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), true)
158160
.putList(ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX.getKey(), TRUSTED_CONNECTOR_ENDPOINTS_REGEXES)
159161
.build();
160162
threadContext = new ThreadContext(settings);
161163
ClusterSettings clusterSettings = clusterSetting(
162164
settings,
163165
ML_COMMONS_TRUSTED_URL_REGEX,
166+
ML_COMMONS_ALLOW_MODEL_URL,
164167
ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX
165168
);
166169
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
@@ -294,6 +297,50 @@ public void testDoExecute_invalidURL() {
294297
assertEquals("URL can't match trusted url regex", argumentCaptor.getValue().getMessage());
295298
}
296299

300+
public void testRegisterModelUrlNotAllowed() throws Exception {
301+
Settings settings = Settings
302+
.builder()
303+
.put(ML_COMMONS_TRUSTED_URL_REGEX.getKey(), trustedUrlRegex)
304+
.put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), false)
305+
.putList(ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX.getKey(), TRUSTED_CONNECTOR_ENDPOINTS_REGEXES)
306+
.build();
307+
ClusterSettings clusterSettings = clusterSetting(
308+
settings,
309+
ML_COMMONS_TRUSTED_URL_REGEX,
310+
ML_COMMONS_ALLOW_MODEL_URL,
311+
ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX
312+
);
313+
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
314+
when(clusterService.getSettings()).thenReturn(settings);
315+
transportRegisterModelAction = new TransportRegisterModelAction(
316+
transportService,
317+
actionFilters,
318+
modelHelper,
319+
mlIndicesHandler,
320+
mlModelManager,
321+
mlTaskManager,
322+
clusterService,
323+
settings,
324+
threadPool,
325+
client,
326+
nodeFilter,
327+
mlTaskDispatcher,
328+
mlStats,
329+
modelAccessControlHelper,
330+
connectorAccessControlHelper,
331+
mlModelGroupManager
332+
);
333+
334+
IllegalArgumentException e = assertThrows(
335+
IllegalArgumentException.class,
336+
() -> transportRegisterModelAction.doExecute(task, prepareRequest("test url", "testModelGroupsID"), actionListener)
337+
);
338+
assertEquals(
339+
e.getMessage(),
340+
"To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use OpenSearch pre-trained models."
341+
);
342+
}
343+
297344
public void testDoExecute_successWithLocalNodeNotEqualToClusterNode() {
298345
when(node1.getId()).thenReturn("NodeId1");
299346
when(node2.getId()).thenReturn("NodeId2");

‎plugin/src/test/java/org/opensearch/ml/rest/RestMLRegisterModelActionTests.java

-14
Original file line numberDiff line numberDiff line change
@@ -147,20 +147,6 @@ public void testRegisterModelRequestRemoteInferenceDisabled() throws Exception {
147147
restMLRegisterModelAction.handleRequest(request, channel, client);
148148
}
149149

150-
public void testRegisterModelUrlNotAllowed() throws Exception {
151-
settings = Settings.builder().put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), false).build();
152-
ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_ALLOW_MODEL_URL);
153-
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
154-
restMLRegisterModelAction = new RestMLRegisterModelAction(clusterService, settings, mlFeatureEnabledSetting);
155-
exceptionRule.expect(IllegalArgumentException.class);
156-
exceptionRule
157-
.expectMessage(
158-
"To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use opensearch pre-trained models."
159-
);
160-
RestRequest request = getRestRequest();
161-
restMLRegisterModelAction.handleRequest(request, channel, client);
162-
}
163-
164150
public void testRegisterModelRequestWithNullUrlAndUrlNotAllowed() throws Exception {
165151
settings = Settings.builder().put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), false).build();
166152
ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_ALLOW_MODEL_URL);

0 commit comments

Comments
 (0)