Skip to content

Commit 971cfb9

Browse files
valeriy42valeriy42
authored
[ML] Refactor assignment planner code (#104260)
This PR simplifies the code in a few places. In other places where I had the TODO comment, the possible simplification would lead to undesired consequences, so I removed the TODO comment with the reference to issue #101612. Closes #101612
1 parent 1b2e8cf commit 971cfb9

5 files changed

Lines changed: 14 additions & 27 deletions

File tree

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -142,13 +142,9 @@ private static void copyAssignments(
142142
for (Map.Entry<AssignmentPlan.Node, Integer> assignment : nodeAssignments.entrySet()) {
143143
AssignmentPlan.Node originalNode = originalNodeById.get(assignment.getKey().id());
144144
dest.assignModelToNode(m, originalNode, assignment.getValue());
145-
if (m.currentAllocationsByNodeId().containsKey(originalNode.id())) {
146-
// TODO (#101612) requiredMemory should be calculated by the AssignmentPlan.Builder
147-
// As the node has all its available memory we need to manually account memory of models with
148-
// current allocations.
149-
long requiredMemory = m.estimateMemoryUsageBytes(m.currentAllocationsByNodeId().get(originalNode.id()));
150-
dest.accountMemory(m, originalNode, requiredMemory);
151-
}
145+
// As the node has all its available memory we need to manually account memory of models with
146+
// current allocations.
147+
dest.accountMemory(m, originalNode);
152148
}
153149
}
154150
}

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AbstractPreserveAllocations.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ Deployment modifyModelPreservingPreviousAssignments(Deployment m) {
6868
AssignmentPlan mergePreservedAllocations(AssignmentPlan assignmentPlan) {
6969
// As the model/node objects the assignment plan are the modified ones,
7070
// they will not match the models/nodes members we have in this class.
71-
// Therefore, we build a lookup table based on the ids so we can merge the plan
71+
// Therefore, we build a lookup table based on the ids, so we can merge the plan
7272
// with its preserved allocations.
7373
final Map<Tuple<String, String>, Integer> plannedAssignmentsByModelNodeIdPair = new HashMap<>();
7474
for (Deployment m : assignmentPlan.models()) {
@@ -80,7 +80,6 @@ AssignmentPlan mergePreservedAllocations(AssignmentPlan assignmentPlan) {
8080

8181
AssignmentPlan.Builder mergedPlanBuilder = AssignmentPlan.builder(nodes, deployments);
8282
for (Node n : nodes) {
83-
// TODO (#101612) Should the first loop happen in the builder constructor?
8483
for (Deployment deploymentAllocationsToPreserve : deployments) {
8584

8685
// if the model m is already allocated on the node n and I want to preserve this allocation

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlan.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -401,8 +401,7 @@ public Builder assignModelToNode(Deployment deployment, Node node, int allocatio
401401
if (allocations <= 0) {
402402
return this;
403403
}
404-
if (/*isAlreadyAssigned(deployment, node) == false
405-
&&*/ requiredMemory > remainingNodeMemory.get(node)) {
404+
if (requiredMemory > remainingNodeMemory.get(node)) {
406405
throw new IllegalArgumentException(
407406
"not enough memory on node ["
408407
+ node.id()
@@ -448,13 +447,14 @@ private static int getCurrentAllocations(Deployment m, Node n) {
448447
}
449448

450449
public void accountMemory(Deployment m, Node n) {
451-
// TODO (#101612) remove or refactor unused method
452-
long requiredMemory = getDeploymentMemoryRequirement(m, n, getCurrentAllocations(m, n));
453-
accountMemory(m, n, requiredMemory);
450+
if (m.currentAllocationsByNodeId().containsKey(n.id())) {
451+
int allocations = m.currentAllocationsByNodeId().get(n.id());
452+
long requiredMemory = m.estimateMemoryUsageBytes(allocations);
453+
accountMemory(m, n, requiredMemory);
454+
}
454455
}
455456

456-
public void accountMemory(Deployment m, Node n, long requiredMemory) {
457-
// TODO (#101612) computation of required memory should be done internally
457+
private void accountMemory(Deployment m, Node n, long requiredMemory) {
458458
remainingNodeMemory.computeIfPresent(n, (k, v) -> v - requiredMemory);
459459
if (remainingNodeMemory.containsKey(n) && remainingNodeMemory.get(n) < 0) {
460460
throw new IllegalArgumentException("not enough memory on node [" + n.id() + "] to assign model [" + m.id() + "]");

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/RandomizedAssignmentRounding.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -310,8 +310,6 @@ private void unassignOversizedModels(Node n) {
310310
private AssignmentPlan toPlan() {
311311
AssignmentPlan.Builder builder = AssignmentPlan.builder(nodes, deployments);
312312
for (Map.Entry<Tuple<AssignmentPlan.Deployment, Node>, Integer> assignment : tryAssigningRemainingCores().entrySet()) {
313-
// TODO (#101612) The model should be assigned to the node only when it is possible. This means, that canAssign should be
314-
// integrated into the assignModelToNode.
315313
if (builder.canAssign(assignment.getKey().v1(), assignment.getKey().v2(), assignment.getValue())) {
316314
builder.assignModelToNode(assignment.getKey().v1(), assignment.getKey().v2(), assignment.getValue());
317315
}

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/ZoneAwareAssignmentPlanner.java‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -183,15 +183,9 @@ private AssignmentPlan swapOriginalModelsInPlan(
183183
for (Map.Entry<Node, Integer> assignment : nodeAssignments.entrySet()) {
184184
Node originalNode = originalNodeById.get(assignment.getKey().id());
185185
planBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());
186-
if (originalDeployment.currentAllocationsByNodeId().containsKey(originalNode.id())) {
187-
// TODO (#101612) requiredMemory should be calculated by the AssignmentPlan.Builder
188-
// As the node has all its available memory we need to manually account memory of models with
189-
// current allocations.
190-
long requiredMemory = originalDeployment.estimateMemoryUsageBytes(
191-
originalDeployment.currentAllocationsByNodeId().get(originalNode.id())
192-
);
193-
planBuilder.accountMemory(m, originalNode, requiredMemory);
194-
}
186+
// As the node has all its available memory we need to manually account memory of models with
187+
// current allocations.
188+
planBuilder.accountMemory(originalDeployment, originalNode);
195189
}
196190
}
197191
return planBuilder.build();

0 commit comments

Comments
 (0)