do not count remote model in checking the max models per node

Signed-off-by: Xun Zhang <[email protected]>
opensearch-project · Jul 5, 2023 · 08a5834 · 08a5834
1 parent 236052d
commit 08a5834
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java
@@ -151,6 +151,20 @@ public String[] getDeployedModels() {
             .toArray(new String[0]);
     }
 
+    /**
+     * Get deployed local models on node.
+     * @return array of model id
+     */
+    public String[] getLocalDeployedModels() {
+        return modelCaches
+                .entrySet()
+                .stream()
+                .filter(entry -> (entry.getValue().getModelState() == MLModelState.DEPLOYED && entry.getValue().getFunctionName() != FunctionName.REMOTE))
+                .map(entry -> entry.getKey())
+                .collect(Collectors.toList())
+                .toArray(new String[0]);
+    }
+
     /**
      * Check if model is running on node.
      * @param modelId model id

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -709,8 +709,8 @@ public void deployModel(
             listener.onResponse("successful");
             return;
         }
-        if (modelCacheHelper.getDeployedModels().length >= maxModelPerNode) {
-            listener.onFailure(new IllegalArgumentException("Exceed max model per node limit"));
+        if (modelCacheHelper.getLocalDeployedModels().length >= maxModelPerNode) {
+            listener.onFailure(new IllegalArgumentException("Exceed max local model per node limit"));
             return;
         }
         modelCacheHelper.initModelState(modelId, MLModelState.DEPLOYING, functionName, workerNodes, deployToAllNodes);