[ML] Remove irrelevant useNewMemoryFields parameter (#129298)

Author: davidkyle

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

@@ -9,7 +9,6 @@
 
 import org.elasticsearch.ResourceAlreadyExistsException;
 import org.elasticsearch.ResourceNotFoundException;
-import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.cluster.SimpleDiffable;
 import org.elasticsearch.common.Randomness;
@@ -107,10 +106,6 @@ public final class TrainedModelAssignment implements SimpleDiffable<TrainedModel
     private final int maxAssignedAllocations;
     private final AdaptiveAllocationsSettings adaptiveAllocationsSettings;
 
-    public static boolean useNewMemoryFields(TransportVersion minClusterVersion) {
-        return minClusterVersion.onOrAfter(TransportVersions.V_8_11_X);
-    }
-
     public static TrainedModelAssignment fromXContent(XContentParser parser) throws IOException {
         return PARSER.apply(parser, null);
     }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsAction.java

@@ -50,12 +50,10 @@
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelType;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentStats;
-import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
 import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignmentMetadata;
 import org.elasticsearch.xpack.core.ml.inference.persistence.InferenceIndexConstants;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceStats;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TrainedModelSizeStats;
-import org.elasticsearch.xpack.core.ml.utils.TransportVersionUtils;
 import org.elasticsearch.xpack.ml.MachineLearning;
 import org.elasticsearch.xpack.ml.inference.persistence.TrainedModelDefinitionDoc;
 import org.elasticsearch.xpack.ml.inference.persistence.TrainedModelProvider;
@@ -328,15 +326,12 @@ private void modelSizeStats(
                         long totalDefinitionLength = pytorchTotalDefinitionLengthsByModelId.getOrDefault(model.getModelId(), 0L);
                         // We ensure that in the mixed cluster state trained model stats uses the same values for memory estimation
                         // as the rebalancer.
-                        boolean useNewMemoryFields = TrainedModelAssignment.useNewMemoryFields(
-                            TransportVersionUtils.getMinTransportVersion(clusterService.state())
-                        );
                         long estimatedMemoryUsageBytes = totalDefinitionLength > 0L
                             ? StartTrainedModelDeploymentAction.estimateMemoryUsageBytes(
                                 model.getModelId(),
                                 totalDefinitionLength,
-                                useNewMemoryFields ? model.getPerDeploymentMemoryBytes() : 0,
-                                useNewMemoryFields ? model.getPerAllocationMemoryBytes() : 0,
+                                model.getPerDeploymentMemoryBytes(),
+                                model.getPerAllocationMemoryBytes(),
                                 numberOfAllocations
                             )
                             : 0L;

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

@@ -52,7 +52,6 @@
 import org.elasticsearch.xpack.core.ml.job.messages.Messages;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 import org.elasticsearch.xpack.core.ml.utils.MlPlatformArchitecturesUtil;
-import org.elasticsearch.xpack.core.ml.utils.TransportVersionUtils;
 import org.elasticsearch.xpack.ml.MachineLearning;
 import org.elasticsearch.xpack.ml.autoscaling.NodeAvailabilityZoneMapper;
 import org.elasticsearch.xpack.ml.inference.assignment.planning.AllocationReducer;
@@ -651,14 +650,12 @@ private TrainedModelAssignmentMetadata.Builder rebalanceAssignments(
         Map<DiscoveryNode, NodeLoad> nodeLoads = detectNodeLoads(nodes, currentState);
         TrainedModelAssignmentMetadata currentMetadata = TrainedModelAssignmentMetadata.fromState(currentState);
 
-        boolean useNewMemoryFields = TrainedModelAssignment.useNewMemoryFields(TransportVersionUtils.getMinTransportVersion(currentState));
         TrainedModelAssignmentRebalancer rebalancer = new TrainedModelAssignmentRebalancer(
             currentMetadata,
             nodeLoads,
             nodeAvailabilityZoneMapper.buildMlNodesByAvailabilityZone(currentState),
             createAssignmentRequest,
-            allocatedProcessorsScale,
-            useNewMemoryFields
+            allocatedProcessorsScale
         );
 
         Set<String> shuttingDownNodeIds = currentState.metadata().nodeShutdowns().getAllNodeIds();

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

@@ -54,22 +54,18 @@ class TrainedModelAssignmentRebalancer {
     private final Optional<CreateTrainedModelAssignmentAction.Request> createAssignmentRequest;
     private final int allocatedProcessorsScale;
 
-    private final boolean useNewMemoryFields;
-
     TrainedModelAssignmentRebalancer(
         TrainedModelAssignmentMetadata currentMetadata,
         Map<DiscoveryNode, NodeLoad> nodeLoads,
         Map<List<String>, Collection<DiscoveryNode>> mlNodesByZone,
         Optional<CreateTrainedModelAssignmentAction.Request> createAssignmentRequest,
-        int allocatedProcessorsScale,
-        boolean useNewMemoryFields
+        int allocatedProcessorsScale
     ) {
         this.currentMetadata = Objects.requireNonNull(currentMetadata);
         this.nodeLoads = Objects.requireNonNull(nodeLoads);
         this.mlNodesByZone = Objects.requireNonNull(mlNodesByZone);
         this.createAssignmentRequest = Objects.requireNonNull(createAssignmentRequest);
         this.allocatedProcessorsScale = allocatedProcessorsScale;
-        this.useNewMemoryFields = useNewMemoryFields;
     }
 
     TrainedModelAssignmentMetadata.Builder rebalance() {
@@ -179,9 +175,8 @@ private AssignmentPlan computePlanForNormalPriorityModels(
                     currentAssignments,
                     assignment.getMaxAssignedAllocations(),
                     assignment.getAdaptiveAllocationsSettings(),
-                    // in the mixed cluster state use old memory fields to avoid unstable assignment plans
-                    useNewMemoryFields ? assignment.getTaskParams().getPerDeploymentMemoryBytes() : 0,
-                    useNewMemoryFields ? assignment.getTaskParams().getPerAllocationMemoryBytes() : 0
+                    assignment.getTaskParams().getPerDeploymentMemoryBytes(),
+                    assignment.getTaskParams().getPerAllocationMemoryBytes()
                 );
             })
             .forEach(planDeployments::add);
@@ -197,8 +192,8 @@ private AssignmentPlan computePlanForNormalPriorityModels(
                     0,
                     createAssignmentRequest.get().getAdaptiveAllocationsSettings(),
                     // in the mixed cluster state use old memory fields to avoid unstable assignment plans
-                    useNewMemoryFields ? taskParams.getPerDeploymentMemoryBytes() : 0,
-                    useNewMemoryFields ? taskParams.getPerAllocationMemoryBytes() : 0
+                    taskParams.getPerDeploymentMemoryBytes(),
+                    taskParams.getPerAllocationMemoryBytes()
                 )
             );
         }
@@ -237,8 +232,8 @@ private AssignmentPlan computePlanForLowPriorityModels(Set<String> assignableNod
                     assignment.getMaxAssignedAllocations(),
                     assignment.getAdaptiveAllocationsSettings(),
                     Priority.LOW,
-                    (useNewMemoryFields == false) ? assignment.getTaskParams().getPerDeploymentMemoryBytes() : 0,
-                    (useNewMemoryFields == false) ? assignment.getTaskParams().getPerAllocationMemoryBytes() : 0
+                    0,
+                    0
                 )
             )
             .forEach(planDeployments::add);
@@ -254,8 +249,8 @@ private AssignmentPlan computePlanForLowPriorityModels(Set<String> assignableNod
                     0,
                     createAssignmentRequest.get().getAdaptiveAllocationsSettings(),
                     Priority.LOW,
-                    (useNewMemoryFields == false) ? taskParams.getPerDeploymentMemoryBytes() : 0,
-                    (useNewMemoryFields == false) ? taskParams.getPerAllocationMemoryBytes() : 0
+                    0,
+                    0
                 )
             );
         }