run auto deploy remote model in partially deployed status (opensearch-project#3423) (opensearch-project#3959)

opensearch-trigger-bot[bot] · Zhangxunmt · web-flow · commit f1dab23022dd · 2025-07-04T21:45:41.000-07:00
* run auto deploy remote model in partially deployed status * add sync up for planning worker nodes * add more UTs and java doc * rename syncPlanningWorkerNodes from comments --------- (cherry picked from commit 8fff3f3) Signed-off-by: Xun Zhang <xunzh@amazon.com> Co-authored-by: Xun Zhang <xunzh@amazon.com>
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelCache.java b/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 syncPlanningWorkerNodes(Set<String> planningWorkerNodes) {
+        this.targetWorkerNodes.clear();
+        this.targetWorkerNodes.addAll(planningWorkerNodes);
+    }
+
     public boolean isDeployToAllNodes() {
         return this.deployToAllNodes != null && this.deployToAllNodes;
     }
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java
@@ -537,6 +537,20 @@ public String[] getWorkerNodes(String modelId) {
         return modelCache.getWorkerNodes();
     }
 
+    /**
+     * Get target worker nodes of model.
+     *
+     * @param modelId model id
+     * @return array of node id; return null if model not exists in cache
+     */
+    public String[] getTargetWorkerNodes(String modelId) {
+        MLModelCache modelCache = modelCaches.get(modelId);
+        if (modelCache == null) {
+            return null;
+        }
+        return modelCache.getTargetWorkerNodes();
+    }
+
     /**
      * Add worker node of model.
      * 
@@ -609,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.syncPlanningWorkerNodes(entry.getValue());
+        });
+    }
+
     /**
      * Clear worker nodes for all models.
      */
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -2436,6 +2436,16 @@ public int getWorkerNodesSize(String modelId, FunctionName functionName) {
         return getWorkerNodes(modelId, functionName, false).length;
     }
 
+    /**
+     * Get target/planning worker node of a specific model.
+     *
+     * @param modelId      model id
+     * @return list of planning worker node ids
+     */
+    public String[] getTargetWorkerNodes(String modelId) {
+        return modelCacheHelper.getTargetWorkerNodes(modelId);
+    }
+
     /**
      * Get predictable instance with model id.
      *
@@ -2476,6 +2486,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);
     }
 
     /**
diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java b/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java
@@ -180,7 +180,9 @@ public void dispatchTask(
                 }
             }, listener::onFailure);
             String[] workerNodes = mlModelManager.getWorkerNodes(modelId, functionName, true);
-            if (workerNodes == null || workerNodes.length == 0) {
+            String[] targetWorkerNodes = mlModelManager.getTargetWorkerNodes(modelId);
+
+            if (requiresAutoDeployment(workerNodes, targetWorkerNodes)) {
                 if (FunctionName.isAutoDeployEnabled(autoDeploymentEnabled, functionName)) {
                     try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                         mlModelManager.getModel(modelId, request.getTenantId(), ActionListener.runBefore(ActionListener.wrap(model -> {
@@ -651,4 +653,10 @@ 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);
+    }
 }
diff --git a/plugin/src/test/java/org/opensearch/ml/model/MLModelCacheHelperTests.java b/plugin/src/test/java/org/opensearch/ml/model/MLModelCacheHelperTests.java
@@ -298,6 +298,27 @@ public void testSyncWorkerNodes_NullModelState() {
         assertArrayEquals(new String[] { newNodeId }, cacheHelper.getWorkerNodes(modelId));
     }
 
+    public void testGetTargetWorkerNodes() {
+        String[] workerNodes = cacheHelper.getTargetWorkerNodes(modelId);
+        assertNull(workerNodes);
+        String newNodeId = "new_node_id";
+        Map<String, Set<String>> modelPlannningWorkerNodes = new HashMap<>();
+        modelPlannningWorkerNodes.put(modelId, ImmutableSet.of(newNodeId));
+        cacheHelper.syncPlanningWorkerNodes(modelPlannningWorkerNodes);
+        workerNodes = cacheHelper.getTargetWorkerNodes(modelId);
+        assertArrayEquals(new String[] { "new_node_id" }, workerNodes);
+
+    }
+
+    public void testSyncPlanningWorkerNodes() {
+        String newNodeId = "new_node_id";
+        Map<String, Set<String>> modelPlannningWorkerNodes = new HashMap<>();
+        modelPlannningWorkerNodes.put(modelId, ImmutableSet.of(newNodeId));
+        cacheHelper.syncPlanningWorkerNodes(modelPlannningWorkerNodes);
+        assertArrayEquals(new String[] { modelId }, cacheHelper.getAllModels());
+        assertArrayEquals(new String[] { newNodeId }, cacheHelper.getTargetWorkerNodes(modelId));
+    }
+
     public void testSyncWorkerNodes_ModelState() {
         String modelId2 = "model_id2";
         cacheHelper.initModelState(modelId2, MLModelState.DEPLOYED, FunctionName.TEXT_EMBEDDING, targetWorkerNodes, true);
diff --git a/plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java b/plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java
@@ -21,6 +21,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.cluster.node.DiscoveryNodeRole.CLUSTER_MANAGER_ROLE;
 import static org.opensearch.ml.common.MLTask.FUNCTION_NAME_FIELD;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_BATCH_INGESTION_BULK_SIZE;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MAX_BATCH_INFERENCE_TASKS;
@@ -50,6 +51,7 @@
 import static org.opensearch.ml.utils.TestHelper.copyFile;
 
 import java.io.IOException;
+import java.net.InetAddress;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
@@ -65,6 +67,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import org.junit.Before;
 import org.junit.Ignore;
@@ -75,11 +78,13 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.opensearch.OpenSearchStatusException;
+import org.opensearch.Version;
 import org.opensearch.action.get.GetRequest;
 import org.opensearch.action.get.GetResponse;
 import org.opensearch.action.index.IndexResponse;
 import org.opensearch.action.update.UpdateRequest;
 import org.opensearch.action.update.UpdateResponse;
+import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.cluster.service.ClusterApplierService;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.ClusterSettings;
@@ -90,6 +95,7 @@
 import org.opensearch.core.common.breaker.CircuitBreaker;
 import org.opensearch.core.common.breaker.CircuitBreakingException;
 import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.common.transport.TransportAddress;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
@@ -986,6 +992,28 @@ public void testSyncModelWorkerNodes() {
         verify(modelCacheHelper).syncWorkerNodes(eq(modelWorkerNodes));
     }
 
+    public void testSyncModelPlanningWorkerNodes() {
+        DiscoveryNode localNode = new DiscoveryNode(
+            "foo1",
+            "node1",
+            new TransportAddress(InetAddress.getLoopbackAddress(), 9300),
+            Collections.emptyMap(),
+            Collections.singleton(CLUSTER_MANAGER_ROLE),
+            Version.CURRENT
+        );
+
+        Map<String, Set<String>> modelWorkerNodes = ImmutableMap.of(modelId, ImmutableSet.of("node1"));
+        when(modelCacheHelper.getFunctionName(modelId)).thenReturn(FunctionName.TEXT_EMBEDDING);
+        when(modelCacheHelper.getDeployToAllNodes(modelId)).thenReturn(true);
+        DiscoveryNode[] planningWorkerNodes = new DiscoveryNode[] { localNode };
+        when(nodeHelper.getEligibleNodes(FunctionName.TEXT_EMBEDDING)).thenReturn(planningWorkerNodes);
+        modelManager.syncModelPlanningWorkerNodes(modelWorkerNodes);
+        verify(modelCacheHelper)
+            .syncPlanningWorkerNodes(
+                Map.of(modelId, Arrays.stream(planningWorkerNodes).map(DiscoveryNode::getId).collect(Collectors.toSet()))
+            );
+    }
+
     public void testClearRoutingTable() {
         modelManager.clearRoutingTable();
         verify(modelCacheHelper).clearWorkerNodes();

Original file line number	Diff line number	Diff line change
`@@ -146,6 +146,11 @@ public void syncWorkerNode(Set<String> workerNodes) {`
`146`	`146`	`this.workerNodes.addAll(workerNodes);`
`147`	`147`	`}`
`148`	`148`
	`149`	`+ public void syncPlanningWorkerNodes(Set<String> planningWorkerNodes) {`
	`150`	`+ this.targetWorkerNodes.clear();`
	`151`	`+ this.targetWorkerNodes.addAll(planningWorkerNodes);`
	`152`	`+ }`
	`153`	`+`
`149`	`154`	`public boolean isDeployToAllNodes() {`
`150`	`155`	`return this.deployToAllNodes != null && this.deployToAllNodes;`
`151`	`156`	`}`
Original file line number	Diff line number	Diff line change
`@@ -180,7 +180,9 @@ public void dispatchTask(`
`180`	`180`	`}`
`181`	`181`	`}, listener::onFailure);`
`182`	`182`	`String[] workerNodes = mlModelManager.getWorkerNodes(modelId, functionName, true);`
`183`		`- if (workerNodes == null \|\| workerNodes.length == 0) {`
	`183`	`+ String[] targetWorkerNodes = mlModelManager.getTargetWorkerNodes(modelId);`
	`184`	`+`
	`185`	`+ if (requiresAutoDeployment(workerNodes, targetWorkerNodes)) {`
`184`	`186`	`if (FunctionName.isAutoDeployEnabled(autoDeploymentEnabled, functionName)) {`
`185`	`187`	`try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {`
`186`	`188`	`mlModelManager.getModel(modelId, request.getTenantId(), ActionListener.runBefore(ActionListener.wrap(model -> {`
`@@ -651,4 +653,10 @@ public void validateOutputSchema(String modelId, ModelTensorOutput output) {`
`651`	`653`	`}`
`652`	`654`	`}`
`653`	`655`	`}`
	`656`	`+`
	`657`	`+ private boolean requiresAutoDeployment(String[] workerNodes, String[] targetWorkerNodes) {`
	`658`	`+ return workerNodes == null`
	`659`	`+ \|\| workerNodes.length == 0`
	`660`	`+ \|\| (targetWorkerNodes != null && workerNodes.length < targetWorkerNodes.length);`
	`661`	`+ }`
`654`	`662`	`}`