fix

Signed-off-by: Yaliang Wu <[email protected]>
ylwu-amzn · Jan 18, 2024 · 07fc8fc · 07fc8fc
1 parent 40c76c7
commit 07fc8fc
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 4 deletions.
diff --git a/...c/main/java/org/opensearch/ml/action/controller/CreateModelControllerTransportAction.java b/...c/main/java/org/opensearch/ml/action/controller/CreateModelControllerTransportAction.java
@@ -15,6 +15,7 @@
 import java.util.List;
 import java.util.Map;
 
+import org.apache.commons.lang3.ArrayUtils;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.DocWriteResponse;
@@ -187,7 +188,7 @@ private void indexAndCreateModelController(
                     if (indexResponse.getResult() == DocWriteResponse.Result.CREATED) {
                         mlModelManager.updateModel(modelId, Map.of(MLModel.IS_MODEL_CONTROLLER_ENABLED_FIELD, true));
                     }
-                    if (mlModelCacheHelper.isModelDeployed(modelId)) {
+                    if (!ArrayUtils.isEmpty(mlModelCacheHelper.getWorkerNodes(modelId))) {
                         log.info("Model {} is deployed. Start to deploy the model controller into cache.", modelId);
                         String[] targetNodeIds = mlModelManager.getWorkerNodes(modelId, mlModel.getAlgorithm());
                         MLDeployModelControllerNodesRequest deployModelControllerNodesRequest = new MLDeployModelControllerNodesRequest(

diff --git a/...c/main/java/org/opensearch/ml/action/controller/DeleteModelControllerTransportAction.java b/...c/main/java/org/opensearch/ml/action/controller/DeleteModelControllerTransportAction.java
@@ -13,6 +13,7 @@
 import java.util.List;
 import java.util.Map;
 
+import org.apache.commons.lang3.ArrayUtils;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.FailedNodeException;
@@ -144,7 +145,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
     // This method is used to handle the condition if we need to undeploy the model controller before deleting it from the index or not.
     private void deleteModelControllerWithDeployedModel(String modelId, ActionListener<DeleteResponse> actionListener) {
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-            if (mlModelCacheHelper.isModelDeployed(modelId)) {
+            if (!ArrayUtils.isEmpty(mlModelCacheHelper.getWorkerNodes(modelId))) {
                 log.info("Model has already been deployed in ML cache, need undeploy model controller before sending delete request.");
                 String[] targetNodeIds = getAllNodes();
                 MLUndeployModelControllerNodesRequest undeployModelControllerNodesRequest = new MLUndeployModelControllerNodesRequest(

diff --git a/...c/main/java/org/opensearch/ml/action/controller/UpdateModelControllerTransportAction.java b/...c/main/java/org/opensearch/ml/action/controller/UpdateModelControllerTransportAction.java
@@ -14,6 +14,7 @@
 import java.util.Iterator;
 import java.util.List;
 
+import org.apache.commons.lang3.ArrayUtils;
 import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.DocWriteResponse;
@@ -177,7 +178,7 @@ private void updateModelController(
                             modelId,
                             updateResponse.getResult()
                         );
-                    if (mlModelCacheHelper.isModelDeployed(modelId) && isDeployRequiredAfterUpdate) {
+                    if (!ArrayUtils.isEmpty(mlModelCacheHelper.getWorkerNodes(modelId)) && isDeployRequiredAfterUpdate) {
                         log
                             .info(
                                 "Model {} is deployed and the user rate limiter config is constructable. Start to deploy the model controller into cache.",

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java
@@ -532,7 +532,7 @@ public MLModelProfile getModelProfile(String modelId) {
         MLModelProfile.MLModelProfileBuilder builder = MLModelProfile.builder();
         builder.modelState(modelCache.getModelState());
         if (modelCache.getPredictor() != null) {
-            builder.predictor(modelCache.getPredictor().toString());
+            builder.predictor(modelCache.getPredictor().toString() + ", " + modelCache.getIsModelEnabled() + ", " + modelCache.getModelRateLimiter() + ", " + modelCache.getUserRateLimiterMap());
         }
         String[] targetWorkerNodes = modelCache.getTargetWorkerNodes();
         if (targetWorkerNodes.length > 0) {