From c2ab5d2a0398bfd7bc7d29131b273b00cfd5400c Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Mon, 16 Sep 2024 11:46:09 -0700 Subject: [PATCH 1/2] Add more missing tenant ids in requests Signed-off-by: Daniel Widdis --- .../agents/TransportRegisterAgentAction.java | 2 +- .../models/DeleteModelTransportAction.java | 30 +++++++++++++++---- .../models/UpdateModelTransportAction.java | 8 ++++- .../ml/helper/ModelAccessControlHelper.java | 7 ++++- 4 files changed, 39 insertions(+), 8 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java b/plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java index 98a15b8c3a..6102a57a60 100644 --- a/plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java +++ b/plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java @@ -86,7 +86,7 @@ private void registerAgent(MLAgent agent, ActionListener { diff --git a/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java index d5907f36fc..6a8458109c 100644 --- a/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java +++ b/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java @@ -124,6 +124,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener actionListener) { - DeleteDataObjectRequest deleteDataObjectRequest = DeleteDataObjectRequest.builder().index(ML_MODEL_INDEX).id(modelId).build(); + private void deleteModel( + String modelId, + String tenantId, + String functionName, + Boolean isHidden, + ActionListener actionListener + ) { + DeleteDataObjectRequest deleteDataObjectRequest = DeleteDataObjectRequest + .builder() + .index(ML_MODEL_INDEX) + .id(modelId) + .tenantId(tenantId) + .build(); sdkClient .deleteDataObjectAsync(deleteDataObjectRequest, client.threadPool().executor(GENERAL_THREAD_POOL)) .whenComplete((r, throwable) -> { @@ -355,7 +373,9 @@ private void deleteModelChunksAndController( // for remote model we don't need to delete model chunks so reducing one latch countdown. countDownLatch.countDown(); } - deleteController(modelId, isHidden, countDownActionListener); + // TODO This uses remote client but we haven't implemented it anywhere else?! + // deleteController(modelId, isHidden, countDownActionListener); + countDownActionListener.onResponse(true); } /** 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 index 096d392e2a..3228d44638 100644 --- a/plugin/src/main/java/org/opensearch/ml/action/models/UpdateModelTransportAction.java +++ b/plugin/src/main/java/org/opensearch/ml/action/models/UpdateModelTransportAction.java @@ -264,6 +264,7 @@ private void updateRemoteOrTextEmbeddingModel( updateModelWithRegisteringToAnotherModelGroup( modelId, newModelGroupId, + tenantId, user, updateModelInput, wrappedListener, @@ -291,6 +292,7 @@ private void updateRemoteOrTextEmbeddingModel( updateModelWithRegisteringToAnotherModelGroup( modelId, newModelGroupId, + tenantId, user, updateModelInput, wrappedListener, @@ -336,6 +338,7 @@ private void updateModelWithNewStandAloneConnector( updateModelWithRegisteringToAnotherModelGroup( modelId, newModelGroupId, + tenantId, user, updateModelInput, wrappedListener, @@ -369,6 +372,7 @@ private void updateModelWithNewStandAloneConnector( private void updateModelWithRegisteringToAnotherModelGroup( String modelId, String newModelGroupId, + String tenantId, User user, MLUpdateModelInput updateModelInput, ActionListener wrappedListener, @@ -414,12 +418,13 @@ private void updateModelWithRegisteringToAnotherModelGroup( wrappedListener.onFailure(exception); })); } else { - buildUpdateRequest(modelId, updateRequest, updateModelInput, wrappedListener, isUpdateModelCache); + buildUpdateRequest(modelId, tenantId, updateRequest, updateModelInput, wrappedListener, isUpdateModelCache); } } private void buildUpdateRequest( String modelId, + String tenantId, UpdateRequest updateRequest, MLUpdateModelInput updateModelInput, ActionListener wrappedListener, @@ -430,6 +435,7 @@ private void buildUpdateRequest( .builder() .index(updateRequest.index()) .id(updateRequest.id()) + .tenantId(tenantId) .dataObject(updateModelInput) .build(); // TODO: This should probably be default on update data object: diff --git a/plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java b/plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java index d4172b024a..962b20f023 100644 --- a/plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java +++ b/plugin/src/main/java/org/opensearch/ml/helper/ModelAccessControlHelper.java @@ -140,7 +140,12 @@ public void validateModelGroupAccess( listener.onResponse(true); return; } - GetDataObjectRequest getModelGroupRequest = GetDataObjectRequest.builder().index(ML_MODEL_GROUP_INDEX).id(modelGroupId).build(); + GetDataObjectRequest getModelGroupRequest = GetDataObjectRequest + .builder() + .index(ML_MODEL_GROUP_INDEX) + .id(modelGroupId) + .tenantId(tenantId) + .build(); try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { ActionListener wrappedListener = ActionListener.runBefore(listener, context::restore); sdkClient From a4de5579a7cc218b08bccd9208d29134979cfd81 Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Mon, 16 Sep 2024 13:32:59 -0700 Subject: [PATCH 2/2] Remove controller workaround as it causes other tests to fail Signed-off-by: Daniel Widdis --- .../ml/action/models/DeleteModelTransportAction.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java index 6a8458109c..2cded3d8c1 100644 --- a/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java +++ b/plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java @@ -373,9 +373,7 @@ private void deleteModelChunksAndController( // for remote model we don't need to delete model chunks so reducing one latch countdown. countDownLatch.countDown(); } - // TODO This uses remote client but we haven't implemented it anywhere else?! - // deleteController(modelId, isHidden, countDownActionListener); - countDownActionListener.onResponse(true); + deleteController(modelId, isHidden, countDownActionListener); } /**