Skip to content

Commit

Permalink
add exclude nodes setting (#813)
Browse files Browse the repository at this point in the history 
* add exclude nodes setting

Signed-off-by: Yaliang Wu <ylwu@amazon.com>

* fix duplicate nodes if node has both ml and data roles

Signed-off-by: Yaliang Wu <ylwu@amazon.com>

---------

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
  • Loading branch information
@ylwu-amzn
ylwu-amzn authored Mar 21, 2023
1 parent 894f983 commit aac0926
Show file tree
Hide file tree
Showing 10 changed files with 281 additions and 22 deletions.
6 changes: 4 additions & 2 deletions plugin/src/main/java/org/opensearch/ml/action/unload/TransportUnloadModelAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -198,8 +198,10 @@ private UnloadModelNodeResponse createUnloadModelNodeResponse(UnloadModelNodesRe
String[] modelIds = unloadModelNodesRequest.getModelIds();

Map<String, Integer> modelWorkerNodeCounts = new HashMap<>();
if (modelIds != null) {
for (String modelId : modelIds) {
boolean specifiedModelIds = modelIds != null && modelIds.length > 0;
String[] removedModelIds = specifiedModelIds ? modelIds : mlModelManager.getAllModelIds();
if (removedModelIds != null) {
for (String modelId : removedModelIds) {
String[] workerNodes = mlModelManager.getWorkerNodes(modelId);
modelWorkerNodeCounts.put(modelId, workerNodes == null ? 0 : workerNodes.length);
}
Expand Down
30 changes: 30 additions & 0 deletions plugin/src/main/java/org/opensearch/ml/cluster/DiscoveryNodeHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

package org.opensearch.ml.cluster;

import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_EXCLUDE_NODE_NAMES;
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ONLY_RUN_ON_ML_NODE;

import java.util.ArrayList;
Expand All @@ -19,6 +20,7 @@
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.node.DiscoveryNode;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.Strings;
import org.opensearch.common.settings.Settings;
import org.opensearch.ml.common.CommonValue;
import org.opensearch.ml.utils.MLNodeUtils;
Expand All @@ -28,12 +30,17 @@ public class DiscoveryNodeHelper {
private final ClusterService clusterService;
private final HotDataNodePredicate eligibleNodeFilter;
private volatile Boolean onlyRunOnMLNode;
private volatile Set<String> excludedNodeNames;

public DiscoveryNodeHelper(ClusterService clusterService, Settings settings) {
this.clusterService = clusterService;
eligibleNodeFilter = new HotDataNodePredicate();
onlyRunOnMLNode = ML_COMMONS_ONLY_RUN_ON_ML_NODE.get(settings);
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_ONLY_RUN_ON_ML_NODE, it -> onlyRunOnMLNode = it);
excludedNodeNames = Strings.commaDelimitedListToSet(ML_COMMONS_EXCLUDE_NODE_NAMES.get(settings));
clusterService
.getClusterSettings()
.addSettingsUpdateConsumer(ML_COMMONS_EXCLUDE_NODE_NAMES, it -> excludedNodeNames = Strings.commaDelimitedListToSet(it));
}

public String[] getEligibleNodeIds() {
Expand All @@ -50,6 +57,9 @@ public DiscoveryNode[] getEligibleNodes() {
final List<DiscoveryNode> eligibleMLNodes = new ArrayList<>();
final List<DiscoveryNode> eligibleDataNodes = new ArrayList<>();
for (DiscoveryNode node : state.nodes()) {
if (excludedNodeNames != null && excludedNodeNames.contains(node.getName())) {
continue;
}
if (MLNodeUtils.isMLNode(node)) {
eligibleMLNodes.add(node);
}
Expand All @@ -68,6 +78,26 @@ public DiscoveryNode[] getEligibleNodes() {
}
}

public String[] filterEligibleNodes(String[] nodeIds) {
if (nodeIds == null || nodeIds.length == 0) {
return nodeIds;
}
DiscoveryNode[] nodes = getNodes(nodeIds);
final Set<String> eligibleNodes = new HashSet<>();
for (DiscoveryNode node : nodes) {
if (excludedNodeNames != null && excludedNodeNames.contains(node.getName())) {
continue;
}
if (MLNodeUtils.isMLNode(node)) {
eligibleNodes.add(node.getId());
}
if (!onlyRunOnMLNode && node.isDataNode() && isEligibleDataNode(node)) {
eligibleNodes.add(node.getId());
}
}
return eligibleNodes.toArray(new String[0]);
}

public DiscoveryNode[] getAllNodes() {
ClusterState state = this.clusterService.state();
final List<DiscoveryNode> nodes = new ArrayList<>();
Expand Down
33 changes: 30 additions & 3 deletions plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@
import org.opensearch.index.reindex.DeleteByQueryAction;
import org.opensearch.index.reindex.DeleteByQueryRequest;
import org.opensearch.ml.breaker.MLCircuitBreakerService;
import org.opensearch.ml.cluster.DiscoveryNodeHelper;
import org.opensearch.ml.common.FunctionName;
import org.opensearch.ml.common.MLModel;
import org.opensearch.ml.common.MLTask;
Expand Down Expand Up @@ -126,6 +127,7 @@ public class MLModelManager {
private final MLIndicesHandler mlIndicesHandler;
private final MLTaskManager mlTaskManager;
private final MLEngine mlEngine;
private final DiscoveryNodeHelper nodeHelper;

private volatile Integer maxModelPerNode;
private volatile Integer maxUploadTasksPerNode;
Expand Down Expand Up @@ -153,7 +155,8 @@ public MLModelManager(
MLIndicesHandler mlIndicesHandler,
MLTaskManager mlTaskManager,
MLModelCacheHelper modelCacheHelper,
MLEngine mlEngine
MLEngine mlEngine,
DiscoveryNodeHelper nodeHelper
) {
this.client = client;
this.threadPool = threadPool;
Expand All @@ -166,6 +169,7 @@ public MLModelManager(
this.mlIndicesHandler = mlIndicesHandler;
this.mlTaskManager = mlTaskManager;
this.mlEngine = mlEngine;
this.nodeHelper = nodeHelper;

this.maxModelPerNode = ML_COMMONS_MAX_MODELS_PER_NODE.get(settings);
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MAX_MODELS_PER_NODE, it -> maxModelPerNode = it);
Expand Down Expand Up @@ -726,13 +730,36 @@ private void removeModel(String modelId) {
}

/**
* Get worker nodes of specif model.
* Get worker nodes of specific model.
*
* @param modelId model id
* @param onlyEligibleNode return only eligible node
* @return list of worker node ids
*/
public String[] getWorkerNodes(String modelId, boolean onlyEligibleNode) {
String[] workerNodeIds = modelCacheHelper.getWorkerNodes(modelId);
if (!onlyEligibleNode) {
return workerNodeIds;
}
if (workerNodeIds == null || workerNodeIds.length == 0) {
return workerNodeIds;
}

String[] eligibleNodeIds = nodeHelper.filterEligibleNodes(workerNodeIds);
if (eligibleNodeIds == null || eligibleNodeIds.length == 0) {
throw new IllegalArgumentException("No eligible worker node found");
}
return eligibleNodeIds;
}

/**
* Get worker node of specific model without filtering eligible node.
*
* @param modelId model id
* @return list of worker node ids
*/
public String[] getWorkerNodes(String modelId) {
return modelCacheHelper.getWorkerNodes(modelId);
return getWorkerNodes(modelId, false);
}

/**
Expand Down
6 changes: 4 additions & 2 deletions plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,8 @@ public Collection<Object> createComponents(
mlIndicesHandler,
mlTaskManager,
modelCacheHelper,
mlEngine
mlEngine,
nodeHelper
);
mlInputDatasetHandler = new MLInputDatasetHandler(client);

Expand Down Expand Up @@ -513,7 +514,8 @@ public List<Setting<?>> getSettings() {
MLCommonsSettings.ML_COMMONS_MAX_ML_TASK_PER_NODE,
MLCommonsSettings.ML_COMMONS_MAX_LOAD_MODEL_TASKS_PER_NODE,
MLCommonsSettings.ML_COMMONS_TRUSTED_URL_REGEX,
MLCommonsSettings.ML_COMMONS_NATIVE_MEM_THRESHOLD
MLCommonsSettings.ML_COMMONS_NATIVE_MEM_THRESHOLD,
MLCommonsSettings.ML_COMMONS_EXCLUDE_NODE_NAMES
);
return settings;
}
Expand Down
3 changes: 3 additions & 0 deletions plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,4 +56,7 @@ private MLCommonsSettings() {}

public static final Setting<Integer> ML_COMMONS_NATIVE_MEM_THRESHOLD = Setting
.intSetting("plugins.ml_commons.native_memory_threshold", 90, 0, 100, Setting.Property.NodeScope, Setting.Property.Dynamic);

public static final Setting<String> ML_COMMONS_EXCLUDE_NODE_NAMES = Setting
.simpleString("plugins.ml_commons.exclude_nodes._name", Setting.Property.NodeScope, Setting.Property.Dynamic);
}
2 changes: 1 addition & 1 deletion plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ public void dispatchTask(MLPredictionTaskRequest request, TransportService trans
transportService.sendRequest(node, getTransportActionName(), request, getResponseHandler(listener));
}
}, e -> { listener.onFailure(e); });
String[] workerNodes = mlModelManager.getWorkerNodes(modelId);
String[] workerNodes = mlModelManager.getWorkerNodes(modelId, true);
if (workerNodes == null || workerNodes.length == 0) {
if (algorithm == FunctionName.TEXT_EMBEDDING) {
listener.onFailure(new IllegalArgumentException("model not loaded"));
Expand Down
2 changes: 1 addition & 1 deletion plugin/src/main/java/org/opensearch/ml/task/MLTaskDispatcher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ public void dispatch(ActionListener<DiscoveryNode> actionListener) {

public void dispatchPredictTask(String[] nodeIds, ActionListener<DiscoveryNode> actionListener) {
if (nodeIds == null || nodeIds.length == 0) {
throw new IllegalArgumentException("Model not loaded yet");
throw new IllegalArgumentException("no eligible node to run predict request");
}
if (ROUND_ROBIN.equals(dispatchPolicy)) {
dispatchTaskWithRoundRobin(
Expand Down
104 changes: 92 additions & 12 deletions plugin/src/test/java/org/opensearch/ml/cluster/DiscoveryNodeHelperTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,16 @@
import static java.util.Collections.emptySet;
import static org.mockito.Mockito.when;
import static org.opensearch.cluster.node.DiscoveryNodeRole.BUILT_IN_ROLES;
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_EXCLUDE_NODE_NAMES;
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_ONLY_RUN_ON_ML_NODE;
import static org.opensearch.ml.utils.TestHelper.ML_ROLE;
import static org.opensearch.ml.utils.TestHelper.clusterSetting;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;

import org.junit.Before;
import org.mockito.Mock;
Expand All @@ -36,10 +40,14 @@
public class DiscoveryNodeHelperTests extends OpenSearchTestCase {
private final String clusterManagerNodeId = "clusterManagerNode";
private final String dataNode1Id = "dataNode1";
private final String dataNode1Name = "dataNodeName1";
private final String dataNode2Id = "dataNode2";
private final String dataNode2Name = "dataNodeName2";
private final String warmDataNode1Id = "warmDataNode1";
private final String mlNode1Id = "mlNode1";
private final String mlNode1Name = "mlNodeName1";
private final String mlNode2Id = "mlNode2";
private final String mlNode2Name = "mlNodeName2";
private final String clusterName = "multi-node-cluster";

@Mock
Expand All @@ -54,11 +62,13 @@ public class DiscoveryNodeHelperTests extends OpenSearchTestCase {
private DiscoveryNode mlNode1;
private DiscoveryNode mlNode2;
private ClusterState clusterState;
private String nonExistingNodeName;

@Before
public void setup() throws IOException {
MockitoAnnotations.openMocks(this);
mockSettings(true);
nonExistingNodeName = randomAlphaOfLength(5);
mockSettings(true, nonExistingNodeName);

clusterManagerNode = new DiscoveryNode(
clusterManagerNodeId,
Expand All @@ -68,8 +78,16 @@ public void setup() throws IOException {
Version.CURRENT
);

dataNode1 = new DiscoveryNode(dataNode1Id, buildNewFakeTransportAddress(), emptyMap(), BUILT_IN_ROLES, Version.CURRENT);
dataNode1 = new DiscoveryNode(
dataNode1Name,
dataNode1Id,
buildNewFakeTransportAddress(),
emptyMap(),
BUILT_IN_ROLES,
Version.CURRENT
);
dataNode2 = new DiscoveryNode(
dataNode1Name,
dataNode2Id,
buildNewFakeTransportAddress(),
ImmutableMap.of(CommonValue.BOX_TYPE_KEY, CommonValue.HOT_BOX_TYPE),
Expand All @@ -84,8 +102,22 @@ public void setup() throws IOException {
Version.CURRENT
);

mlNode1 = new DiscoveryNode(mlNode1Id, buildNewFakeTransportAddress(), emptyMap(), ImmutableSet.of(ML_ROLE), Version.CURRENT);
mlNode2 = new DiscoveryNode(mlNode2Id, buildNewFakeTransportAddress(), emptyMap(), ImmutableSet.of(ML_ROLE), Version.CURRENT);
mlNode1 = new DiscoveryNode(
mlNode1Name,
mlNode1Id,
buildNewFakeTransportAddress(),
emptyMap(),
ImmutableSet.of(ML_ROLE),
Version.CURRENT
);
mlNode2 = new DiscoveryNode(
mlNode2Name,
mlNode2Id,
buildNewFakeTransportAddress(),
emptyMap(),
ImmutableSet.of(ML_ROLE),
Version.CURRENT
);

DiscoveryNodes nodes = DiscoveryNodes
.builder()
Expand All @@ -102,30 +134,78 @@ public void setup() throws IOException {
discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
}

private void mockSettings(boolean onlyRunOnMLNode) {
settings = Settings.builder().put(ML_COMMONS_ONLY_RUN_ON_ML_NODE.getKey(), onlyRunOnMLNode).build();
ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_ONLY_RUN_ON_ML_NODE);
private void mockSettings(boolean onlyRunOnMLNode, String excludedNodeName) {
settings = Settings
.builder()
.put(ML_COMMONS_ONLY_RUN_ON_ML_NODE.getKey(), onlyRunOnMLNode)
.put(ML_COMMONS_EXCLUDE_NODE_NAMES.getKey(), excludedNodeName)
.build();
ClusterSettings clusterSettings = clusterSetting(settings, ML_COMMONS_ONLY_RUN_ON_ML_NODE, ML_COMMONS_EXCLUDE_NODE_NAMES);
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
}

public void testGetEligibleNodes_MLNode() {
DiscoveryNode[] eligibleNodes = discoveryNodeHelper.getEligibleNodes();
assertEquals(2, eligibleNodes.length);
assertEquals(mlNode1.getName(), eligibleNodes[0].getName());
assertEquals(mlNode2.getName(), eligibleNodes[1].getName());
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()));
}

public void testGetEligibleNodes_DataNode() {
mockSettings(false);
mockSettings(false, nonExistingNodeName);
DiscoveryNodeHelper discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
DiscoveryNodes nodes = DiscoveryNodes.builder().add(clusterManagerNode).add(dataNode1).add(dataNode2).add(warmDataNode1).build();
clusterState = new ClusterState(new ClusterName(clusterName), 123l, "111111", null, null, nodes, null, null, 0, false);
when(clusterService.state()).thenReturn(clusterState);

DiscoveryNode[] eligibleNodes = discoveryNodeHelper.getEligibleNodes();
assertEquals(2, eligibleNodes.length);
assertEquals(mlNode1.getName(), eligibleNodes[0].getName());
assertEquals(mlNode2.getName(), eligibleNodes[1].getName());
assertEquals(dataNode1.getName(), eligibleNodes[0].getName());
assertEquals(dataNode2.getName(), eligibleNodes[1].getName());
}

public void testGetEligibleNodes_MLNode_Excluded() {
mockSettings(false, mlNode1.getName() + "," + mlNode2.getName());
DiscoveryNodeHelper discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
DiscoveryNode[] eligibleNodes = discoveryNodeHelper.getEligibleNodes();
assertEquals(2, eligibleNodes.length);
assertEquals(dataNode1.getName(), eligibleNodes[0].getName());
assertEquals(dataNode1.getName(), eligibleNodes[1].getName());
}

public void testFilterEligibleNodes_Null() {
mockSettings(false, mlNode1.getName() + "," + mlNode2.getName());
DiscoveryNodeHelper discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
String[] eligibleNodes = discoveryNodeHelper.filterEligibleNodes(null);
assertNull(eligibleNodes);
}

public void testFilterEligibleNodes_Empty() {
mockSettings(false, mlNode1.getName() + "," + mlNode2.getName());
DiscoveryNodeHelper discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
String[] eligibleNodes = discoveryNodeHelper.filterEligibleNodes(new String[] {});
assertEquals(0, eligibleNodes.length);
}

public void testFilterEligibleNodes() {
mockSettings(true, mlNode1.getName());
DiscoveryNodeHelper discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
String[] eligibleNodes = discoveryNodeHelper.filterEligibleNodes(new String[] { mlNode1Id, mlNode2Id, dataNode1Id });
assertEquals(1, eligibleNodes.length);
assertEquals(mlNode2Id, eligibleNodes[0]);
}

public void testFilterEligibleNodes_BothMLAndDataNodes() {
mockSettings(false, mlNode1.getName());
DiscoveryNodeHelper discoveryNodeHelper = new DiscoveryNodeHelper(clusterService, settings);
String[] eligibleNodes = discoveryNodeHelper.filterEligibleNodes(new String[] { mlNode1Id, mlNode2Id, dataNode1Id });
assertEquals(2, eligibleNodes.length);
Set<String> nodeIds = new HashSet<>();
nodeIds.addAll(Arrays.asList(eligibleNodes));
assertTrue(nodeIds.contains(dataNode1Id));
assertTrue(nodeIds.contains(mlNode2Id));
}

public void testGetAllNodeIds() {
Expand Down
Loading

0 comments on commit aac0926

Please sign in to comment.