You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "lihaosky (via GitHub)" <gi...@apache.org> on 2023/07/17 22:20:57 UTC
[GitHub] [kafka] lihaosky opened a new pull request, #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
lihaosky opened a new pull request, #14030:
URL: https://github.com/apache/kafka/pull/14030
## Description
This PR is on top of https://github.com/apache/kafka/pull/13996. This PR make use of the graph to compute min cost based on rack awareness for stateful active tasks.
## Testing
TODO
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax merged pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax merged PR #14030:
URL: https://github.com/apache/kafka/pull/14030
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1270106140
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
Review Comment:
Looking into `constructStatefulActiveTaskGraph` it does the same thing. Sounds error prone. Should we not get source and sink ID from `graph` instead of re-computing it? (Seems we do the same thing on two places)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1271152663
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
Well, if it's useful for `constructStatefulActiveTaskGraph` to have such a list, we should construct this list inside `constructStatefulActiveTaskGraph` but not pass it in? Otherwise, we leak an optimization from `constructStatefulActiveTaskGraph` to the caller what seems not ideal?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1271081266
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
Review Comment:
Forgot to change it...
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, taskIds, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignStatefulActiveTaskFromMinCostFlow(graph, taskIds, clientList, taskIdList,
+ clientStates, clientCapacity, taskClientMap);
+
+ return cost;
+ }
+
+ private void constructStatefulActiveTaskGraph(final Graph<Integer> graph,
+ final SortedSet<TaskId> statefulTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> clientCapacity,
Review Comment:
Not the same as `ClientState#capacity`. This is to track how many tasks are assigned to each client originally and we need maintain this number after optimizing cost. The reason is we need to keep assignment balanced. So we assume the assignment passed in from caller is balanced. After optimization, the same number of tasks assigned to each client doesn't change.
So when we construct the graph, from source to each task, there's 1 edge with capacity 1, flow 1 and cost 0, from each task, there's 1 edge to each client with capacity 1, cost computed and flow 0 or 1 (1 means current assignment). From each client to sink, there's one edge with capacity equal to number of original tasks assigned to the client, cost 0 and flow equal to capacity.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -38,29 +43,34 @@
public class RackAwareTaskAssignor {
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
Review Comment:
For stateless tasks, there's no `non_overlap_cost`, so traffic cost doesn't matter as long as it's positive. The reason is that it's less expensive to relocate stateless tasks.
For stateful tasks, there's `non_overlap_cost` which means there's cost if you are moving the task to other clients. Since we prefer not to move stateful tasks as much as possible, there's non-zero value for this cost. The value for `traffic_cost` vs `non_overlap_cost` means how much we favor one compared to another. Setting to 10 basically means we value `traffic_cost` much more than `non_overlap_cost`. Other default could also be possible. Also, since this class is also going to be used by `StickyAssignor`, my plan is to assign `non_overlap_cost` higher value than `traffic_cost` for sticky assignor to favor stickiness much more
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
Here I meant `processId` using the UUID.
Stateless tasks have 0 `non_overlap_cost` while stateful tasks have postive `non_overlap_cost`
`inCurrentAssignment` means the task is assigned to current client referred to by `clientId`. If not in current assignment, there's possible `non_overlap_cost`
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
Review Comment:
See above explanation of `non_overlap_cost`
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java:
##########
@@ -268,24 +268,44 @@ public static class AssignmentConfigs {
public final long probingRebalanceIntervalMs;
public final List<String> rackAwareAssignmentTags;
+ // TODO: get from streamsConfig after we add the config
Review Comment:
It's discussed here: https://cwiki.apache.org/confluence/display/KAFKA/KIP-925%3A+Rack+aware+task+assignment+in+Kafka+Streams#KIP925:RackawaretaskassignmentinKafkaStreams-5.PublicInterfaces
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
This is just for getting index easily later. Since node id in graph is integer (index), it's easier get reference back to UUID and ClientState using index. Otherwise, we need to maintain an index to UUID map I think
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
Review Comment:
Good idea!
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
Review Comment:
Same reason as above to get reference to taskId from nodeId easily
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -38,29 +43,34 @@
public class RackAwareTaskAssignor {
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
+ private static final int DEFAULT_STATEFUL_NON_OVERLAP_COST = 1;
Review Comment:
See above
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1275434335
##########
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java:
##########
@@ -69,12 +71,38 @@ public final class AssignmentTestUtils {
public static final UUID UUID_8 = uuidForInt(8);
public static final UUID UUID_9 = uuidForInt(9);
- public static final TopicPartition TP_0_0 = new TopicPartition("topic0", 0);
- public static final TopicPartition TP_0_1 = new TopicPartition("topic0", 1);
- public static final TopicPartition TP_0_2 = new TopicPartition("topic0", 2);
- public static final TopicPartition TP_1_0 = new TopicPartition("topic1", 0);
- public static final TopicPartition TP_1_1 = new TopicPartition("topic1", 1);
- public static final TopicPartition TP_1_2 = new TopicPartition("topic1", 2);
+ public static final String RACK_0 = "rock0";
Review Comment:
```suggestion
public static final String RACK_0 = "rack0";
```
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +188,213 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID processId, final boolean inCurrentAssignment, final int trafficCost, final int nonOverlapCost) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(processId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + processId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + processId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final int trafficCost, final int nonOverlapCost) {
Review Comment:
Can we add JavaDocs? It's a little unclear what this method does. Also maybe move `activeTasks` as first parameter, as they are the main input (all others are metadata)?
```
/*
Compute the cost for the provided {@code activeTasks}. The passed in active tasks must be contained in {@code clientState}`.
*/
```
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -74,11 +81,7 @@ public synchronized boolean canEnableRackAwareAssignorForActiveTasks() {
}
return validateTopicPartitionRack();
- }
-
- public boolean canEnableRackAwareAssignorForStandbyTasks() {
- // TODO
- return false;
+ // TODO: add changelog topic, standby task validation
}
// Visible for testing. This method also checks if all TopicPartitions exist in cluster
Review Comment:
Typo one line below: `D[e]scribe`
##########
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java:
##########
@@ -69,12 +71,38 @@ public final class AssignmentTestUtils {
public static final UUID UUID_8 = uuidForInt(8);
public static final UUID UUID_9 = uuidForInt(9);
- public static final TopicPartition TP_0_0 = new TopicPartition("topic0", 0);
- public static final TopicPartition TP_0_1 = new TopicPartition("topic0", 1);
- public static final TopicPartition TP_0_2 = new TopicPartition("topic0", 2);
- public static final TopicPartition TP_1_0 = new TopicPartition("topic1", 0);
- public static final TopicPartition TP_1_1 = new TopicPartition("topic1", 1);
- public static final TopicPartition TP_1_2 = new TopicPartition("topic1", 2);
+ public static final String RACK_0 = "rock0";
Review Comment:
Same below.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
Review Comment:
Could work.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1272724945
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
Review Comment:
The balance requirement I meant here is actually keeping `originalAssignedTaskNum` even though some clients having 0 task assigned seems not making it balance... I'll reword it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1270094068
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -38,29 +43,34 @@
public class RackAwareTaskAssignor {
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
+ private static final int DEFAULT_STATEFUL_NON_OVERLAP_COST = 1;
Review Comment:
Why is this not zero (as it is for stateless)?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -38,29 +43,34 @@
public class RackAwareTaskAssignor {
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
Review Comment:
Why is this set to `10`? In particular, why is it higher than for the stateless case? In the end, my understanding was that we try to optimize for input partitions, and for this case, there is no difference if a task has state or not, but only the number of input topic partitions for a task matter (each with equal cost)
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java:
##########
@@ -268,24 +268,44 @@ public static class AssignmentConfigs {
public final long probingRebalanceIntervalMs;
public final List<String> rackAwareAssignmentTags;
+ // TODO: get from streamsConfig after we add the config
Review Comment:
I cannot remember such parameters being defined in the KIP. Can you elaborate?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
Seem unnecessary to extract and pass expliclity, as we pass `clientStates` anyway.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
There is a `isStateful` flag -- do we need it? My understanding is that costs (as we only consider input topic partitions) are independent of stateful/stateless?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
Review Comment:
Looking into `constructStatefulActiveTaskGraph` it does the same thing. Sound error prone. Should we not get source and sink ID from `graph` instead of re-computing it? (Seems we do the same thing on two places)
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/Graph.java:
##########
@@ -126,7 +126,7 @@ public void setSinkNode(final V node) {
sinkNode = node;
}
- public int totalCost() {
+ public long totalCost() {
int totalCost = 0;
Review Comment:
Should this be `long`, given that we return a `long`?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
Review Comment:
Should we also verify that it is not empty?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
What does `inCurrentAssignment` exactly mean?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
I don't think that `clientId` is a `UUID`, but it should be rather a `String`? Or is this supposed to be the `processId` ?
("client" is a somewhat overloaded term unfortunately...)
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, taskIds, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignStatefulActiveTaskFromMinCostFlow(graph, taskIds, clientList, taskIdList,
+ clientStates, clientCapacity, taskClientMap);
+
+ return cost;
+ }
+
+ private void constructStatefulActiveTaskGraph(final Graph<Integer> graph,
+ final SortedSet<TaskId> statefulTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> clientCapacity,
Review Comment:
Not sure if I understand what `clientCapacity` is? Is it the same as `ClientState#capacity`? If yes, why track it twice?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
Review Comment:
Why do we pass an empty `Graph` into ` constructStatefulActiveTaskGraph`? Might be cleaner to just let the method return a `Graph` instead of make it `void` and modify an input parameter?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
Review Comment:
Not sure if I can follow?
My understanding was, that we say cost == 0 if we can assign to the same rack, and otherwise cost is 1.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
Review Comment:
`constructActiveTaskGraph` ?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
Review Comment:
Why `statefulTasks`? Seems the provided tasks would be stateless, if `isStateful == false` ?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
Review Comment:
Why do we deep copy?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, taskIds, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignStatefulActiveTaskFromMinCostFlow(graph, taskIds, clientList, taskIdList,
+ clientStates, clientCapacity, taskClientMap);
+
+ return cost;
+ }
+
+ private void constructStatefulActiveTaskGraph(final Graph<Integer> graph,
+ final SortedSet<TaskId> statefulTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> clientCapacity,
+ final boolean isStateful) {
+ for (final TaskId taskId : statefulTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ clientCapacity.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
Review Comment:
As above.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
Review Comment:
Should we make `sourceId = -1` and `sinkId = taskIdList.size() + clientList.size() + 1` ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1271152663
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
Well, if it's useful for `constructStatefulActiveTaskGraph` to have such a list, we should construct this list inside ` constructStatefulActiveTaskGraph` but not pass it in? Otherwise, we leak an optimization from ` constructStatefulActiveTaskGraph` to the caller what seems not ideal?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1271151967
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(taskIds);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, taskIds, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignStatefulActiveTaskFromMinCostFlow(graph, taskIds, clientList, taskIdList,
+ clientStates, clientCapacity, taskClientMap);
+
+ return cost;
+ }
+
+ private void constructStatefulActiveTaskGraph(final Graph<Integer> graph,
+ final SortedSet<TaskId> statefulTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> clientCapacity,
Review Comment:
> This is to track how many tasks are assigned to each client originally
Should we call it `numberOriginallyAssignedTasks` to avoid overloading the term "capacity" (and it's not really a "capacity" anyway).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1272693494
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
Sure. I will change it to `processId`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1272791951
##########
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java:
##########
@@ -145,147 +188,517 @@ public void disableActiveSinceRackMissingInClient() {
// False since process1 doesn't have rackId
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableActiveSinceRackDiffersInSameProcess() {
+ public void shouldDisableActiveWhenRackDiffersInSameProcess() {
final Map<UUID, Map<String, Optional<String>>> processRacks = new HashMap<>();
// Different consumers in same process have different rack ID. This shouldn't happen.
// If happens, there's a bug somewhere
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
processRacks,
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithoutDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithoutDescribingTopics() {
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
// partitionWithoutInfo00 has rackInfo in cluster metadata
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithDescribingTopics() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
doReturn(
Collections.singletonMap(
- TOPIC0,
+ TP_0_NAME,
Collections.singletonList(
- new TopicPartitionInfo(0, node0, Arrays.asList(replicas1), Collections.emptyList())
+ new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
)
)
- ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
+ public void shouldDisableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
- doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(
+ TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
assertTrue(assignor.populateTopicsToDiscribe(new HashSet<>()));
}
+ @Test
+ public void shouldOptimizeEmptyActiveTasks() {
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet();
Review Comment:
This is to test nothing is changed if empty set is passed in
##########
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java:
##########
@@ -145,147 +188,517 @@ public void disableActiveSinceRackMissingInClient() {
// False since process1 doesn't have rackId
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableActiveSinceRackDiffersInSameProcess() {
+ public void shouldDisableActiveWhenRackDiffersInSameProcess() {
final Map<UUID, Map<String, Optional<String>>> processRacks = new HashMap<>();
// Different consumers in same process have different rack ID. This shouldn't happen.
// If happens, there's a bug somewhere
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
processRacks,
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithoutDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithoutDescribingTopics() {
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
// partitionWithoutInfo00 has rackInfo in cluster metadata
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithDescribingTopics() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
doReturn(
Collections.singletonMap(
- TOPIC0,
+ TP_0_NAME,
Collections.singletonList(
- new TopicPartitionInfo(0, node0, Arrays.asList(replicas1), Collections.emptyList())
+ new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
)
)
- ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
+ public void shouldDisableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
- doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(
+ TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
assertTrue(assignor.populateTopicsToDiscribe(new HashSet<>()));
}
+ @Test
+ public void shouldOptimizeEmptyActiveTasks() {
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet();
+
+ if (assignor.canEnableRackAwareAssignor()) {
+ final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, stateful);
+ assertEquals(0, originalCost);
+
+ final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, stateful);
+ assertEquals(0, cost);
+ }
+
+ assertEquals(mkSet(TASK_0_1, TASK_1_1), clientState0.activeTasks());
+ }
+
+ @Test
+ public void shouldOptimizeActiveTasks() {
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+ clientState1.assignActive(TASK_1_0);
+ clientState2.assignActive(TASK_0_0);
+
+ // task_0_0 has same rack as UUID_1
+ // task_0_1 has same rack as UUID_2 and UUID_3
+ // task_1_0 has same rack as UUID_1 and UUID_3
+ // task_1_1 has same rack as UUID_2
+ // Optimal assignment is UUID_1: {0_0, 1_0}, UUID_2: {1_1}, UUID_3: {0_1} which result in no cross rack traffic
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
+
+ if (assignor.canEnableRackAwareAssignor()) {
+ int expected = stateful ? 40 : 4;
+ final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, stateful);
+ assertEquals(expected, originalCost);
+
+ expected = stateful ? 4 : 0;
+ final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, stateful);
+ assertEquals(expected, cost);
+ }
+
+ assertEquals(mkSet(TASK_0_0, TASK_1_0), clientState0.activeTasks());
+ assertEquals(mkSet(TASK_1_1), clientState1.activeTasks());
+ assertEquals(mkSet(TASK_0_1), clientState2.activeTasks());
+ }
+
+ @Test
+ public void shouldOptimizeActiveTasksWithWeightOverride() {
+ final AssignmentConfigs assignmentConfigs = new AssignmentConfigs(1L, 2, 2, 60000L, Collections.emptyList(), 1, 10);
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ assignmentConfigs
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+ clientState1.assignActive(TASK_1_0);
+ clientState2.assignActive(TASK_0_0);
+
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
+
+ // Because non_overlap_cost is very high, this basically will stick to original assignment
+ if (assignor.canEnableRackAwareAssignor()) {
+ final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, stateful);
+ assertEquals(4, originalCost);
Review Comment:
In this test, `trafficCost` is 1 and `nonOverlapCost` is 10. So this cost 4 is the trafficCost. nonOverlapCost in original assignment should be 0 since tasks haven't been moved.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1271539596
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
It's also used on line 286.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1272692608
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java:
##########
@@ -268,24 +268,44 @@ public static class AssignmentConfigs {
public final long probingRebalanceIntervalMs;
public final List<String> rackAwareAssignmentTags;
+ // TODO: get from streamsConfig after we add the config
Review Comment:
I refactored this a bit and will add comments in `RackAwareTaskAssignor` method
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
Review Comment:
Or `moveTaskCost`? `moveAssignmentCost`? Have no strong preference here...
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
Yes. List is trying to make it deterministic. `keySet` of `SortedMap` to list should maintain order. That's why line 241 requires sorted map.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = taskIdList.size() + i;
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ private void assignActiveTaskFromMinCostFlow(final Graph<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalClientCapacity,
+ final Map<TaskId, UUID> taskClientMap) {
+ int taskAssigned = 0;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ if (edge.flow > 0) {
+ taskAssigned++;
+ final int clientIndex = edge.destination - taskIdList.size();
+ final UUID clientId = clientList.get(clientIndex);
+ final UUID originalClientId = taskClientMap.get(taskId);
+
+ // Don't need to assign this task to other client
+ if (clientId.equals(originalClientId)) {
+ break;
+ }
+
+ final ClientState clientState = clientStates.get(clientId);
+ final ClientState originalClientState = clientStates.get(originalClientId);
+ originalClientState.unassignActive(taskId);
+ clientState.assignActive(taskId);
+ }
Review Comment:
Yeah. There should be only one edge. I didn't break here for the validations below to catch anything wrong
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
Review Comment:
I'll refactor `isStateful` alway and pass in `trafficCost` and `nonOverlapCost` directly here.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
Review Comment:
The `keySet` or List we get should be deterministic which can ensure out min-cost assignment is deterministic. If it's not sorted when passed in, we need to sort it ourselves.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
Review Comment:
This is also to make it deterministic. `activeTasks` is a `SortedSet`
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
Review Comment:
Yeah. This is to mute some warning in Intellij, Checkstyle or spotBugs
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = taskIdList.size() + i;
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ private void assignActiveTaskFromMinCostFlow(final Graph<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalClientCapacity,
+ final Map<TaskId, UUID> taskClientMap) {
+ int taskAssigned = 0;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ if (edge.flow > 0) {
+ taskAssigned++;
+ final int clientIndex = edge.destination - taskIdList.size();
+ final UUID clientId = clientList.get(clientIndex);
+ final UUID originalClientId = taskClientMap.get(taskId);
+
+ // Don't need to assign this task to other client
+ if (clientId.equals(originalClientId)) {
+ break;
+ }
+
+ final ClientState clientState = clientStates.get(clientId);
+ final ClientState originalClientState = clientStates.get(originalClientId);
+ originalClientState.unassignActive(taskId);
+ clientState.assignActive(taskId);
+ }
+ }
+ }
+
+ // Validate task assigned
+ if (taskAssigned != activeTasks.size()) {
+ throw new IllegalStateException("Computed active task assignment number "
+ + taskAssigned + " is different size " + activeTasks.size());
+ }
+
+ // Validate capacity constraint
Review Comment:
This is the `originalAssignedTaskNum`. "load" is more accurate than `capacity`
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
Sure. I can add a comment to say it's `processId`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] lihaosky commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "lihaosky (via GitHub)" <gi...@apache.org>.
lihaosky commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1271098823
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -38,29 +43,34 @@
public class RackAwareTaskAssignor {
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
Review Comment:
One thing missing here is that I intend to have different default stateful traffic cost and non-overlap cost for `StickyAssignor` and `HAAssignor`, so instead of hardcoding here, maybe it makes more sense to pass in from constructor or assign() call.
This brings to another question of where should `RackAwareTaskAssignor` be constructed. Before `StickyAssignor` and `HAAssignor` are constructed using reflection, we can't pass `RackAwareTaskAssignor` to them via their constructor. We can construct `RackAwareTaskAssignor` inside them, but constructor of `RackAwareTaskAssignor` requires a lot more params than `HAAssignor`'s assign() call. So I'm thinking construct `RackAwareTaskAssignor` inside `StreamsPartitionAssignor` and pass it to `HAAssignor` or `StickyAssignor` using `assignmentConfig` variable...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [kafka] mjsax commented on a diff in pull request #14030: KAFKA-15022: [3/N] use graph to compute rack aware assignment for active stateful tasks
Posted by "mjsax (via GitHub)" <gi...@apache.org>.
mjsax commented on code in PR #14030:
URL: https://github.com/apache/kafka/pull/14030#discussion_r1272606193
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
Review Comment:
We `filter()` already for `isPresent` -- seems we only seen the first check?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
Review Comment:
Given that we don't read `taskClientMap` nor `originalAssignedTaskNumber`, should we just pass `new HashMap` directly?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
Review Comment:
nit: double whitespace `taskNodeId = 0`
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
Review Comment:
Same question as for `clientList`.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
Review Comment:
`clientId` is `processId` here, right? (renaming or add comment?)
##########
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java:
##########
@@ -145,147 +188,517 @@ public void disableActiveSinceRackMissingInClient() {
// False since process1 doesn't have rackId
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableActiveSinceRackDiffersInSameProcess() {
+ public void shouldDisableActiveWhenRackDiffersInSameProcess() {
final Map<UUID, Map<String, Optional<String>>> processRacks = new HashMap<>();
// Different consumers in same process have different rack ID. This shouldn't happen.
// If happens, there's a bug somewhere
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
processRacks,
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithoutDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithoutDescribingTopics() {
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
// partitionWithoutInfo00 has rackInfo in cluster metadata
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithDescribingTopics() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
doReturn(
Collections.singletonMap(
- TOPIC0,
+ TP_0_NAME,
Collections.singletonList(
- new TopicPartitionInfo(0, node0, Arrays.asList(replicas1), Collections.emptyList())
+ new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
)
)
- ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
+ public void shouldDisableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
- doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(
+ TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
assertTrue(assignor.populateTopicsToDiscribe(new HashSet<>()));
}
+ @Test
+ public void shouldOptimizeEmptyActiveTasks() {
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet();
+
+ if (assignor.canEnableRackAwareAssignor()) {
+ final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, stateful);
+ assertEquals(0, originalCost);
+
+ final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, stateful);
+ assertEquals(0, cost);
+ }
+
+ assertEquals(mkSet(TASK_0_1, TASK_1_1), clientState0.activeTasks());
+ }
+
+ @Test
+ public void shouldOptimizeActiveTasks() {
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+ clientState1.assignActive(TASK_1_0);
+ clientState2.assignActive(TASK_0_0);
+
+ // task_0_0 has same rack as UUID_1
+ // task_0_1 has same rack as UUID_2 and UUID_3
+ // task_1_0 has same rack as UUID_1 and UUID_3
+ // task_1_1 has same rack as UUID_2
+ // Optimal assignment is UUID_1: {0_0, 1_0}, UUID_2: {1_1}, UUID_3: {0_1} which result in no cross rack traffic
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
+
+ if (assignor.canEnableRackAwareAssignor()) {
+ int expected = stateful ? 40 : 4;
+ final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, stateful);
+ assertEquals(expected, originalCost);
+
+ expected = stateful ? 4 : 0;
+ final long cost = assignor.optimizeActiveTasks(clientStateMap, taskIds, stateful);
+ assertEquals(expected, cost);
+ }
+
+ assertEquals(mkSet(TASK_0_0, TASK_1_0), clientState0.activeTasks());
+ assertEquals(mkSet(TASK_1_1), clientState1.activeTasks());
+ assertEquals(mkSet(TASK_0_1), clientState2.activeTasks());
+ }
+
+ @Test
+ public void shouldOptimizeActiveTasksWithWeightOverride() {
+ final AssignmentConfigs assignmentConfigs = new AssignmentConfigs(1L, 2, 2, 60000L, Collections.emptyList(), 1, 10);
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ assignmentConfigs
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+ final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+ clientState1.assignActive(TASK_1_0);
+ clientState2.assignActive(TASK_0_0);
+
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0),
+ mkEntry(UUID_2, clientState1),
+ mkEntry(UUID_3, clientState2)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
+
+ // Because non_overlap_cost is very high, this basically will stick to original assignment
+ if (assignor.canEnableRackAwareAssignor()) {
+ final long originalCost = assignor.activeTasksCost(clientStateMap, taskIds, stateful);
+ assertEquals(4, originalCost);
Review Comment:
Why is expected cost not 40 for stateful as in the test above?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = taskIdList.size() + i;
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ private void assignActiveTaskFromMinCostFlow(final Graph<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalClientCapacity,
+ final Map<TaskId, UUID> taskClientMap) {
+ int taskAssigned = 0;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ if (edge.flow > 0) {
+ taskAssigned++;
+ final int clientIndex = edge.destination - taskIdList.size();
+ final UUID clientId = clientList.get(clientIndex);
+ final UUID originalClientId = taskClientMap.get(taskId);
+
+ // Don't need to assign this task to other client
+ if (clientId.equals(originalClientId)) {
+ break;
+ }
+
+ final ClientState clientState = clientStates.get(clientId);
+ final ClientState originalClientState = clientStates.get(originalClientId);
+ originalClientState.unassignActive(taskId);
+ clientState.assignActive(taskId);
+ }
+ }
+ }
+
+ // Validate task assigned
+ if (taskAssigned != activeTasks.size()) {
+ throw new IllegalStateException("Computed active task assignment number "
+ + taskAssigned + " is different size " + activeTasks.size());
+ }
+
+ // Validate capacity constraint
Review Comment:
Is this really a "capacity" or rather the "load" of a client?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> statefulTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(statefulTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> clientCapacity = new HashMap<>();
+ final Graph<Integer> graph = new Graph<>();
+
+ constructStatefulActiveTaskGraph(graph, statefulTasks, clientList, taskIdList,
+ clientStates, taskClientMap, clientCapacity, isStateful);
+
+ final int sourceId = taskIdList.size() + clientList.size();
+ final int sinkId = sourceId + 1;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(sourceId, taskNodeId, 1, 0, 1);
+ }
+ for (int i = 0; i < clientList.size(); i++) {
+ final int capacity = clientCapacity.getOrDefault(clientList.get(i), 0);
+ final int clientNodeId = taskIdList.size() + i;
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+ graph.setSourceNode(sourceId);
+ graph.setSinkNode(sinkId);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active stateful task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param taskIds Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> taskIds,
+ final boolean isStateful) {
+ if (taskIds.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
Question as above: why do we need to make a deep copy into a list? Can't we just pass `clientStates.keySet()` instead? Seems both methods only read but don't modify.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
Review Comment:
So all tasks passed in are stateful? If yes, maybe add "all" to the comments (also, if it's `false` does it mean all tasks are stateless?)
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
Review Comment:
Is it important for the method that the maps are sorted? Should we use `Map` (least restrictive) interface instead)?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
Review Comment:
Seems `constructActiveTaskGraph` is only reading but not modifying `clientList` -- why do we need to pass a deep-copy? Could we pass a `Set` instead of a `List`?
Looking into `constructActiveTaskGraph`, it seems we access by index and try to make thing deterministic. Is this the reason why we need a list here? If yes, given that we go from `keySet` to list, is this translation actually deterministic itself (could be given that it's a sorted-map).
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = taskIdList.size() + i;
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ private void assignActiveTaskFromMinCostFlow(final Graph<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalClientCapacity,
+ final Map<TaskId, UUID> taskClientMap) {
+ int taskAssigned = 0;
Review Comment:
`task[s]Assigned` ?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
Review Comment:
Not sure if I understand? How would a client with zero tasks improve balance?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = taskIdList.size() + i;
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ private void assignActiveTaskFromMinCostFlow(final Graph<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalClientCapacity,
+ final Map<TaskId, UUID> taskClientMap) {
+ int taskAssigned = 0;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ if (edge.flow > 0) {
+ taskAssigned++;
+ final int clientIndex = edge.destination - taskIdList.size();
+ final UUID clientId = clientList.get(clientIndex);
+ final UUID originalClientId = taskClientMap.get(taskId);
+
+ // Don't need to assign this task to other client
+ if (clientId.equals(originalClientId)) {
+ break;
+ }
+
+ final ClientState clientState = clientStates.get(clientId);
+ final ClientState originalClientState = clientStates.get(originalClientId);
+ originalClientState.unassignActive(taskId);
+ clientState.assignActive(taskId);
Review Comment:
Could we simplify this 4 lines to:
```
clientStates.get(originalClientId).unassignActive(taskId);
clientStates.get(clientId).assignActive(taskId);
```
##########
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignorTest.java:
##########
@@ -145,147 +188,517 @@ public void disableActiveSinceRackMissingInClient() {
// False since process1 doesn't have rackId
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableActiveSinceRackDiffersInSameProcess() {
+ public void shouldDisableActiveWhenRackDiffersInSameProcess() {
final Map<UUID, Map<String, Optional<String>>> processRacks = new HashMap<>();
// Different consumers in same process have different rack ID. This shouldn't happen.
// If happens, there's a bug somewhere
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
- processRacks.computeIfAbsent(process0UUID, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
+ processRacks.computeIfAbsent(UUID_1, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
processRacks,
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
assertFalse(assignor.validateClientRack());
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithoutDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithoutDescribingTopics() {
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterForTopic0(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
mockInternalTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
// partitionWithoutInfo00 has rackInfo in cluster metadata
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void enableRackAwareAssignorForActiveWithDescribingTopics() {
+ public void shouldEnableRackAwareAssignorForActiveWithDescribingTopics() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
doReturn(
Collections.singletonMap(
- TOPIC0,
+ TP_0_NAME,
Collections.singletonList(
- new TopicPartitionInfo(0, node0, Arrays.asList(replicas1), Collections.emptyList())
+ new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
)
)
- ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertTrue(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertTrue(assignor.canEnableRackAwareAssignor());
}
@Test
- public void disableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
+ public void shouldDisableRackAwareAssignorForActiveWithDescribingTopicsFailure() {
final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
- doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TOPIC0));
+ doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(
+ TP_0_NAME));
final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
getClusterWithNoNode(),
- getTaskTopicPartitionMapForTask1(),
+ getTaskTopicPartitionMapForTask0(),
getTopologyGroupTaskMap(),
getProcessRacksForProcess0(),
spyTopicManager,
new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
);
- assertFalse(assignor.canEnableRackAwareAssignorForActiveTasks());
+ assertFalse(assignor.canEnableRackAwareAssignor());
assertTrue(assignor.populateTopicsToDiscribe(new HashSet<>()));
}
+ @Test
+ public void shouldOptimizeEmptyActiveTasks() {
+ final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
+ getClusterForTopic0And1(),
+ getTaskTopicPartitionMapForAllTasks(),
+ getTopologyGroupTaskMap(),
+ getProcessRacksForAllProcess(),
+ mockInternalTopicManager,
+ new AssignorConfiguration(streamsConfig.originals()).assignmentConfigs()
+ );
+
+ final ClientState clientState0 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
+
+ clientState0.assignActiveTasks(mkSet(TASK_0_1, TASK_1_1));
+
+ final SortedMap<UUID, ClientState> clientStateMap = new TreeMap<>(mkMap(
+ mkEntry(UUID_1, clientState0)
+ ));
+ final SortedSet<TaskId> taskIds = mkSortedSet();
Review Comment:
Why is this set empty?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignorConfiguration.java:
##########
@@ -268,24 +268,44 @@ public static class AssignmentConfigs {
public final long probingRebalanceIntervalMs;
public final List<String> rackAwareAssignmentTags;
+ // TODO: get from streamsConfig after we add the config
Review Comment:
Might be good to add an extensive comment to explain what both do.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +193,212 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null || topicPartitions.isEmpty()) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
+ }
+
+ return cost;
+ }
+
+ private static int getSinkID(final List<UUID> clientList, final List<TaskId> taskIdList) {
+ return clientList.size() + taskIdList.size();
+ }
+
+ // For testing. canEnableRackAwareAssignor must be called first
+ long activeTasksCost(final SortedMap<UUID, ClientState> clientStates, final SortedSet<TaskId> activeTasks, final boolean isStateful) {
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+ return graph.totalCost();
+ }
+
+ /**
+ * Optimize active task assignment for rack awareness. canEnableRackAwareAssignor must be called first
+ * @param clientStates Client states
+ * @param activeTasks Tasks to reassign if needed. They must be assigned already in clientStates
+ * @param isStateful Whether the tasks are stateful
+ * @return Total cost after optimization
+ */
+ public long optimizeActiveTasks(final SortedMap<UUID, ClientState> clientStates,
+ final SortedSet<TaskId> activeTasks,
+ final boolean isStateful) {
+ if (activeTasks.isEmpty()) {
+ return 0;
+ }
+
+ final List<UUID> clientList = new ArrayList<>(clientStates.keySet());
+ final List<TaskId> taskIdList = new ArrayList<>(activeTasks);
+ final Map<TaskId, UUID> taskClientMap = new HashMap<>();
+ final Map<UUID, Integer> originalAssignedTaskNumber = new HashMap<>();
+ final Graph<Integer> graph = constructActiveTaskGraph(activeTasks, clientList, taskIdList,
+ clientStates, taskClientMap, originalAssignedTaskNumber, isStateful);
+
+ graph.solveMinCostFlow();
+ final long cost = graph.totalCost();
+
+ assignActiveTaskFromMinCostFlow(graph, activeTasks, clientList, taskIdList,
+ clientStates, originalAssignedTaskNumber, taskClientMap);
+
+ return cost;
+ }
+
+ private Graph<Integer> constructActiveTaskGraph(final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<TaskId, UUID> taskClientMap,
+ final Map<UUID, Integer> originalAssignedTaskNumber,
+ final boolean isStateful) {
+ final Graph<Integer> graph = new Graph<>();
+
+ for (final TaskId taskId : activeTasks) {
+ for (final Entry<UUID, ClientState> clientState : clientStates.entrySet()) {
+ if (clientState.getValue().hasAssignedTask(taskId)) {
+ originalAssignedTaskNumber.merge(clientState.getKey(), 1, Integer::sum);
+ }
+ }
+ }
+
+ // Make task and client Node id in graph deterministic
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ for (int j = 0; j < clientList.size(); j++) {
+ final int clientNodeId = taskIdList.size() + j;
+ final UUID clientId = clientList.get(j);
+
+ final int flow = clientStates.get(clientId).hasAssignedTask(taskId) ? 1 : 0;
+ final int cost = getCost(taskId, clientId, flow == 1, isStateful);
+ if (flow == 1) {
+ if (taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " assigned to multiple clients "
+ + clientId + ", " + taskClientMap.get(taskId));
+ }
+ taskClientMap.put(taskId, clientId);
+ }
+
+ graph.addEdge(taskNodeId, clientNodeId, 1, cost, flow);
+ }
+ if (!taskClientMap.containsKey(taskId)) {
+ throw new IllegalArgumentException("Task " + taskId + " not assigned to any client");
+ }
+ }
+
+ final int sinkId = getSinkID(clientList, taskIdList);
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ graph.addEdge(SOURCE_ID, taskNodeId, 1, 0, 1);
+ }
+
+ // It's possible that some clients have 0 task assign. These clients will have 0 tasks assigned
+ // even though it may have lower traffic cost. This is to maintain the balance requirement.
+ for (int i = 0; i < clientList.size(); i++) {
+ final int clientNodeId = taskIdList.size() + i;
+ final int capacity = originalAssignedTaskNumber.getOrDefault(clientList.get(i), 0);
+ // Flow equals to capacity for edges to sink
+ graph.addEdge(clientNodeId, sinkId, capacity, 0, capacity);
+ }
+
+ graph.setSourceNode(SOURCE_ID);
+ graph.setSinkNode(sinkId);
+
+ return graph;
+ }
+
+ private void assignActiveTaskFromMinCostFlow(final Graph<Integer> graph,
+ final SortedSet<TaskId> activeTasks,
+ final List<UUID> clientList,
+ final List<TaskId> taskIdList,
+ final Map<UUID, ClientState> clientStates,
+ final Map<UUID, Integer> originalClientCapacity,
+ final Map<TaskId, UUID> taskClientMap) {
+ int taskAssigned = 0;
+ for (int taskNodeId = 0; taskNodeId < taskIdList.size(); taskNodeId++) {
+ final TaskId taskId = taskIdList.get(taskNodeId);
+ final Map<Integer, Graph<Integer>.Edge> edges = graph.edges(taskNodeId);
+ for (final Graph<Integer>.Edge edge : edges.values()) {
+ if (edge.flow > 0) {
+ taskAssigned++;
+ final int clientIndex = edge.destination - taskIdList.size();
+ final UUID clientId = clientList.get(clientIndex);
+ final UUID originalClientId = taskClientMap.get(taskId);
+
+ // Don't need to assign this task to other client
+ if (clientId.equals(originalClientId)) {
+ break;
+ }
+
+ final ClientState clientState = clientStates.get(clientId);
+ final ClientState originalClientState = clientStates.get(originalClientId);
+ originalClientState.unassignActive(taskId);
+ clientState.assignActive(taskId);
+ }
Review Comment:
Should we `break` here (ie, inside the `if` block)? There should be only one edge with flow 1?
Or even replace the `for` loop over the edges with a `graph.edges(taskNodeId).values().stream().filter(e.flow == 1).findFirst` (and throw if we don't find any edge with flow == 1)?
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
Review Comment:
> Here I meant processId using the UUID.
Should we rename it? It's unfortunate that it's overloaded, but `clientId` sounds like `client.id` config, but it not. (Or add a comment?)
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -38,29 +43,34 @@
public class RackAwareTaskAssignor {
private static final Logger log = LoggerFactory.getLogger(RackAwareTaskAssignor.class);
+ private static final int DEFAULT_STATEFUL_TRAFFIC_COST = 10;
Review Comment:
Thanks for the in-person sync. To summarize what we discussed:
For RackAware assignment, we want to avoid to move a task from client A to client B if both clients are in the same rack (because the cross-rack-traffic cost is still the same anyway). -- However, if we can move a task from client A to client B and it will reduce cross-rack-traffic cost, we should move the task and neglect non-overlap-cost.
It make sense to me, to pass in the values as propose, and use different ones for sticky vs HA.
About where to construct it: I would say whatever is simplest :) -- I guess we could pass it via `assign()` directly, or just create inside `assign()` -- whatever is less complex to do.
##########
streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/RackAwareTaskAssignor.java:
##########
@@ -185,4 +191,224 @@ public boolean validateClientRack() {
}
return true;
}
+
+ private int getCost(final TaskId taskId, final UUID clientId, final boolean inCurrentAssignment, final boolean isStateful) {
+ final Map<String, Optional<String>> clientRacks = racksForProcess.get(clientId);
+ if (clientRacks == null) {
+ throw new IllegalStateException("Client " + clientId + " doesn't exist in processRacks");
+ }
+ final Optional<Optional<String>> clientRackOpt = clientRacks.values().stream().filter(Optional::isPresent).findFirst();
+ if (!clientRackOpt.isPresent() || !clientRackOpt.get().isPresent()) {
+ throw new IllegalStateException("Client " + clientId + " doesn't have rack configured. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+
+ final String clientRack = clientRackOpt.get().get();
+ final Set<TopicPartition> topicPartitions = partitionsForTask.get(taskId);
+ if (topicPartitions == null) {
+ throw new IllegalStateException("Task " + taskId + " has no TopicPartitions");
+ }
+
+ final int trafficCost = assignmentConfigs.trafficCost == null ? (isStateful ? DEFAULT_STATEFUL_TRAFFIC_COST : DEFAULT_STATELESS_TRAFFIC_COST)
+ : assignmentConfigs.trafficCost;
+ final int nonOverlapCost = assignmentConfigs.nonOverlapCost == null ? (isStateful ? DEFAULT_STATEFUL_NON_OVERLAP_COST : DEFAULT_STATELESS_NON_OVERLAP_COST)
+ : assignmentConfigs.nonOverlapCost;
+
+ int cost = 0;
+ for (final TopicPartition tp : topicPartitions) {
+ final Set<String> tpRacks = racksForPartition.get(tp);
+ if (tpRacks == null || tpRacks.isEmpty()) {
+ throw new IllegalStateException("TopicPartition " + tp + " has no rack information. Maybe forgot to call canEnableRackAwareAssignor first");
+ }
+ if (!tpRacks.contains(clientRack)) {
+ cost += trafficCost;
+ }
+ }
+
+ if (!inCurrentAssignment) {
+ cost += nonOverlapCost;
Review Comment:
Left already a comment above about adding a better explanation. Also wondering, if we could find a more descriptive name? Unfortunately, I don't have a good idea either.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org