run auto deploy remote model in partially deployed status

Signed-off-by: Xun Zhang <[email protected]>
Zhangxunmt · Jan 23, 2025 · 0671963 · 0671963
1 parent 1659a60
commit 0671963
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 1 deletion.
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,21 @@ 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.
      * 

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -2460,6 +2460,10 @@ public int getWorkerNodesSize(String modelId, FunctionName functionName) {
         return getWorkerNodes(modelId, functionName, false).length;
     }
 
+    public String[] getTargetWorkerNodes(String modelId) {
+        return modelCacheHelper.getTargetWorkerNodes(modelId);
+    }
+
     /**
      * Get predictable instance with model id.
      *

diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java b/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java
@@ -158,7 +158,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, ActionListener.runBefore(ActionListener.wrap(model -> {
@@ -568,4 +570,9 @@ 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);
+    }
 }