Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add exclude nodes setting #813

Merged
merged 2 commits into from
Mar 21, 2023
Merged

add exclude nodes setting #813

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
93d7d9a
add exclude nodes setting
ylwu-amzn Mar 21, 2023
e9d8054
fix duplicate nodes if node has both ml and data roles
ylwu-amzn Mar 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions plugin/src/main/java/org/opensearch/ml/action/unload/TransportUnloadModelAction.java
View file
Open in desktop
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
Open in desktop
ylwu-amzn marked this conversation as resolved.
Show resolved Hide resolved
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());
}
ylwu-amzn marked this conversation as resolved.
Show resolved Hide resolved
}
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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");
ylwu-amzn marked this conversation as resolved.
Show resolved Hide resolved
}
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
Open in desktop
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