[ML] Fix double-counting of inference memory in the assignment rebalancer (#133919) (#134052)

The static method TrainedModelAssignmentRebalancer.getNodeFreeMemoryExcludingPerNodeOverheadAndNativeInference was used to subtract load.getAssignedNativeInferenceMemory() from load.getFreeMemoryExcludingPerNodeOverhead(). However, in NodeLoad.getFreeMemoryExcludingPerNodeOverhead(), native inference memory was already subtracted as part of the getAssignedJobMemoryExcludingPerNodeOverhead() calculation. This led to double-counting of the native inference memory. Avoiding this double-counting allows us to remove the private method getNodeFreeMemoryExcludingPerNodeOverheadAndNativeInference() entirely.

Author: valeriy42

docs/changelog/133919.yaml

@@ -0,0 +1,5 @@
+pr: 133919
+summary: Fix double-counting of inference memory in the assignment rebalancer
+area: Machine Learning
+type: bug
+issues: []

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

@@ -303,9 +303,7 @@ private Map<List<String>, List<AssignmentPlan.Node>> createNodesByZoneMap() {
  nodes.add(
  new AssignmentPlan.Node(
  discoveryNode.getId(),
- // We subtract native inference memory as the planner expects available memory for
- // native inference including current assignments.
- getNodeFreeMemoryExcludingPerNodeOverheadAndNativeInference(load),
+ load.getFreeMemoryExcludingPerNodeOverhead(),
  MlProcessors.get(discoveryNode, allocatedProcessorsScale).roundUp()
  )
  );
@@ -322,10 +320,6 @@ private Map<List<String>, List<AssignmentPlan.Node>> createNodesByZoneMap() {
  }));
  }
 
- private static long getNodeFreeMemoryExcludingPerNodeOverheadAndNativeInference(NodeLoad load) {
- return load.getFreeMemoryExcludingPerNodeOverhead() - load.getAssignedNativeInferenceMemory();
- }
-
  private TrainedModelAssignmentMetadata.Builder buildAssignmentsFromPlan(AssignmentPlan assignmentPlan) {
  TrainedModelAssignmentMetadata.Builder builder = TrainedModelAssignmentMetadata.Builder.empty();
  for (AssignmentPlan.Deployment deployment : assignmentPlan.deployments()) {