diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/TrainedModelAssignment.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/TrainedModelAssignment.java index e8621a68e5c22..9b4ffb7a95120 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/TrainedModelAssignment.java +++ b/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 0L ? StartTrainedModelDeploymentAction.estimateMemoryUsageBytes( model.getModelId(), totalDefinitionLength, - useNewMemoryFields ? model.getPerDeploymentMemoryBytes() : 0, - useNewMemoryFields ? model.getPerAllocationMemoryBytes() : 0, + model.getPerDeploymentMemoryBytes(), + model.getPerAllocationMemoryBytes(), numberOfAllocations ) : 0L; diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java index ae6a8934c3d62..b359daebd5368 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java +++ b/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 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 shuttingDownNodeIds = currentState.metadata().nodeShutdowns().getAllNodeIds(); diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java index 8804d588988b2..20cbdaecaa222 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java @@ -54,22 +54,18 @@ class TrainedModelAssignmentRebalancer { private final Optional createAssignmentRequest; private final int allocatedProcessorsScale; - private final boolean useNewMemoryFields; - TrainedModelAssignmentRebalancer( TrainedModelAssignmentMetadata currentMetadata, Map nodeLoads, Map, Collection> mlNodesByZone, Optional 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 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 assignableNod 0, createAssignmentRequest.get().getAdaptiveAllocationsSettings(), Priority.LOW, - (useNewMemoryFields == false) ? taskParams.getPerDeploymentMemoryBytes() : 0, - (useNewMemoryFields == false) ? taskParams.getPerAllocationMemoryBytes() : 0 + 0, + 0 ) ); } diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancerTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancerTests.java index 65a974e04045e..a2b321a332ae1 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancerTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancerTests.java @@ -46,8 +46,7 @@ public void testRebalance_GivenNoAssignments() { Map.of(), Map.of(), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments().isEmpty(), is(true)); } @@ -82,8 +81,7 @@ public void testRebalance_GivenAllAssignmentsAreSatisfied_ShouldMakeNoChanges() nodeLoads, Map.of(), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(currentMetadata, equalTo(result)); @@ -122,8 +120,7 @@ public void testRebalance_GivenAllAssignmentsAreSatisfied_GivenOutdatedRoutingEn nodeLoads, Map.of(List.of(), List.of(node1, node2)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments(), is(aMapWithSize(2))); @@ -152,8 +149,7 @@ public void testRebalance_GivenModelToAddAlreadyExists() { Map.of(), Map.of(), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance() ); } @@ -168,8 +164,7 @@ public void testRebalance_GivenFirstModelToAdd_NoMLNodes() throws Exception { Map.of(), Map.of(), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -196,8 +191,7 @@ public void testRebalance_GivenFirstModelToAdd_NotEnoughProcessors() throws Exce nodeLoads, Map.of(List.of(), List.of(node)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -233,8 +227,7 @@ public void testRebalance_GivenFirstModelToAdd_NotEnoughMemory() throws Exceptio nodeLoads, Map.of(), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -270,8 +263,7 @@ public void testRebalance_GivenFirstModelToAdd_ErrorDetectingNodeLoad() throws E nodeLoads, Map.of(), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -307,8 +299,7 @@ public void testRebalance_GivenProblemsOnMultipleNodes() throws Exception { nodeLoads, Map.of(List.of(), List.of(node1, node2)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -341,8 +332,7 @@ public void testRebalance_GivenFirstModelToAdd_FitsFully() throws Exception { nodeLoads, Map.of(List.of(), List.of(node1)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -381,8 +371,7 @@ public void testRebalance_GivenModelToAdd_AndPreviousAssignments_AndTwoNodes_All nodeLoads, Map.of(List.of(), List.of(node1, node2)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments(), is(aMapWithSize(2))); @@ -446,8 +435,7 @@ public void testRebalance_GivenPreviousAssignments_AndNewNode() throws Exception nodeLoads, Map.of(List.of(), List.of(node1, node2, node3)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments(), is(aMapWithSize(2))); @@ -511,8 +499,7 @@ public void testRebalance_GivenPreviousAssignments_AndRemovedNode_AndRemainingNo nodeLoads, Map.of(List.of(), List.of(node1)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments(), is(aMapWithSize(2))); @@ -582,8 +569,7 @@ public void testRebalance_GivenPreviousAssignments_AndRemovedNode_AndRemainingNo nodeLoads, Map.of(List.of(), List.of(node1)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments(), is(aMapWithSize(2))); @@ -632,8 +618,7 @@ public void testRebalance_GivenFailedAssignment_RestartsAssignment() throws Exce nodeLoads, Map.of(List.of(), List.of(node1)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); assertThat(result.allAssignments(), is(aMapWithSize(1))); @@ -667,8 +652,7 @@ public void testRebalance_GivenLowPriorityModelToAdd_OnlyModel_NotEnoughMemory() nodeLoads, Map.of(), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(deploymentId); @@ -714,8 +698,7 @@ public void testRebalance_GivenLowPriorityModelToAdd_NotEnoughMemoryNorProcessor nodeLoads, Map.of(List.of("zone-1"), List.of(node1), List.of("zone-2"), List.of(node2)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams1, null)), - 1, - false + 1 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(deployment1); @@ -754,8 +737,7 @@ public void testRebalance_GivenMixedPriorityModels_NotEnoughMemoryForLowPriority nodeLoads, Map.of(List.of(), List.of(node1)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); { @@ -809,8 +791,7 @@ public void testRebalance_GivenMixedPriorityModels_TwoZones_EachNodeCanHoldOneMo nodeLoads, Map.of(List.of("zone-1"), List.of(node1), List.of("zone-2"), List.of(node2)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); List assignedNodes = new ArrayList<>(); @@ -864,8 +845,7 @@ public void testRebalance_GivenModelUsingAllCpu_FittingLowPriorityModelCanStart( nodeLoads, Map.of(List.of(), List.of(node1)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); { @@ -915,8 +895,7 @@ public void testRebalance_GivenMultipleLowPriorityModels_AndMultipleNodes() thro nodeLoads, Map.of(List.of(), List.of(node1, node2)), Optional.empty(), - 1, - false + 1 ).rebalance().build(); { @@ -967,8 +946,7 @@ public void testRebalance_GivenNormalPriorityModelToLoad_EvictsLowPriorityModel( nodeLoads, Map.of(List.of(), List.of(node1)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams2, null)), - 1, - false + 1 ).rebalance().build(); { @@ -1021,8 +999,7 @@ public void testRebalance_GivenNormalPriorityModelToLoad_AndLowPriorityModelCanS nodeLoads, Map.of(List.of(), List.of(node1, node2)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams2, null)), - 1, - false + 1 ).rebalance().build(); { @@ -1075,8 +1052,7 @@ public void testRebalance_GivenNormalPriorityModelToLoad_AndLowPriorityModelMust nodeLoads, Map.of(List.of(), List.of(node1, node2)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams2, null)), - 1, - false + 1 ).rebalance().build(); { @@ -1122,8 +1098,7 @@ public void testRebalance_GivenFirstModelToAdd_GivenScalingProcessorSetting() { nodeLoads, Map.of(List.of(), List.of(node)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 2, - false + 2 ).rebalance().build(); TrainedModelAssignment assignment = result.getDeploymentAssignment(modelId); @@ -1145,8 +1120,7 @@ public void testRebalance_GivenFirstModelToAdd_GivenScalingProcessorSetting() { nodeLoads, Map.of(List.of(), List.of(node)), Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)), - 1, - false + 1 ).rebalance().build(); assignment = result.getDeploymentAssignment(modelId);