Fix memory usage estimation for ELSER models (elastic#131630)

* Pass model ID instead of deployment ID to memory estimator

* Update docs/changelog/131630.yaml

Author: jan-elastic

docs/changelog/131630.yaml

@@ -0,0 +1,5 @@
+pr: 131630
+summary: Fix memory usage estimation for ELSER models
+area: Machine Learning
+type: bug
+issues: []

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

@@ -169,6 +169,7 @@ private AssignmentPlan computePlanForNormalPriorityModels(
                     .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().getTargetAllocations()));
                 return new AssignmentPlan.Deployment(
                     assignment.getDeploymentId(),
+                    assignment.getModelId(),
                     assignment.getTaskParams().getModelBytes(),
                     assignment.getTaskParams().getNumberOfAllocations(),
                     assignment.getTaskParams().getThreadsPerAllocation(),
@@ -185,6 +186,7 @@ private AssignmentPlan computePlanForNormalPriorityModels(
             planDeployments.add(
                 new AssignmentPlan.Deployment(
                     taskParams.getDeploymentId(),
+                    taskParams.getModelId(),
                     taskParams.getModelBytes(),
                     taskParams.getNumberOfAllocations(),
                     taskParams.getThreadsPerAllocation(),
@@ -225,6 +227,7 @@ private AssignmentPlan computePlanForLowPriorityModels(Set<String> assignableNod
             .map(
                 assignment -> new AssignmentPlan.Deployment(
                     assignment.getDeploymentId(),
+                    assignment.getModelId(),
                     assignment.getTaskParams().getModelBytes(),
                     assignment.getTaskParams().getNumberOfAllocations(),
                     assignment.getTaskParams().getThreadsPerAllocation(),
@@ -242,6 +245,7 @@ private AssignmentPlan computePlanForLowPriorityModels(Set<String> assignableNod
             planDeployments.add(
                 new AssignmentPlan.Deployment(
                     taskParams.getDeploymentId(),
+                    taskParams.getModelId(),
                     taskParams.getModelBytes(),
                     taskParams.getNumberOfAllocations(),
                     taskParams.getThreadsPerAllocation(),

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

@@ -55,6 +55,7 @@ Deployment modifyModelPreservingPreviousAssignments(Deployment m) {
 
         return new Deployment(
             m.deploymentId(),
+            m.modelId(),
             m.memoryBytes(),
             m.allocations() - calculatePreservedAllocations(m),
             m.threadsPerAllocation(),

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

@@ -47,6 +47,7 @@ public class AssignmentPlan implements Comparable<AssignmentPlan> {
      */
     public record Deployment(
         String deploymentId,
+        String modelId,
         long memoryBytes,
         int allocations,
         int threadsPerAllocation,
@@ -59,6 +60,7 @@ public record Deployment(
     ) {
         public Deployment(
             String deploymentId,
+            String modelId,
             long modelBytes,
             int allocations,
             int threadsPerAllocation,
@@ -70,6 +72,7 @@ public Deployment(
         ) {
             this(
                 deploymentId,
+                modelId,
                 modelBytes,
                 allocations,
                 threadsPerAllocation,
@@ -96,7 +99,7 @@ boolean hasEverBeenAllocated() {
 
         public long estimateMemoryUsageBytes(int allocations) {
             return StartTrainedModelDeploymentAction.estimateMemoryUsageBytes(
-                deploymentId,
+                modelId,
                 memoryBytes,
                 perDeploymentMemoryBytes,
                 perAllocationMemoryBytes,
@@ -106,24 +109,23 @@ public long estimateMemoryUsageBytes(int allocations) {
 
         long estimateAdditionalMemoryUsageBytes(int allocationsOld, int allocationsNew) {
             return StartTrainedModelDeploymentAction.estimateMemoryUsageBytes(
-                deploymentId,
+                modelId,
                 memoryBytes,
                 perDeploymentMemoryBytes,
                 perAllocationMemoryBytes,
                 allocationsNew
             ) - StartTrainedModelDeploymentAction.estimateMemoryUsageBytes(
-                deploymentId,
+                modelId,
                 memoryBytes,
                 perDeploymentMemoryBytes,
                 perAllocationMemoryBytes,
                 allocationsOld
             );
-
         }
 
         long minimumMemoryRequiredBytes() {
             return StartTrainedModelDeploymentAction.estimateMemoryUsageBytes(
-                deploymentId,
+                modelId,
                 memoryBytes,
                 perDeploymentMemoryBytes,
                 perAllocationMemoryBytes,

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

@@ -115,6 +115,7 @@ private AssignmentPlan solveAllocatingAtLeastOnceModelsThatWerePreviouslyAllocat
             .map(
                 m -> new AssignmentPlan.Deployment(
                     m.deploymentId(),
+                    m.modelId(),
                     m.memoryBytes(),
                     1,
                     m.threadsPerAllocation(),
@@ -148,6 +149,7 @@ private AssignmentPlan solveAllocatingAtLeastOnceModelsThatWerePreviouslyAllocat
                 : Map.of();
             return new AssignmentPlan.Deployment(
                 m.deploymentId(),
+                m.modelId(),
                 m.memoryBytes(),
                 m.allocations(),
                 m.threadsPerAllocation(),

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

@@ -127,6 +127,7 @@ private AssignmentPlan computeZonePlan(
                 d -> new AssignmentPlan.Deployment(
                     // replace each deployment with a new deployment
                     d.deploymentId(),
+                    d.modelId(),
                     d.memoryBytes(),
                     deploymentIdToTargetAllocationsPerZone.get(d.deploymentId()),
                     d.threadsPerAllocation(),
@@ -163,6 +164,7 @@ private AssignmentPlan computePlanAcrossAllNodes(List<AssignmentPlan> plans) {
             .map(
                 d -> new AssignmentPlan.Deployment(
                     d.deploymentId(),
+                    d.modelId(),
                     d.memoryBytes(),
                     d.allocations(),
                     d.threadsPerAllocation(),