[ML] Fix model assignment error handling and assignment explanation generation

docs/changelog/133916.yaml

@@ -0,0 +1,5 @@
+pr: 133916
+summary: Fix model assignment error handling and assignment explanation generation
+area: Machine Learning
+type: bug
+issues: []

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());
+                }
             }
         }
     }

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

@@ -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

@@ -1198,7 +1198,7 @@ public void testCopyAssignments() {
         assertThat(deployment2Assignments.get().get(node2), equalTo(1));
     }
 
-    public void testRebalance_GivenDeploymentWithMemoryRequirements_ConsidersNativeExecutableOverhead() {
+    public void testRebalance_GivenDeploymentWithMemoryRequirements_ExplainMissingAllocations() {
         // 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();