Preserve deployments with zero allocations during assignment planning.

Author: jan-elastic

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

@@ -16,6 +16,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -190,15 +191,13 @@ public String toString() {
 
     private AssignmentPlan(
         Map<Deployment, Map<Node, Integer>> assignments,
-        Map<Node, Long> remainingNodeMemory,
-        Map<Node, Integer> remainingNodeCores,
+        Map<String, Long> remainingNodeMemory,
+        Map<String, Integer> remainingNodeCores,
         Map<Deployment, Integer> remainingModelAllocations
     ) {
         this.assignments = Objects.requireNonNull(assignments);
-        this.remainingNodeMemory = remainingNodeMemory.entrySet()
-            .stream()
-            .collect(Collectors.toMap(e -> e.getKey().id(), e -> e.getValue()));
-        this.remainingNodeCores = remainingNodeCores.entrySet().stream().collect(Collectors.toMap(e -> e.getKey().id(), e -> e.getValue()));
+        this.remainingNodeMemory = Objects.requireNonNull(remainingNodeMemory);
+        this.remainingNodeCores = Objects.requireNonNull(remainingNodeCores);
         this.remainingModelAllocations = Objects.requireNonNull(remainingModelAllocations);
     }
 
@@ -315,6 +314,14 @@ private Quality computeQuality() {
         return new Quality(isSatisfyingPreviousAssignments, weighedAllocationsScore, memoryScore);
     }
 
+    public AssignmentPlan withDeploymentsWithZeroAllocations(Collection<Deployment> deploymentsWithZeroAllocations) {
+        Map<Deployment, Map<Node, Integer>> newAssignments = new HashMap<>(assignments);
+        for (Deployment deployment : deploymentsWithZeroAllocations) {
+            newAssignments.put(deployment, Collections.emptyMap());
+        }
+        return new AssignmentPlan(newAssignments, remainingNodeMemory, remainingNodeCores, remainingModelAllocations);
+    }
+
     public String prettyPrint() {
         if (assignments.isEmpty()) {
             return "Empty plan";
@@ -477,7 +484,12 @@ public AssignmentPlan build() {
                 }
                 finalAssignments.put(m, allocationsPerNode);
             }
-            return new AssignmentPlan(finalAssignments, remainingNodeMemory, remainingNodeCores, remainingModelAllocations);
+            return new AssignmentPlan(
+                finalAssignments,
+                remainingNodeMemory.entrySet().stream().collect(Collectors.toMap(e -> e.getKey().id(), Map.Entry::getValue)),
+                remainingNodeCores.entrySet().stream().collect(Collectors.toMap(e -> e.getKey().id(), Map.Entry::getValue)),
+                remainingModelAllocations
+            );
         }
     }
 

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

@@ -47,13 +47,18 @@ public class AssignmentPlanner {
 
     private final List<Node> nodes;
     private final List<AssignmentPlan.Deployment> deployments;
+    private final List<AssignmentPlan.Deployment> deploymentsWithZeroAllocations;
 
     public AssignmentPlanner(List<Node> nodes, List<AssignmentPlan.Deployment> deployments) {
         this.nodes = nodes.stream().sorted(Comparator.comparing(Node::id)).toList();
         this.deployments = deployments.stream()
             .filter(deployment -> deployment.allocations() > 0)
             .sorted(Comparator.comparing(AssignmentPlan.Deployment::deploymentId))
             .toList();
+        this.deploymentsWithZeroAllocations = deployments.stream()
+            .filter(deployment -> deployment.allocations() == 0)
+            .sorted(Comparator.comparing(AssignmentPlan.Deployment::deploymentId))
+            .toList();
     }
 
     public AssignmentPlan computePlan() {
@@ -99,7 +104,8 @@ public AssignmentPlan computePlan(boolean tryAssigningAllPreviouslyAllocatedMode
 
         logger.debug(() -> "Best plan =\n" + bestPlan.prettyPrint());
         logger.debug(() -> prettyPrintOverallStats(bestPlan));
-        return bestPlan;
+
+        return bestPlan.withDeploymentsWithZeroAllocations(deploymentsWithZeroAllocations);
     }
 
     /**

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlannerTests.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.StopWatch;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.ml.inference.assignment.planning.AssignmentPlan.Deployment;
 import org.elasticsearch.xpack.ml.inference.assignment.planning.AssignmentPlan.Node;
 
@@ -24,6 +25,7 @@
 import java.util.stream.Stream;
 
 import static org.elasticsearch.test.hamcrest.OptionalMatchers.isEmpty;
+import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.everyItem;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
@@ -1165,7 +1167,28 @@ public void testGivenClusterResize_ShouldRemoveAllocatedDeployments_NewMemoryFie
         assertThat(indexedBasedPlan.get("m_1"), equalTo(Map.of("n_1", 2)));
         assertThat(assignmentPlan.getRemainingNodeMemory(node1.id()), greaterThanOrEqualTo(0L));
         assertThat(assignmentPlan.getRemainingNodeCores(node1.id()), greaterThanOrEqualTo(0));
+    }
+
+    public void testDeploymentWithZeroAllocationsIsPreserved() {
+        Node node = new Node("n_1", ByteSizeValue.ofGb(4).getBytes(), 4);
+        Deployment deployment = new Deployment(
+            "m_1",
+            "m_1",
+            ByteSizeValue.ofMb(500).getBytes(),
+            0,
+            1,
+            Map.of(),
+            0,
+            new AdaptiveAllocationsSettings(true, 0, 42),
+            0,
+            0
+        );
+
+        AssignmentPlan assignmentPlan = new AssignmentPlanner(List.of(node), List.of(deployment)).computePlan();
 
+        assertThat(assignmentPlan.deployments(), contains(deployment));
+        assertThat(assignmentPlan.getRemainingNodeMemory(node.id()), equalTo(ByteSizeValue.ofGb(4).getBytes()));
+        assertThat(assignmentPlan.getRemainingNodeCores(node.id()), equalTo(4));
     }
 
     public static List<Deployment> createDeploymentsFromPlan(AssignmentPlan plan) {