Test before model assignment.

Author: valeriy42

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

@@ -138,7 +138,9 @@ static void copyAssignments(AssignmentPlan source, AssignmentPlan.Builder dest,
             Map<AssignmentPlan.Node, Integer> sourceNodeAssignments = source.assignments(deployment).orElse(Map.of());
             for (Map.Entry<AssignmentPlan.Node, Integer> sourceAssignment : sourceNodeAssignments.entrySet()) {
                 AssignmentPlan.Node node = originalNodeById.get(sourceAssignment.getKey().id());
-                dest.assignModelToNode(deployment, node, sourceAssignment.getValue());
+                if(dest.canAssign(deployment, node, sourceAssignment.getValue())) {
+                    dest.assignModelToNode(deployment, node, sourceAssignment.getValue());
+                }
             }
         }
     }
@@ -318,7 +320,12 @@ private Map<List<String>, List<AssignmentPlan.Node>> createNodesByZoneMap() {
     }
 
     private static long getNodeFreeMemoryExcludingPerNodeOverheadAndNativeInference(NodeLoad load) {
-        return load.getFreeMemoryExcludingPerNodeOverhead() - load.getAssignedNativeInferenceMemory();
+        // load.getFreeMemoryExcludingPerNodeOverhead() = maxMemory - assignedJobMemoryExcludingPerNodeOverhead - 30MB native executable code overhead
+        // assignedJobMemoryExcludingPerNodeOverhead = assignedAnomalyDetectorMemory + assignedDataFrameAnalyticsMemory + assignedNativeInferenceMemory
+        // load.getAssignedNativeInferenceMemory() = assignedNativeInferenceMemory
+        // TODO: (valeriy) assignedNativeInferenceMemory is double counted in the current calculation.
+        return load.getFreeMemoryExcludingPerNodeOverhead()/* - load.getAssignedNativeInferenceMemory()*/;
+
     }
 
     private TrainedModelAssignmentMetadata.Builder buildAssignmentsFromPlan(AssignmentPlan assignmentPlan) {
@@ -405,8 +412,17 @@ private Optional<String> explainAssignment(
         if (Strings.isNullOrEmpty(load.getError()) == false) {
             return Optional.of(load.getError());
         }
-
-        if (deployment.memoryBytes() > assignmentPlan.getRemainingNodeMemory(node.getId())) {
+        // TODO (valeriy): this test should be actually true, but it is false, because we use the "naked" deployment footprint
+        // Get existing allocations for this node to avoid double counting
+        int existingAllocationsOnNode = assignmentPlan.assignments(deployment)
+            .flatMap(assignments -> assignments.entrySet().stream()
+                .filter(entry -> entry.getKey().id().equals(node.getId()))
+                .findFirst()
+                .map(Map.Entry::getValue))
+            .orElse(0);
+        int notYetAssignedAllocations = deployment.allocations() - assignmentPlan.totalAllocations(deployment);
+//        if (deployment.estimateMemoryUsageBytes(deployment.allocations() - existingAllocationsOnNode) > assignmentPlan.getRemainingNodeMemory(node.getId())) {
+        if (deployment.estimateAdditionalMemoryUsageBytes(existingAllocationsOnNode, existingAllocationsOnNode + notYetAssignedAllocations) > assignmentPlan.getRemainingNodeMemory(node.getId())) {
             // If any ML processes are running on a node we require some space to load the shared libraries.
             // So if none are currently running then this per-node overhead must be added to the requirement.
             // From node load we know if we had any jobs or models assigned before the rebalance.

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

@@ -107,7 +107,7 @@ public long estimateMemoryUsageBytes(int allocations) {
             );
         }
 
-        long estimateAdditionalMemoryUsageBytes(int allocationsOld, int allocationsNew) {
+        public long estimateAdditionalMemoryUsageBytes(int allocationsOld, int allocationsNew) {
             return StartTrainedModelDeploymentAction.estimateMemoryUsageBytes(
                 modelId,
                 memoryBytes,
@@ -417,7 +417,7 @@ int getRemainingAllocations(Deployment m) {
             return remainingModelAllocations.get(m);
         }
 
-        boolean canAssign(Deployment deployment, Node node, int allocations) {
+        public boolean canAssign(Deployment deployment, Node node, int allocations) {
             long requiredMemory = getDeploymentMemoryRequirement(deployment, node, allocations);
             return canAssign(deployment, node, allocations, requiredMemory);
         }

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

@@ -211,7 +211,9 @@ private AssignmentPlan swapOriginalDeploymentsInPlan(
             Map<Node, Integer> nodeAssignments = plan.assignments(planDeployment).orElse(Map.of());
             for (Map.Entry<Node, Integer> assignment : nodeAssignments.entrySet()) {
                 Node originalNode = originalNodeById.get(assignment.getKey().id());
-                finalPlanBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());
+                if (finalPlanBuilder.canAssign(originalDeployment, originalNode, assignment.getValue())) {
+                    finalPlanBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());
+                }
             }
         }
         return finalPlanBuilder.build();

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

@@ -1262,4 +1262,72 @@ private static DiscoveryNode buildNode(String name, long nativeMemory, int alloc
             )
             .build();
     }
+
+    public void testRebalance_GivenDeploymentWithMemoryRequirements_ConsidersNativeExecutableOverhead() {
+        // Create a node with just enough memory to fit the model plus native executable overhead
+        long modelMemory = ByteSizeValue.ofMb(200).getBytes();
+        long memoryOverhead = ByteSizeValue.ofMb(240).getBytes();
+        long JVMOverhead = ByteSizeValue.ofMb(50).getBytes();
+        long nodeMemory = memoryOverhead + modelMemory*2 + JVMOverhead;
+
+        DiscoveryNode node = buildNode("node-1", nodeMemory, 4);
+
+        String deploymentId = "model-with-overhead-test";
+        StartTrainedModelDeploymentAction.TaskParams taskParams = normalPriorityParams(
+            deploymentId,
+            deploymentId,
+            modelMemory,
+            1,
+            1
+        );
+
+        TrainedModelAssignmentMetadata currentMetadata = TrainedModelAssignmentMetadata.Builder.empty().build();
+        Map<DiscoveryNode, NodeLoad> nodeLoads = new HashMap<>();
+
+        // This node has no jobs or models yet, so the overhead should be accounted for
+        nodeLoads.put(node, NodeLoad.builder("node-1")
+            .setMaxMemory(nodeMemory)
+            .build());
+
+        TrainedModelAssignmentMetadata result = new TrainedModelAssignmentRebalancer(
+            currentMetadata,
+            nodeLoads,
+            Map.of(List.of(), List.of(node)),
+            Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)),
+            1
+        ).rebalance().build();
+
+        // Verify the deployment was successful
+        TrainedModelAssignment assignment = result.getDeploymentAssignment(deploymentId);
+        assertThat(assignment, is(notNullValue()));
+        assertThat(assignment.getAssignmentState(), equalTo(AssignmentState.STARTING));
+        assertThat(assignment.getNodeRoutingTable(), is(aMapWithSize(1)));
+        assertThat(assignment.getNodeRoutingTable(), hasKey("node-1"));
+        assertThat(assignment.getReason().isPresent(), is(false));
+
+        // Now try with a node that has slightly less memory - this should fail
+        long insufficientNodeMemory = nodeMemory - ByteSizeValue.ofMb(21).getBytes();
+        DiscoveryNode insufficientNode = buildNode("node-2", insufficientNodeMemory, 4);
+
+        Map<DiscoveryNode, NodeLoad> insufficientNodeLoads = Map.of(
+            insufficientNode, NodeLoad.builder("node-2")
+                .setMaxMemory(insufficientNodeMemory)
+                .build()
+        );
+
+        TrainedModelAssignmentMetadata insufficientResult = new TrainedModelAssignmentRebalancer(
+            TrainedModelAssignmentMetadata.Builder.empty().build(),
+            insufficientNodeLoads,
+            Map.of(List.of(), List.of(insufficientNode)),
+            Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)),
+            1)
+            .rebalance().build();
+
+        TrainedModelAssignment insufficientAssignment = insufficientResult.getDeploymentAssignment(deploymentId);
+        assertThat(insufficientAssignment, is(notNullValue()));
+        assertThat(insufficientAssignment.getAssignmentState(), equalTo(AssignmentState.STARTING));
+        assertThat(insufficientAssignment.getNodeRoutingTable(), is(anEmptyMap()));
+        assertThat(insufficientAssignment.getReason().isPresent(), is(true));
+        assertThat(insufficientAssignment.getReason().get(), containsString("insufficient available memory"));
+    }
 }