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>
austintlee · Mar 18, 2024 · 082b249 · 082b249
1 parent aa4f1ed
commit 082b249
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 23 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/register/TransportRegisterModelAction.java b/plugin/src/main/java/org/opensearch/ml/action/register/TransportRegisterModelAction.java
@@ -7,6 +7,7 @@
 
 import static org.opensearch.ml.common.MLTask.STATE_FIELD;
 import static org.opensearch.ml.common.MLTaskState.FAILED;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX;
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_URL_REGEX;
 import static org.opensearch.ml.task.MLTaskManager.TASK_SEMAPHORE_TIMEOUT;
@@ -89,6 +90,7 @@ public class TransportRegisterModelAction extends HandledTransportAction<ActionR
     private List<String> trustedConnectorEndpointsRegex;
 
     ModelAccessControlHelper modelAccessControlHelper;
+    private volatile boolean isModelUrlAllowed;
 
     ConnectorAccessControlHelper connectorAccessControlHelper;
     MLModelGroupManager mlModelGroupManager;
@@ -132,6 +134,9 @@ public TransportRegisterModelAction(
         trustedUrlRegex = ML_COMMONS_TRUSTED_URL_REGEX.get(settings);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_TRUSTED_URL_REGEX, it -> trustedUrlRegex = it);
 
+        isModelUrlAllowed = ML_COMMONS_ALLOW_MODEL_URL.get(settings);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_ALLOW_MODEL_URL, it -> isModelUrlAllowed = it);
+
         trustedConnectorEndpointsRegex = ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX.get(settings);
         clusterService
             .getClusterSettings()
@@ -142,6 +147,11 @@ public TransportRegisterModelAction(
     protected void doExecute(Task task, ActionRequest request, ActionListener<MLRegisterModelResponse> listener) {
         MLRegisterModelRequest registerModelRequest = MLRegisterModelRequest.fromActionRequest(request);
         MLRegisterModelInput registerModelInput = registerModelRequest.getRegisterModelInput();
+        if (registerModelInput.getUrl() != null && !isModelUrlAllowed) {
+            throw new IllegalArgumentException(
+                "To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use OpenSearch pre-trained models."
+            );
+        }
         registerModelInput.setIsHidden(RestActionUtils.isSuperAdminUser(clusterService, client));
         if (StringUtils.isEmpty(registerModelInput.getModelGroupId())) {
             mlModelGroupManager.validateUniqueModelGroupName(registerModelInput.getModelName(), ActionListener.wrap(modelGroups -> {

diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLRegisterModelAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLRegisterModelAction.java
@@ -7,7 +7,6 @@
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
-import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
 import static org.opensearch.ml.utils.MLExceptionUtils.REMOTE_INFERENCE_DISABLED_ERR_MSG;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_DEPLOY_MODEL;
 import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_ID;
@@ -35,7 +34,6 @@
 
 public class RestMLRegisterModelAction extends BaseRestHandler {
     private static final String ML_REGISTER_MODEL_ACTION = "ml_register_model_action";
-    private volatile boolean isModelUrlAllowed;
     private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     /**
@@ -51,8 +49,6 @@ public RestMLRegisterModelAction(MLFeatureEnabledSetting mlFeatureEnabledSetting
      * @param settings settings
      */
     public RestMLRegisterModelAction(ClusterService clusterService, Settings settings, MLFeatureEnabledSetting mlFeatureEnabledSetting) {
-        isModelUrlAllowed = ML_COMMONS_ALLOW_MODEL_URL.get(settings);
-        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_ALLOW_MODEL_URL, it -> isModelUrlAllowed = it);
         this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
     }
 
@@ -103,11 +99,6 @@ MLRegisterModelRequest getRequest(RestRequest request) throws IOException {
         if (mlInput.getFunctionName() == FunctionName.REMOTE && !mlFeatureEnabledSetting.isRemoteInferenceEnabled()) {
             throw new IllegalStateException(REMOTE_INFERENCE_DISABLED_ERR_MSG);
         }
-        if (mlInput.getUrl() != null && !isModelUrlAllowed) {
-            throw new IllegalArgumentException(
-                "To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use opensearch pre-trained models."
-            );
-        }
         return new MLRegisterModelRequest(mlInput);
     }
 }
diff --git a/plugin/src/test/java/org/opensearch/ml/action/models/SearchModelITTests.java b/plugin/src/test/java/org/opensearch/ml/action/models/SearchModelITTests.java
@@ -5,11 +5,14 @@
 
 package org.opensearch.ml.action.models;
 
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.rules.ExpectedException;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.ml.action.MLCommonsIntegTestCase;
@@ -187,4 +190,9 @@ private void test_matchPhrase_search() {
         assertEquals(1, response.getHits().getTotalHits().value);
     }
 
+    @Override
+    protected Settings nodeSettings(int ordinal) {
+        return Settings.builder().put(super.nodeSettings(ordinal)).put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), true).build();
+    }
+
 }
diff --git a/...in/src/test/java/org/opensearch/ml/action/register/TransportRegisterModelActionTests.java b/...in/src/test/java/org/opensearch/ml/action/register/TransportRegisterModelActionTests.java
@@ -14,6 +14,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ALLOW_MODEL_URL;
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX;
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_URL_REGEX;
 import static org.opensearch.ml.utils.TestHelper.clusterSetting;
@@ -155,12 +156,14 @@ public void setup() throws IOException {
         settings = Settings
             .builder()
             .put(ML_COMMONS_TRUSTED_URL_REGEX.getKey(), trustedUrlRegex)
+            .put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), true)
             .putList(ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX.getKey(), TRUSTED_CONNECTOR_ENDPOINTS_REGEXES)
             .build();
         threadContext = new ThreadContext(settings);
         ClusterSettings clusterSettings = clusterSetting(
             settings,
             ML_COMMONS_TRUSTED_URL_REGEX,
+            ML_COMMONS_ALLOW_MODEL_URL,
             ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX
         );
         when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
@@ -294,6 +297,50 @@ public void testDoExecute_invalidURL() {
         assertEquals("URL can't match trusted url regex", argumentCaptor.getValue().getMessage());
     }
 
+    public void testRegisterModelUrlNotAllowed() throws Exception {
+        Settings settings = Settings
+            .builder()
+            .put(ML_COMMONS_TRUSTED_URL_REGEX.getKey(), trustedUrlRegex)
+            .put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), false)
+            .putList(ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX.getKey(), TRUSTED_CONNECTOR_ENDPOINTS_REGEXES)
+            .build();
+        ClusterSettings clusterSettings = clusterSetting(
+            settings,
+            ML_COMMONS_TRUSTED_URL_REGEX,
+            ML_COMMONS_ALLOW_MODEL_URL,
+            ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX
+        );
+        when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
+        when(clusterService.getSettings()).thenReturn(settings);
+        transportRegisterModelAction = new TransportRegisterModelAction(
+            transportService,
+            actionFilters,
+            modelHelper,
+            mlIndicesHandler,
+            mlModelManager,
+            mlTaskManager,
+            clusterService,
+            settings,
+            threadPool,
+            client,
+            nodeFilter,
+            mlTaskDispatcher,
+            mlStats,
+            modelAccessControlHelper,
+            connectorAccessControlHelper,
+            mlModelGroupManager
+        );
+
+        IllegalArgumentException e = assertThrows(
+            IllegalArgumentException.class,
+            () -> transportRegisterModelAction.doExecute(task, prepareRequest("test url", "testModelGroupsID"), actionListener)
+        );
+        assertEquals(
+            e.getMessage(),
+            "To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use OpenSearch pre-trained models."
+        );
+    }
+
     public void testDoExecute_successWithLocalNodeNotEqualToClusterNode() {
         when(node1.getId()).thenReturn("NodeId1");
         when(node2.getId()).thenReturn("NodeId2");

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLRegisterModelActionTests.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLRegisterModelActionTests.java
@@ -147,20 +147,6 @@ public void testRegisterModelRequestRemoteInferenceDisabled() throws Exception {
         restMLRegisterModelAction.handleRequest(request, channel, client);
     }
 
-    public void testRegisterModelUrlNotAllowed() throws Exception {
-        settings = Settings.builder().put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), false).build();
-        ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_ALLOW_MODEL_URL);
-        when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
-        restMLRegisterModelAction = new RestMLRegisterModelAction(clusterService, settings, mlFeatureEnabledSetting);
-        exceptionRule.expect(IllegalArgumentException.class);
-        exceptionRule
-            .expectMessage(
-                "To upload custom model user needs to enable allow_registering_model_via_url settings. Otherwise please use opensearch pre-trained models."
-            );
-        RestRequest request = getRestRequest();
-        restMLRegisterModelAction.handleRequest(request, channel, client);
-    }
-
     public void testRegisterModelRequestWithNullUrlAndUrlNotAllowed() throws Exception {
         settings = Settings.builder().put(ML_COMMONS_ALLOW_MODEL_URL.getKey(), false).build();
         ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_ALLOW_MODEL_URL);