|  | @@ -29,7 +29,6 @@ import org.elasticsearch.xpack.ml.MachineLearning;
 | 
	
		
			
				|  |  |  import org.elasticsearch.xpack.ml.inference.assignment.TrainedModelAssignmentMetadata;
 | 
	
		
			
				|  |  |  import org.junit.Before;
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -import java.util.OptionalInt;
 | 
	
		
			
				|  |  |  import java.util.Set;
 | 
	
		
			
				|  |  |  import java.util.function.LongSupplier;
 | 
	
		
			
				|  |  |  
 | 
	
	
		
			
				|  | @@ -41,13 +40,10 @@ import static org.mockito.Mockito.when;
 | 
	
		
			
				|  |  |  public class MlProcessorAutoscalingDeciderTests extends ESTestCase {
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |      private ScaleTimer scaleTimer;
 | 
	
		
			
				|  |  | -    private NodeAvailabilityZoneMapper nodeAvailabilityZoneMapper;
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |      @Before
 | 
	
		
			
				|  |  |      public void setup() {
 | 
	
		
			
				|  |  |          scaleTimer = new ScaleTimer(System::currentTimeMillis);
 | 
	
		
			
				|  |  | -        nodeAvailabilityZoneMapper = mock(NodeAvailabilityZoneMapper.class);
 | 
	
		
			
				|  |  | -        when(nodeAvailabilityZoneMapper.getNumMlAvailabilityZones()).thenReturn(OptionalInt.empty());
 | 
	
		
			
				|  |  |      }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |      public void testScale_GivenCurrentCapacityIsUsedExactly() {
 | 
	
	
		
			
				|  | @@ -152,112 +148,6 @@ public class MlProcessorAutoscalingDeciderTests extends ESTestCase {
 | 
	
		
			
				|  |  |          assertThat(capacity.reason(), equalTo("requesting scale up as there are unsatisfied deployments"));
 | 
	
		
			
				|  |  |      }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -    public void testScale_GivenUnsatisfiedDeployments_AndThreeMlAvailabilityZones_AndNodeProcessorsMoreThanTierProcessors() {
 | 
	
		
			
				|  |  | -        givenMlAvailabilityZones(3);
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        String modelId1 = "model-id-1";
 | 
	
		
			
				|  |  | -        String modelId2 = "model-id-2";
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        String mlNodeId1 = "ml-node-id-1";
 | 
	
		
			
				|  |  | -        String mlNodeId2 = "ml-node-id-2";
 | 
	
		
			
				|  |  | -        String dataNodeId = "data-node-id";
 | 
	
		
			
				|  |  | -        DiscoveryNode mlNode1 = buildNode(mlNodeId1, true, 4);
 | 
	
		
			
				|  |  | -        DiscoveryNode mlNode2 = buildNode(mlNodeId2, true, 4);
 | 
	
		
			
				|  |  | -        DiscoveryNode dataNode = buildNode(dataNodeId, false, 24);
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        ClusterState clusterState = ClusterState.builder(new ClusterName("test"))
 | 
	
		
			
				|  |  | -            .nodes(DiscoveryNodes.builder().add(mlNode1).add(mlNode2).add(dataNode).build())
 | 
	
		
			
				|  |  | -            .metadata(
 | 
	
		
			
				|  |  | -                Metadata.builder()
 | 
	
		
			
				|  |  | -                    .putCustom(
 | 
	
		
			
				|  |  | -                        TrainedModelAssignmentMetadata.NAME,
 | 
	
		
			
				|  |  | -                        TrainedModelAssignmentMetadata.Builder.empty()
 | 
	
		
			
				|  |  | -                            .addNewAssignment(
 | 
	
		
			
				|  |  | -                                modelId1,
 | 
	
		
			
				|  |  | -                                TrainedModelAssignment.Builder.empty(
 | 
	
		
			
				|  |  | -                                    new StartTrainedModelDeploymentAction.TaskParams(modelId1, 42L, 8, 1, 1024, ByteSizeValue.ONE)
 | 
	
		
			
				|  |  | -                                )
 | 
	
		
			
				|  |  | -                            )
 | 
	
		
			
				|  |  | -                            .addNewAssignment(
 | 
	
		
			
				|  |  | -                                modelId2,
 | 
	
		
			
				|  |  | -                                TrainedModelAssignment.Builder.empty(
 | 
	
		
			
				|  |  | -                                    new StartTrainedModelDeploymentAction.TaskParams(modelId2, 42L, 4, 3, 1024, ByteSizeValue.ONE)
 | 
	
		
			
				|  |  | -                                )
 | 
	
		
			
				|  |  | -                                    .addRoutingEntry(mlNodeId1, new RoutingInfo(1, 1, RoutingState.STARTED, ""))
 | 
	
		
			
				|  |  | -                                    .addRoutingEntry(mlNodeId2, new RoutingInfo(1, 1, RoutingState.STARTED, ""))
 | 
	
		
			
				|  |  | -                            )
 | 
	
		
			
				|  |  | -                            .build()
 | 
	
		
			
				|  |  | -                    )
 | 
	
		
			
				|  |  | -                    .build()
 | 
	
		
			
				|  |  | -            )
 | 
	
		
			
				|  |  | -            .build();
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        MlProcessorAutoscalingDecider decider = newDecider();
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        MlProcessorAutoscalingCapacity capacity = decider.scale(
 | 
	
		
			
				|  |  | -            Settings.EMPTY,
 | 
	
		
			
				|  |  | -            newContext(clusterState),
 | 
	
		
			
				|  |  | -            new MlAutoscalingContext(clusterState)
 | 
	
		
			
				|  |  | -        );
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        assertThat(capacity.nodeProcessors(), equalTo(Processors.of(8.0)));
 | 
	
		
			
				|  |  | -        assertThat(capacity.tierProcessors(), equalTo(Processors.of(8.0)));
 | 
	
		
			
				|  |  | -        assertThat(capacity.reason(), equalTo("requesting scale up as there are unsatisfied deployments"));
 | 
	
		
			
				|  |  | -    }
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -    public void testScale_GivenUnsatisfiedDeployments_AndThreeMlAvailabilityZones_AndNodeProcessorsLessThanTierProcessors() {
 | 
	
		
			
				|  |  | -        givenMlAvailabilityZones(3);
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        String modelId1 = "model-id-1";
 | 
	
		
			
				|  |  | -        String modelId2 = "model-id-2";
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        String mlNodeId1 = "ml-node-id-1";
 | 
	
		
			
				|  |  | -        String mlNodeId2 = "ml-node-id-2";
 | 
	
		
			
				|  |  | -        String dataNodeId = "data-node-id";
 | 
	
		
			
				|  |  | -        DiscoveryNode mlNode1 = buildNode(mlNodeId1, true, 4);
 | 
	
		
			
				|  |  | -        DiscoveryNode mlNode2 = buildNode(mlNodeId2, true, 4);
 | 
	
		
			
				|  |  | -        DiscoveryNode dataNode = buildNode(dataNodeId, false, 24);
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        ClusterState clusterState = ClusterState.builder(new ClusterName("test"))
 | 
	
		
			
				|  |  | -            .nodes(DiscoveryNodes.builder().add(mlNode1).add(mlNode2).add(dataNode).build())
 | 
	
		
			
				|  |  | -            .metadata(
 | 
	
		
			
				|  |  | -                Metadata.builder()
 | 
	
		
			
				|  |  | -                    .putCustom(
 | 
	
		
			
				|  |  | -                        TrainedModelAssignmentMetadata.NAME,
 | 
	
		
			
				|  |  | -                        TrainedModelAssignmentMetadata.Builder.empty()
 | 
	
		
			
				|  |  | -                            .addNewAssignment(
 | 
	
		
			
				|  |  | -                                modelId1,
 | 
	
		
			
				|  |  | -                                TrainedModelAssignment.Builder.empty(
 | 
	
		
			
				|  |  | -                                    new StartTrainedModelDeploymentAction.TaskParams(modelId1, 42L, 8, 1, 1024, ByteSizeValue.ONE)
 | 
	
		
			
				|  |  | -                                )
 | 
	
		
			
				|  |  | -                            )
 | 
	
		
			
				|  |  | -                            .addNewAssignment(
 | 
	
		
			
				|  |  | -                                modelId2,
 | 
	
		
			
				|  |  | -                                TrainedModelAssignment.Builder.empty(
 | 
	
		
			
				|  |  | -                                    new StartTrainedModelDeploymentAction.TaskParams(modelId2, 42L, 4, 6, 1024, ByteSizeValue.ONE)
 | 
	
		
			
				|  |  | -                                )
 | 
	
		
			
				|  |  | -                                    .addRoutingEntry(mlNodeId1, new RoutingInfo(1, 1, RoutingState.STARTED, ""))
 | 
	
		
			
				|  |  | -                                    .addRoutingEntry(mlNodeId2, new RoutingInfo(1, 1, RoutingState.STARTED, ""))
 | 
	
		
			
				|  |  | -                            )
 | 
	
		
			
				|  |  | -                            .build()
 | 
	
		
			
				|  |  | -                    )
 | 
	
		
			
				|  |  | -                    .build()
 | 
	
		
			
				|  |  | -            )
 | 
	
		
			
				|  |  | -            .build();
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        MlProcessorAutoscalingDecider decider = newDecider();
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        MlProcessorAutoscalingCapacity capacity = decider.scale(
 | 
	
		
			
				|  |  | -            Settings.EMPTY,
 | 
	
		
			
				|  |  | -            newContext(clusterState),
 | 
	
		
			
				|  |  | -            new MlAutoscalingContext(clusterState)
 | 
	
		
			
				|  |  | -        );
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -        assertThat(capacity.nodeProcessors(), equalTo(Processors.of(8.0)));
 | 
	
		
			
				|  |  | -        assertThat(capacity.tierProcessors(), equalTo(Processors.of(11.0)));
 | 
	
		
			
				|  |  | -        assertThat(capacity.reason(), equalTo("requesting scale up as there are unsatisfied deployments"));
 | 
	
		
			
				|  |  | -    }
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  |      public void testScale_GivenMoreThanHalfProcessorsAreUsed() {
 | 
	
		
			
				|  |  |          String modelId1 = "model-id-1";
 | 
	
		
			
				|  |  |          String modelId2 = "model-id-2";
 | 
	
	
		
			
				|  | @@ -429,7 +319,7 @@ public class MlProcessorAutoscalingDeciderTests extends ESTestCase {
 | 
	
		
			
				|  |  |      }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |      private MlProcessorAutoscalingDecider newDecider() {
 | 
	
		
			
				|  |  | -        return new MlProcessorAutoscalingDecider(scaleTimer, nodeAvailabilityZoneMapper);
 | 
	
		
			
				|  |  | +        return new MlProcessorAutoscalingDecider(scaleTimer);
 | 
	
		
			
				|  |  |      }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |      private AutoscalingDeciderContext newContext(ClusterState clusterState) {
 | 
	
	
		
			
				|  | @@ -438,10 +328,6 @@ public class MlProcessorAutoscalingDeciderTests extends ESTestCase {
 | 
	
		
			
				|  |  |          return context;
 | 
	
		
			
				|  |  |      }
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -    private void givenMlAvailabilityZones(int zones) {
 | 
	
		
			
				|  |  | -        when(nodeAvailabilityZoneMapper.getNumMlAvailabilityZones()).thenReturn(OptionalInt.of(zones));
 | 
	
		
			
				|  |  | -    }
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  |      private static class TimeMachine implements LongSupplier {
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |          private long offsetMillis;
 |