Set tenant ID for predict request (#2619)

* AWS DDB SDK client support for remote data store Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * AWS DDB SDK client support for remote data store Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * multi-tenancy for models (create, get, delete, update) + update connector (#2546) * multi-tenancy for models (create, get, delete) Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> * added update connector + update model Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> --------- Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * [Feature/multi_tenancy] Add source map to GetDataObjectResponse (#2489) * Add source map to GetDataObjectResponse Signed-off-by: Daniel Widdis <widdis@gmail.com> * Add test for map getter in clients Signed-off-by: Daniel Widdis <widdis@gmail.com> --------- Signed-off-by: Daniel Widdis <widdis@gmail.com> Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * [Feature/multi_tenancy] Add UpdateDataObject interface, Client, and Connector Implementations (#2520) * Restore original exception handling expectations Signed-off-by: Daniel Widdis <widdis@gmail.com> * Add UpdateDataObject to interface and implementations Signed-off-by: Daniel Widdis <widdis@gmail.com> * Implement UpdateConnector action Signed-off-by: Daniel Widdis <widdis@gmail.com> * Move CompletionException handling to a common method Signed-off-by: Daniel Widdis <widdis@gmail.com> * Add tests for SDKClient exceptions refactored from Transport Action Signed-off-by: Daniel Widdis <widdis@gmail.com> --------- Signed-off-by: Daniel Widdis <widdis@gmail.com> Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * Addressed CR comment Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * Added javadoc based on feedback Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * Set tenant ID for predict request Signed-off-by: Arjun kumar Giri <arjung@amazon.com> * Simplify instantiating Data Object Request/Response builders (#2608) Signed-off-by: Daniel Widdis <widdis@gmail.com> * Addressed comments Signed-off-by: Arjun kumar Giri <arjung@amazon.com> --------- Signed-off-by: Arjun kumar Giri <arjung@amazon.com> Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> Signed-off-by: Daniel Widdis <widdis@gmail.com> Signed-off-by: arjunkumargiri <142054468+arjunkumargiri@users.noreply.github.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com> Co-authored-by: Daniel Widdis <widdis@gmail.com>
opensearch-project · Jul 11, 2024 · 8c80f43 · 8c80f43
1 parent a421f49
commit 8c80f43
Show file tree

Hide file tree

Showing 33 changed files with 411 additions and 324 deletions.
diff --git a/.github/workflows/CI-workflow.yml b/.github/workflows/CI-workflow.yml
@@ -147,8 +147,8 @@ jobs:
       - name: Generate Password For Admin
         id: genpass
         run: |
-            PASSWORD=$(openssl rand -base64 20 | tr -dc 'A-Za-z0-9!@#$%^&*()_+=-')
-            echo "password={$PASSWORD}" >> $GITHUB_OUTPUT
+          PASSWORD=$(openssl rand -base64 20 | tr -dc 'A-Za-z0-9!@#$%^&*()_+=-')
+          echo "password={$PASSWORD}" >> $GITHUB_OUTPUT
       - name: Run Docker Image
         if: env.imagePresent == 'true'
         run: |

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java
@@ -180,7 +180,7 @@ public void initMLIndexIfAbsent(MLIndex index, ActionListener<Boolean> listener)
      */
     public void shouldUpdateIndex(String indexName, Integer newVersion, ActionListener<Boolean> listener) {
         IndexMetadata indexMetaData = clusterService.state().getMetadata().indices().get(indexName);
-        if (indexMetaData == null) {
+        if (indexMetaData == null || indexMetaData.mapping() == null) {
             listener.onResponse(Boolean.FALSE);
             return;
         }

diff --git a/...in/src/main/java/org/opensearch/ml/action/controller/CreateControllerTransportAction.java b/...in/src/main/java/org/opensearch/ml/action/controller/CreateControllerTransportAction.java
@@ -99,7 +99,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLCrea
 
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<MLCreateControllerResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
-            mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
+            // TODO: Add support for multi tenancy
+            mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {
                 FunctionName functionName = mlModel.getAlgorithm();
                 Boolean isHidden = mlModel.getIsHidden();
                 if (functionName == TEXT_EMBEDDING || functionName == REMOTE) {

diff --git a/...in/src/main/java/org/opensearch/ml/action/controller/DeleteControllerTransportAction.java b/...in/src/main/java/org/opensearch/ml/action/controller/DeleteControllerTransportAction.java
@@ -86,7 +86,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
         String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
-            mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
+            // TODO: Add support for multi tenancy
+            mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {
                 Boolean isHidden = mlModel.getIsHidden();
                 modelAccessControlHelper
                     .validateModelGroupAccess(user, mlModel.getModelGroupId(), client, ActionListener.wrap(hasPermission -> {

diff --git a/plugin/src/main/java/org/opensearch/ml/action/controller/GetControllerTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/controller/GetControllerTransportAction.java
@@ -85,7 +85,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLCont
                     try (XContentParser parser = createXContentParserFromRegistry(xContentRegistry, r.getSourceAsBytesRef())) {
                         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
                         MLController controller = MLController.parse(parser);
-                        mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
+                        // TODO: Add support for multi tenancy
+                        mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {
                             Boolean isHidden = mlModel.getIsHidden();
                             modelAccessControlHelper
                                 .validateModelGroupAccess(user, mlModel.getModelGroupId(), client, ActionListener.wrap(hasPermission -> {

diff --git a/...in/src/main/java/org/opensearch/ml/action/controller/UpdateControllerTransportAction.java b/...in/src/main/java/org/opensearch/ml/action/controller/UpdateControllerTransportAction.java
@@ -91,7 +91,8 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
 
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
-            mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
+            // TODO: Add support for multi tenancy
+            mlModelManager.getModel(modelId, null, null, excludes, ActionListener.wrap(mlModel -> {
                 FunctionName functionName = mlModel.getAlgorithm();
                 Boolean isHidden = mlModel.getIsHidden();
                 if (functionName == TEXT_EMBEDDING || functionName == REMOTE) {

diff --git a/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelAction.java b/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelAction.java
@@ -149,7 +149,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLDepl
 
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<MLDeployModelResponse> wrappedListener = ActionListener.runBefore(listener, context::restore);
-            mlModelManager.getModel(modelId, null, excludes, ActionListener.wrap(mlModel -> {
+            mlModelManager.getModel(modelId, tenantId, null, excludes, ActionListener.wrap(mlModel -> {
                 FunctionName functionName = mlModel.getAlgorithm();
                 Boolean isHidden = mlModel.getIsHidden();
                 if (!TenantAwareHelper.validateTenantResource(mlFeatureEnabledSetting, tenantId, mlModel.getTenantId(), listener)) {
@@ -285,6 +285,10 @@ private void deployModel(
             String taskId = response.getId();
             mlTask.setTaskId(taskId);
             if (algorithm == FunctionName.REMOTE) {
+                if (mlFeatureEnabledSetting.isMultiTenancyEnabled()) {
+                    listener.onResponse(new MLDeployModelResponse(taskId, MLTaskType.DEPLOY_MODEL, MLTaskState.CREATED.name()));
+                    return;
+                }
                 mlTaskManager.add(mlTask, eligibleNodeIds);
                 deployRemoteModel(mlModel, mlTask, localNodeId, eligibleNodes, deployToAllNodes, listener);
                 return;

diff --git a/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java b/plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java
@@ -224,6 +224,7 @@ private void deployModel(
             mlModelManager
                 .deployModel(
                     modelId,
+                    null,
                     modelContentHash,
                     functionName,
                     deployToAllNodes,

diff --git a/plugin/src/main/java/org/opensearch/ml/action/models/GetModelTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/models/GetModelTransportAction.java
@@ -103,6 +103,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLMode
             .builder()
             .index(ML_MODEL_INDEX)
             .id(modelId)
+            .tenantId(tenantId)
             .fetchSourceContext(fetchSourceContext)
             .build();
         User user = RestActionUtils.getUserContext(client);

diff --git a/plugin/src/main/java/org/opensearch/ml/action/models/UpdateModelTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/models/UpdateModelTransportAction.java
@@ -138,7 +138,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
 
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<UpdateResponse> wrappedListener = ActionListener.runBefore(actionListener, context::restore);
-            mlModelManager.getModel(sdkClient, modelId, null, excludes, ActionListener.wrap(mlModel -> {
+            mlModelManager.getModel(modelId, tenantId, null, excludes, ActionListener.wrap(mlModel -> {
                 if (TenantAwareHelper.validateTenantResource(mlFeatureEnabledSetting, tenantId, mlModel.getTenantId(), actionListener)) {
                     if (!isModelDeploying(mlModel.getModelState())) {
                         FunctionName functionName = mlModel.getAlgorithm();

diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
@@ -211,7 +211,7 @@ public void onFailure(Exception e) {
                 modelActionListener.onResponse(cachedMlModel);
             } else {
                 // For multi-node cluster, the function name is null in cache, so should always get model first.
-                mlModelManager.getModel(modelId, modelActionListener);
+                mlModelManager.getModel(modelId, tenantId, modelActionListener);
             }
         }
     }

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
@@ -419,6 +419,7 @@ private MLRegisterModelGroupInput createRegisterModelGroupRequest(MLRegisterMode
             .backendRoles(registerModelInput.getBackendRoles())
             .modelAccessMode(registerModelInput.getAccessMode())
             .isAddAllBackendRoles(registerModelInput.getAddAllBackendRoles())
+            .tenantId(registerModelInput.getTenantId())
             .build();
     }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelsAction.java b/plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelsAction.java
@@ -182,7 +182,7 @@ private void validateAccess(String modelId, String tenantId, ActionListener<Bool
         boolean isSuperAdmin = isSuperAdminUserWrapper(clusterService, client);
         String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-            mlModelManager.getModel(modelId, null, excludes, ActionListener.runBefore(ActionListener.wrap(mlModel -> {
+            mlModelManager.getModel(modelId, tenantId, null, excludes, ActionListener.runBefore(ActionListener.wrap(mlModel -> {
                 if (!TenantAwareHelper.validateTenantResource(mlFeatureEnabledSetting, tenantId, mlModel.getTenantId(), listener)) {
                     return;
                 }

diff --git a/plugin/src/main/java/org/opensearch/ml/helper/ConnectorAccessControlHelper.java b/plugin/src/main/java/org/opensearch/ml/helper/ConnectorAccessControlHelper.java
@@ -112,6 +112,7 @@ public void validateConnectorAccess(
             GetDataObjectRequest getDataObjectRequest = GetDataObjectRequest
                 .builder()
                 .index(ML_CONNECTOR_INDEX)
+                .tenantId(tenantId)
                 .id(connectorId)
                 .fetchSourceContext(fetchSourceContext)
                 .build();