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

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

@@ -138,9 +138,7 @@ 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());
-                if (dest.canAssign(deployment, node, sourceAssignment.getValue())) {
-                    dest.assignModelToNode(deployment, node, sourceAssignment.getValue());
-                }
+                dest.assignModelToNode(deployment, node, sourceAssignment.getValue());
             }
         }
     }

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

@@ -103,8 +103,7 @@ AssignmentPlan mergePreservedAllocations(AssignmentPlan assignmentPlan) {
                     0
                 );
 
-                long requiredMemory = mergedPlanBuilder.getDeploymentMemoryRequirement(deploymentNewAllocations, n, newAllocations);
-                if (newAllocations > 0 && mergedPlanBuilder.canAssign(deploymentNewAllocations, n, newAllocations, requiredMemory)) {
+                if (newAllocations > 0) {
                     mergedPlanBuilder.assignModelToNode(deploymentNewAllocations, n, newAllocations);
                 }
             }

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);
         }
 
-        public boolean canAssign(Deployment deployment, Node node, int allocations) {
+        boolean canAssign(Deployment deployment, Node node, int allocations) {
             long requiredMemory = getDeploymentMemoryRequirement(deployment, node, allocations);
             return canAssign(deployment, node, allocations, requiredMemory);
         }
@@ -441,7 +441,7 @@ public Builder assignModelToNode(Deployment deployment, Node node, int allocatio
         }
 
         public Builder assignModelToNode(Deployment deployment, Node node, int allocations, long requiredMemory) {
-            if (allocations <= 0) {
+            if (allocations <= 0 || canAssign(deployment, node, allocations, requiredMemory) == false) {
                 return this;
             }
             if (requiredMemory > remainingNodeMemory.get(node)) {

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

@@ -310,11 +310,7 @@ private void unassignOversizedModels(Node n) {
         private AssignmentPlan toPlan() {
             AssignmentPlan.Builder builder = AssignmentPlan.builder(nodes, deployments);
             for (Map.Entry<Tuple<AssignmentPlan.Deployment, Node>, Integer> assignment : tryAssigningRemainingCores().entrySet()) {
-                // TODO (#101612) The model should be assigned to the node only when it is possible. This means, that canAssign should be
-                // integrated into the assignModelToNode.
-                if (builder.canAssign(assignment.getKey().v1(), assignment.getKey().v2(), assignment.getValue())) {
-                    builder.assignModelToNode(assignment.getKey().v1(), assignment.getKey().v2(), assignment.getValue());
-                }
+                builder.assignModelToNode(assignment.getKey().v1(), assignment.getKey().v2(), assignment.getValue());
             }
             return builder.build();
         }

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

@@ -91,6 +91,8 @@ private AssignmentPlan computePlan(boolean tryAssigningPreviouslyAssignedModels)
                 remainingZones,
                 tryAssigningPreviouslyAssignedModels
             );
+
+            // Update remaining allocations to account for allocations satisfied in this zone
             plan.deployments()
                 .forEach(
                     d -> deploymentIdToRemainingAllocations.computeIfPresent(
@@ -211,14 +213,20 @@ 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());
-                if (finalPlanBuilder.canAssign(originalDeployment, originalNode, assignment.getValue())) {
-                    finalPlanBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());
-                }
+                finalPlanBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());
             }
         }
         return finalPlanBuilder.build();
     }
 
+    /**
+     * The mergeAllocationsByNodeIdByDeploymentId method is responsible for consolidating allocation data
+     * from multiple AssignmentPlan objects into a single structure. This structure maps deployment IDs
+     * to their respective node allocations, allowing the system to track how resources are distributed
+     * across nodes for each deployment.
+     * @param plans List of AssignmentPlan objects to merge allocations from
+     * @return
+     */
     private Map<String, Map<String, Integer>> mergeAllocationsByNodeIdByDeploymentId(List<AssignmentPlan> plans) {
         Map<String, Map<String, Integer>> allocationsByNodeIdByDeploymentId = new HashMap<>();
         deployments.forEach(d -> allocationsByNodeIdByDeploymentId.put(d.deploymentId(), new HashMap<>()));