add sync up for planning worker nodes

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Author: Zhangxunmt

plugin/src/main/java/org/opensearch/ml/model/MLModelCache.java

@@ -146,6 +146,11 @@ public void syncWorkerNode(Set<String> workerNodes) {
         this.workerNodes.addAll(workerNodes);
     }
 
+    public void syncPlanningWorkerNode(Set<String> planningWorkerNodes) {
+        this.targetWorkerNodes.clear();
+        this.targetWorkerNodes.addAll(planningWorkerNodes);
+    }
+
     public boolean isDeployToAllNodes() {
         return this.deployToAllNodes != null && this.deployToAllNodes;
     }

plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java

@@ -551,7 +551,6 @@ public String[] getTargetWorkerNodes(String modelId) {
         return modelCache.getTargetWorkerNodes();
     }
 
-
     /**
      * Add worker node of model.
      * 
@@ -624,6 +623,19 @@ public void syncWorkerNodes(Map<String, Set<String>> modelWorkerNodes) {
         });
     }
 
+    /**
+     * Sync planning worker nodes for all models.
+     *
+     * @param modelPlanningWorkerNodes planning worker nodes of all models
+     */
+    public void syncPlanningWorkerNodes(Map<String, Set<String>> modelPlanningWorkerNodes) {
+        log.debug("sync model planning worker nodes");
+        modelPlanningWorkerNodes.entrySet().forEach(entry -> {
+            MLModelCache modelCache = getOrCreateModelCache(entry.getKey());
+            modelCache.syncPlanningWorkerNode(entry.getValue());
+        });
+    }
+
     /**
      * Clear worker nodes for all models.
      */

plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java

@@ -2480,6 +2480,23 @@ public String[] getExpiredModels() {
      */
     public synchronized void syncModelWorkerNodes(Map<String, Set<String>> modelWorkerNodes) {
         modelCacheHelper.syncWorkerNodes(modelWorkerNodes);
+
+        syncModelPlanningWorkerNodes(modelWorkerNodes);
+    }
+
+    public synchronized void syncModelPlanningWorkerNodes(Map<String, Set<String>> modelWorkerNodes) {
+        Map<String, Set<String>> modelPlanningWorkerNodes = new HashMap<>();
+        modelWorkerNodes.keySet().forEach(modelId -> {
+            FunctionName functionName = modelCacheHelper.getFunctionName(modelId);
+            boolean isDeployToAll = modelCacheHelper.getDeployToAllNodes(modelId);
+            if (!isDeployToAll) {
+                return;
+            }
+            DiscoveryNode[] eligibleNodes = nodeHelper.getEligibleNodes(functionName);
+            Set<String> eligibleNodeIds = Arrays.stream(eligibleNodes).map(DiscoveryNode::getId).collect(Collectors.toSet());
+            modelPlanningWorkerNodes.put(modelId, eligibleNodeIds);
+        });
+        modelCacheHelper.syncPlanningWorkerNodes(modelPlanningWorkerNodes);
     }
 
     /**

plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java

@@ -655,7 +655,8 @@ public void validateOutputSchema(String modelId, ModelTensorOutput output) {
     }
 
     private boolean requiresAutoDeployment(String[] workerNodes, String[] targetWorkerNodes) {
-        return workerNodes == null || workerNodes.length == 0 ||
-                (targetWorkerNodes != null && workerNodes.length < targetWorkerNodes.length);
+        return workerNodes == null
+            || workerNodes.length == 0
+            || (targetWorkerNodes != null && workerNodes.length < targetWorkerNodes.length);
     }
 }