agent out of the setting only_run_on_ml_node (#2299)

* agent out of the setting only_run_on_ml_node Signed-off-by: Jing Zhang <[email protected]> * address comments Signed-off-by: Jing Zhang <[email protected]> --------- Signed-off-by: Jing Zhang <[email protected]> (cherry picked from commit e3e72f7)
opensearch-project · Apr 8, 2024 · 602c774 · 602c774
1 parent 1809b20
commit 602c774
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/cluster/DiscoveryNodeHelper.java b/plugin/src/main/java/org/opensearch/ml/cluster/DiscoveryNodeHelper.java
@@ -73,7 +73,7 @@ public DiscoveryNode[] getEligibleNodes(FunctionName functionName) {
             if (excludedNodeNames != null && excludedNodeNames.contains(node.getName())) {
                 continue;
             }
-            if (functionName == FunctionName.REMOTE) {// remote model
+            if (functionName == FunctionName.REMOTE || functionName == FunctionName.AGENT) {// non-local model (remote model, agent)
                 getEligibleNode(remoteModelEligibleNodeRoles, eligibleNodes, node);
             } else { // local model
                 if (onlyRunOnMLNode) {

diff --git a/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java b/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java
@@ -37,7 +37,7 @@ private MLCommonsSettings() {}
     public static final Setting<Integer> ML_COMMONS_MAX_ML_TASK_PER_NODE = Setting
         .intSetting("plugins.ml_commons.max_ml_task_per_node", 10, 0, 10000, Setting.Property.NodeScope, Setting.Property.Dynamic);
     public static final Setting<Boolean> ML_COMMONS_ONLY_RUN_ON_ML_NODE = Setting
-        .boolSetting("plugins.ml_commons.only_run_on_ml_node", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+        .boolSetting("plugins.ml_commons.only_run_on_ml_node", true, Setting.Property.NodeScope, Setting.Property.Dynamic);
 
     public static final Setting<Boolean> ML_COMMONS_ENABLE_INHOUSE_PYTHON_MODEL = Setting
         .boolSetting("plugins.ml_commons.enable_inhouse_python_model", false, Setting.Property.NodeScope, Setting.Property.Dynamic);

diff --git a/plugin/src/test/java/org/opensearch/ml/cluster/DiscoveryNodeHelperTests.java b/plugin/src/test/java/org/opensearch/ml/cluster/DiscoveryNodeHelperTests.java
@@ -182,6 +182,19 @@ public void testGetEligibleNodes_MLNode_RemoteModel() {
         assertFalse(nodeIds.contains(warmDataNode1.getId()));
     }
 
+    public void testGetEligibleNodes_MLNode_Agent() {
+        DiscoveryNode[] eligibleNodes = discoveryNodeHelper.getEligibleNodes(FunctionName.AGENT);
+        assertEquals(5, eligibleNodes.length);
+        Set<String> nodeIds = new HashSet<>();
+        nodeIds.addAll(Arrays.asList(eligibleNodes).stream().map(n -> n.getId()).collect(Collectors.toList()));
+        assertTrue(nodeIds.contains(mlNode1.getId()));
+        assertTrue(nodeIds.contains(mlNode2.getId()));
+        assertTrue(nodeIds.contains(dataNode1.getId()));
+        assertTrue(nodeIds.contains(dataNode2.getId()));
+        assertTrue(nodeIds.contains(allRoleNode.getId()));
+        assertFalse(nodeIds.contains(warmDataNode1.getId()));
+    }
+
     public void testGetEligibleNodes_MLNode_LocalModel() {
         DiscoveryNode[] eligibleNodes = discoveryNodeHelper.getEligibleNodes(FunctionName.TEXT_EMBEDDING);
         assertEquals(3, eligibleNodes.length);