enable auto redeploy for hidden model

Signed-off-by: Bhavana Ramaram <[email protected]>
rbhavna · Feb 13, 2024 · 80d3a75 · 80d3a75
1 parent 67e904e
commit 80d3a75
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 6 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/deploy/MLDeployModelRequest.java b/common/src/main/java/org/opensearch/ml/common/transport/deploy/MLDeployModelRequest.java
@@ -38,24 +38,29 @@ public class MLDeployModelRequest extends MLTaskRequest {
     private String modelId;
     private String[] modelNodeIds;
     boolean async;
+    // This is to identify if the get request is initiated by user or not. During auto redeploy, we also perform deploy operation. This field is to distinguish between
+    // these two situations.
+    private final boolean isUserInitiatedDeployRequest;
 
     @Builder
-    public MLDeployModelRequest(String modelId, String[] modelNodeIds, boolean async, boolean dispatchTask) {
+    public MLDeployModelRequest(String modelId, String[] modelNodeIds, boolean async, boolean dispatchTask, boolean isUserInitiatedDeployRequest) {
         super(dispatchTask);
         this.modelId = modelId;
         this.modelNodeIds = modelNodeIds;
         this.async = async;
+        this.isUserInitiatedDeployRequest = isUserInitiatedDeployRequest;
     }
 
     public MLDeployModelRequest(String modelId, boolean async) {
-        this(modelId, null, async, true);
+        this(modelId, null, async, true, true);
     }
 
     public MLDeployModelRequest(StreamInput in) throws IOException {
         super(in);
         this.modelId = in.readString();
         this.modelNodeIds = in.readOptionalStringArray();
         this.async = in.readBoolean();
+        this.isUserInitiatedDeployRequest = in.readOptionalBoolean();
     }
 
     @Override
@@ -74,6 +79,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeString(modelId);
         out.writeOptionalStringArray(modelNodeIds);
         out.writeBoolean(async);
+        out.writeOptionalBoolean(isUserInitiatedDeployRequest);
     }
 
     public static MLDeployModelRequest parse(XContentParser parser, String modelId) throws IOException {
@@ -96,7 +102,7 @@ public static MLDeployModelRequest parse(XContentParser parser, String modelId)
             }
         }
         String[] nodeIds = nodeIdList == null ? null : nodeIdList.toArray(new String[0]);
-        return new MLDeployModelRequest(modelId, nodeIds, false, true);
+        return new MLDeployModelRequest(modelId, nodeIds, false, true, true);
     }
 
     public static MLDeployModelRequest fromActionRequest(ActionRequest actionRequest) {

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
@@ -131,6 +131,7 @@ public TransportDeployModelAction(
     protected void doExecute(Task task, ActionRequest request, ActionListener<MLDeployModelResponse> listener) {
         MLDeployModelRequest deployModelRequest = MLDeployModelRequest.fromActionRequest(request);
         String modelId = deployModelRequest.getModelId();
+        Boolean isUserInitiatedDeployRequest = deployModelRequest.isUserInitiatedDeployRequest();
         User user = RestActionUtils.getUserContext(client);
         boolean isSuperAdmin = isSuperAdminUserWrapper(clusterService, client);
         String[] excludes = new String[] { MLModel.MODEL_CONTENT_FIELD, MLModel.OLD_MODEL_CONTENT_FIELD };
@@ -144,7 +145,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLDepl
                     throw new IllegalStateException(REMOTE_INFERENCE_DISABLED_ERR_MSG);
                 }
                 if (isHidden != null && isHidden) {
-                    if (isSuperAdmin) {
+                    if (isSuperAdmin || !isUserInitiatedDeployRequest) {
                         deployModel(deployModelRequest, mlModel, modelId, wrappedListener, listener);
                     } else {
                         wrappedListener

diff --git a/plugin/src/main/java/org/opensearch/ml/autoredeploy/MLModelAutoReDeployer.java b/plugin/src/main/java/org/opensearch/ml/autoredeploy/MLModelAutoReDeployer.java
@@ -308,7 +308,7 @@ private void triggerModelRedeploy(ModelAutoRedeployArrangement modelAutoRedeploy
                 ImmutableMap.of(MLModel.AUTO_REDEPLOY_RETRY_TIMES_FIELD, Optional.ofNullable(autoRedeployRetryTimes).orElse(0) + 1)
             );
 
-        MLDeployModelRequest deployModelRequest = new MLDeployModelRequest(modelId, nodeIds, false, true);
+        MLDeployModelRequest deployModelRequest = new MLDeployModelRequest(modelId, nodeIds, false, true, false);
         client.execute(MLDeployModelAction.INSTANCE, deployModelRequest, listener);
     }
 

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -857,7 +857,7 @@ private void updateModelRegisterStateAsDone(
     void deployModelAfterRegistering(MLRegisterModelInput registerModelInput, String modelId) {
         String[] modelNodeIds = registerModelInput.getModelNodeIds();
         log.debug("start deploying model after registering, modelId: {} on nodes: {}", modelId, Arrays.toString(modelNodeIds));
-        MLDeployModelRequest request = new MLDeployModelRequest(modelId, modelNodeIds, false, true);
+        MLDeployModelRequest request = new MLDeployModelRequest(modelId, modelNodeIds, false, true, true);
         ActionListener<MLDeployModelResponse> listener = ActionListener
             .wrap(r -> log.debug("model deployed, response {}", r), e -> log.error("Failed to deploy model", e));
         client.execute(MLDeployModelAction.INSTANCE, request, listener);

diff --git a/plugin/src/test/java/org/opensearch/ml/action/deploy/TransportDeployModelActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/deploy/TransportDeployModelActionTests.java
@@ -286,6 +286,7 @@ public void testDoExecute_no_permission_hidden_model() {
         MLModel mlModel = mock(MLModel.class);
         when(mlModel.getAlgorithm()).thenReturn(FunctionName.ANOMALY_LOCALIZATION);
         when(mlModel.getIsHidden()).thenReturn(true);
+        when(mlDeployModelRequest.isUserInitiatedDeployRequest()).thenReturn(true);
         doAnswer(invocation -> {
             ActionListener<MLModel> listener = invocation.getArgument(3);
             listener.onResponse(mlModel);