ClusterAgentAutoScaler { @VisibleForTesting Completable doAgentScaling() { return Completable.defer(() -> { if (!configuration.isAutoScalingAgentsEnabled()) { logger.debug("Auto scaling agents is not enabled"); return Completable.complete(); } List<Completable> actions = new ArrayList<>(); Map<String, Job> allJobs = getAllJobs(); Map<String, Task> allTasks = getAllTasks(); List<AgentInstanceGroup> activeInstanceGroups = getActiveInstanceGroups(); Map<AgentInstanceGroup, List<AgentInstance>> instancesForActiveInstanceGroups = getInstancesForInstanceGroups(activeInstanceGroups); Map<String, List<AgentInstance>> instancesForActiveInstanceGroupsById = instancesForActiveInstanceGroups.entrySet().stream() .collect(Collectors.toMap(e -> e.getKey().getId(), Map.Entry::getValue)); Map<String, Long> numberOfTasksOnAgents = getNumberOfTasksOnAgents(allTasks.values()); Map<Tier, Set<String>> failedTaskIdsByTier = getFailedTaskIdsByTier( schedulingService.getLastTaskPlacementFailures(), FailureKind.NEVER_TRIGGER_AUTOSCALING ); long now = clock.wallTime(); for (Tier tier : Tier.values()) { logger.info("Starting scaling actions for tier: {}", tier); TierAutoScalingConfiguration tierConfiguration = ClusterOperationFunctions.getTierConfiguration(tier, configuration); TierAutoScalerExecution tierAutoScalerExecution = tierTierAutoScalerExecutions.computeIfAbsent( tier, k -> new TierAutoScalerExecution(tier, titusRuntime.getRegistry()) ); String primaryInstanceType = tierConfiguration.getPrimaryInstanceType(); ResourceDimension tierResourceDimension = agentManagementService.getResourceLimits(primaryInstanceType); List<AgentInstanceGroup> activeScalableInstanceGroupsForTier = activeInstanceGroups.stream() .filter(ig -> ig.getTier() == tier && ig.getInstanceType().equals(primaryInstanceType)) .collect(Collectors.toList()); logger.info("{} active instance groups({}): {}", tier, activeScalableInstanceGroupsForTier.size(), activeScalableInstanceGroupsForTier); List<AgentInstance> idleInstancesForTier = getIdleInstancesForTier(tier, primaryInstanceType, instancesForActiveInstanceGroups, numberOfTasksOnAgents, now, tierConfiguration.getIdleInstanceGracePeriodMs()); tierAutoScalerExecution.getTotalIdleInstancesGauge().set(idleInstancesForTier.size()); logger.info("{} idle instances({}): {}", tier, idleInstancesForTier.size(), idleInstancesForTier); Set<String> failedTaskIds = failedTaskIdsByTier.getOrDefault(tier, Collections.emptySet()); tierAutoScalerExecution.getTotalFailedTasksGauge().set(failedTaskIds.size()); logger.info("{} failed tasks({}): {}", tier, failedTaskIds.size(), failedTaskIds); int agentCountToScaleUp = 0; Set<String> potentialTaskIdsForScaleUp = new HashSet<>(); boolean usedScaleUpCooldown = false; if (hasTimeElapsed(tierAutoScalerExecution.getLastScaleUp().get(), now, tierConfiguration.getScaleUpCoolDownMs())) { int minIdleForTier = tierConfiguration.getMinIdle(); if (idleInstancesForTier.size() < minIdleForTier) { int instancesNeededForMinIdle = minIdleForTier - idleInstancesForTier.size(); logger.info("{} needs {} instances to satisfy min idle {}", tier, instancesNeededForMinIdle, minIdleForTier); agentCountToScaleUp += instancesNeededForMinIdle; } Set<String> placementFailureTaskIds = getFailedTaskIdsByTier( schedulingService.getLastTaskPlacementFailures(), IGNORED_FAILURE_KINDS_WITH_LAUNCHGUARD ).getOrDefault(tier, Collections.emptySet()); logger.info("{} had the placement excluding launch guard failures({}): {}", tier, placementFailureTaskIds.size(), placementFailureTaskIds); Set<String> scalablePlacementFailureTaskIds = filterOutTaskIdsForScaling(placementFailureTaskIds, allJobs, allTasks, tierResourceDimension); logger.info("{} had the scalable placement failures({}): {}", tier, scalablePlacementFailureTaskIds.size(), scalablePlacementFailureTaskIds); potentialTaskIdsForScaleUp.addAll(scalablePlacementFailureTaskIds); if (agentCountToScaleUp > 0 || !scalablePlacementFailureTaskIds.isEmpty()) { usedScaleUpCooldown = true; } } Set<String> tasksPastSlo = getTasksPastSlo(failedTaskIds, allTasks, now, tierConfiguration.getTaskSloMs()); Set<String> scalableTasksPastSlo = filterOutTaskIdsForScaling(tasksPastSlo, allJobs, allTasks, tierResourceDimension); tierAutoScalerExecution.getTotalTasksPastSloGauge().set(scalableTasksPastSlo.size()); logger.info("{} had tasks past slo({}): {}", tier, scalableTasksPastSlo.size(), scalableTasksPastSlo); potentialTaskIdsForScaleUp.addAll(scalableTasksPastSlo); Set<String> taskIdsForScaleUp = new HashSet<>(); for (String taskId : potentialTaskIdsForScaleUp) { boolean previouslyScaledFor = taskIdsForPreviousScaleUps.getIfPresent(taskId) != null; if (!previouslyScaledFor) { taskIdsForScaleUp.add(taskId); taskIdsForPreviousScaleUps.put(taskId, taskId); } } tierAutoScalerExecution.getTotalTasksForScaleUpGauge().set(taskIdsForScaleUp.size()); logger.info("{} had tasks to scale up({}): {}", tier, taskIdsForScaleUp.size(), taskIdsForScaleUp); int agentScaleUpCountByDominantResource = calculateAgentScaleUpCountByDominantResource(taskIdsForScaleUp, allJobs, allTasks, tierResourceDimension); logger.info("{} needs {} instances based on dominant resource", tier, agentScaleUpCountByDominantResource); agentCountToScaleUp += agentScaleUpCountByDominantResource; logger.info("{} needs {} instances", tier, agentCountToScaleUp); tierAutoScalerExecution.getTotalAgentsToScaleUpGauge().set(agentCountToScaleUp); boolean scalingUp = false; if (agentCountToScaleUp > 0) { long maxTokensToTake = Math.min(SCALE_UP_TOKEN_BUCKET_CAPACITY, agentCountToScaleUp); Optional<Pair<Long, ImmutableTokenBucket>> takeOpt = tierAutoScalerExecution.getLastScaleUpTokenBucket().tryTake(1, maxTokensToTake); if (takeOpt.isPresent()) { Pair<Long, ImmutableTokenBucket> takePair = takeOpt.get(); tierAutoScalerExecution.setLastScaleUpTokenBucket(takePair.getRight()); long tokensAvailable = takePair.getLeft(); Pair<Integer, Completable> scaleUpPair = createScaleUpCompletable(activeScalableInstanceGroupsForTier, (int) tokensAvailable); Integer agentCountBeingScaled = scaleUpPair.getLeft(); tierAutoScalerExecution.getTotalAgentsBeingScaledUpGauge().set(agentCountBeingScaled); if (agentCountBeingScaled > 0) { actions.add(scaleUpPair.getRight()); logger.info("Attempting to scale up {} tier by {} agent instances", tier, agentCountBeingScaled); scalingUp = true; if (usedScaleUpCooldown) { tierAutoScalerExecution.getLastScaleUp().set(clock.wallTime()); } } } } if (!scalingUp && hasTimeElapsed(tierAutoScalerExecution.getLastScaleDown().get(), now, tierConfiguration.getScaleDownCoolDownMs())) { int agentCountToScaleDown = 0; int maxIdleForTier = tierConfiguration.getMaxIdle(); if (idleInstancesForTier.size() > maxIdleForTier) { int instancesNotNeededForMaxIdle = idleInstancesForTier.size() - maxIdleForTier; logger.info("{} can remove {} instances to satisfy max idle {}", tier, instancesNotNeededForMaxIdle, maxIdleForTier); agentCountToScaleDown += instancesNotNeededForMaxIdle; } tierAutoScalerExecution.getTotalAgentsToScaleDownGauge().set(agentCountToScaleDown); if (agentCountToScaleDown > 0) { long maxTokensToTake = Math.min(SCALE_DOWN_TOKEN_BUCKET_CAPACITY, agentCountToScaleDown); Optional<Pair<Long, ImmutableTokenBucket>> takeOpt = tierAutoScalerExecution.getLastScaleDownTokenBucket().tryTake(1, maxTokensToTake); if (takeOpt.isPresent()) { Pair<Long, ImmutableTokenBucket> takePair = takeOpt.get(); tierAutoScalerExecution.setLastScaleDownTokenBucket(takePair.getRight()); long tokensAvailable = takePair.getLeft(); Pair<Integer, Completable> scaleDownPair = createSetRemovableOverrideStatusesCompletable(idleInstancesForTier, activeScalableInstanceGroupsForTier, instancesForActiveInstanceGroupsById, (int) tokensAvailable); Integer agentCountBeingScaledDown = scaleDownPair.getLeft(); tierAutoScalerExecution.getTotalAgentsBeingScaledDownGauge().set(agentCountBeingScaledDown); if (agentCountBeingScaledDown > 0) { actions.add(scaleDownPair.getRight()); logger.info("Attempting to scale down {} tier by {} agent instances", tier, agentCountBeingScaledDown); tierAutoScalerExecution.getLastScaleDown().set(clock.wallTime()); } } } } logger.info("Finishing scaling actions for tier: {}", tier); } List<AgentInstance> removableInstancesPastElapsedTime = getRemovableInstancesPastElapsedTime(instancesForActiveInstanceGroups, now, configuration.getAgentInstanceRemovableTimeoutMs()); logger.info("Removable instances past elapsed time({}): {}", removableInstancesPastElapsedTime.size(), removableInstancesPastElapsedTime); if (!removableInstancesPastElapsedTime.isEmpty()) { actions.add(createResetOverrideStatusesCompletable(removableInstancesPastElapsedTime)); logger.info("Resetting agent instances({}): {}", removableInstancesPastElapsedTime.size(), removableInstancesPastElapsedTime); } return actions.isEmpty() ? Completable.complete() : Completable.concat(actions); }).doOnCompleted(() -> logger.debug("Completed scaling agents")) .timeout(CLUSTER_AGENT_AUTO_SCALE_COMPLETABLE_TIMEOUT_MS, TimeUnit.MILLISECONDS); } @Inject  ClusterAgentAutoScaler(TitusRuntime titusRuntime,
                                  ClusterOperationsConfiguration configuration,
                                  AgentManagementService agentManagementService,
                                  V3JobOperations v3JobOperations,
                                  SchedulingService<? extends TaskRequest> schedulingService);  ClusterAgentAutoScaler(TitusRuntime titusRuntime,
                                  ClusterOperationsConfiguration configuration,
                                  AgentManagementService agentManagementService,
                                  V3JobOperations v3JobOperations,
                                  SchedulingService<? extends TaskRequest> schedulingService,
                                  Scheduler scheduler); @Activator void enterActiveMode(); @PreDestroy void shutdown();  }
@Test public void testOnlyFenzoPartitionIsScaled() { List<AgentInstanceGroup> instanceGroups = Arrays.asList( createPartition("fenzoPartition1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 1), createPartition("kubeSchedulerPartition", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 1), createPartition("fenzoPartition2", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 1) ); when(agentManagementService.getInstanceGroups()).thenReturn(instanceGroups); when(agentManagementService.isOwnedByFenzo(any(AgentInstanceGroup.class))).thenAnswer(invocation -> { AgentInstanceGroup instanceGroup = invocation.getArgument(0); return instanceGroup.getId().contains("fenzo"); }); when(agentManagementService.scaleUp(eq("fenzoPartition1"), anyInt())).thenReturn(Completable.complete()); when(agentManagementService.scaleUp(eq("fenzoPartition2"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService).scaleUp("fenzoPartition1", 1); verify(agentManagementService, times(0)).scaleUp("kubeSchedulerPartition", 1); verify(agentManagementService).scaleUp("fenzoPartition2", 1); }
@Test public void testScaleUpMinIdle() { AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 10); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService).scaleUp("instanceGroup1", 5); }
@Test public void testScaleDownMaxIdle() { AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 12, 20); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); List<AgentInstance> agentInstances = createAgents(12, "instanceGroup1", false); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(agentInstances); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(2)).updateAgentInstanceAttributes(any(), any()); }
@Test public void testScaleUpForFailingTasks() { when(configuration.getFlexMinIdle()).thenReturn(0); Job job = createJob(); when(v3JobOperations.getJobs()).thenReturn(Collections.singletonList(job)); List<Task> tasks = createTasks(10, "jobId"); when(v3JobOperations.getTasks()).thenReturn(tasks); Map<FailureKind, Map<StubTaskRequest, List<TaskPlacementFailure>>> taskPlacementFailures = createTaskPlacementFailures(ImmutableMap.of( AllAgentsFull, 10 ), Tier.Flex); doReturn(taskPlacementFailures).when(schedulingService).getLastTaskPlacementFailures(); AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 20); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService).scaleUp("instanceGroup1", 10); }
@Test public void testNoScaleUpForSomeFailureKinds() { when(configuration.getFlexMinIdle()).thenReturn(0); Job job = createJob(); when(v3JobOperations.getJobs()).thenReturn(Collections.singletonList(job)); List<Task> tasks = createTasks(10, "jobId"); when(v3JobOperations.getTasks()).thenReturn(tasks); Map<FailureKind, Map<StubTaskRequest, List<TaskPlacementFailure>>> taskPlacementFailures = createTaskPlacementFailures(ImmutableMap.of( WaitingForInUseIpAllocation, 2, OpportunisticResource, 8 ), Tier.Flex); doReturn(taskPlacementFailures).when(schedulingService).getLastTaskPlacementFailures(); AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 20); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, never()).scaleUp(anyString(), anyInt()); }
@Test public void testScaleUpForTasksPastSlo() { when(configuration.getFlexMinIdle()).thenReturn(0); when(configuration.getFlexScaleUpCoolDownMs()).thenReturn(72000000L); when(configuration.getFlexTaskSloMs()).thenReturn(3600000L); AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 10); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); Job job = createJob(); when(v3JobOperations.getJobs()).thenReturn(Collections.singletonList(job)); List<Task> tasks = createTasks(10, "jobId"); when(v3JobOperations.getTasks()).thenReturn(tasks); Map<FailureKind, Map<StubTaskRequest, List<TaskPlacementFailure>>> taskPlacementFailures = createTaskPlacementFailures(ImmutableMap.of( AllAgentsFull, 10 ), Tier.Flex); doReturn(taskPlacementFailures).when(schedulingService).getLastTaskPlacementFailures(); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(1, TimeUnit.HOURS); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService).scaleUp("instanceGroup1", 10); }
@Test public void testNoScaleUpForTasksPastSloWithSomeFailureKinds() { when(configuration.getFlexMinIdle()).thenReturn(0); when(configuration.getFlexScaleUpCoolDownMs()).thenReturn(72000000L); when(configuration.getFlexTaskSloMs()).thenReturn(3600000L); AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 10); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); Job job = createJob(); when(v3JobOperations.getJobs()).thenReturn(Collections.singletonList(job)); List<Task> tasks = createTasks(10, "jobId"); when(v3JobOperations.getTasks()).thenReturn(tasks); Map<FailureKind, Map<StubTaskRequest, List<TaskPlacementFailure>>> taskPlacementFailures = createTaskPlacementFailures(ImmutableMap.of( WaitingForInUseIpAllocation, 2, OpportunisticResource, 8 ), Tier.Flex); doReturn(taskPlacementFailures).when(schedulingService).getLastTaskPlacementFailures(); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(1, TimeUnit.HOURS); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, never()).scaleUp(anyString(), anyInt()); }
@Test public void testScaleUpForNonPrimaryInstanceType() { AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "m4.16xlarge", 0, 0, 10); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(0)).scaleUp(anyString(), anyInt()); }
@Test public void testScaleUpForActiveAndPhasedOutInstanceGroups() { AgentInstanceGroup activeInstanceGroup = createPartition("activeInstanceGroup", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 0, 10); AgentInstanceGroup phasedOutInstanceGroup = createPartition("phasedOutInstanceGroup", InstanceGroupLifecycleState.PhasedOut, "r4.16xlarge", 0, 0, 10); List<AgentInstanceGroup> instanceGroups = Lists.newArrayList(phasedOutInstanceGroup, activeInstanceGroup); when(agentManagementService.getInstanceGroups()).thenReturn(instanceGroups); when(agentManagementService.getAgentInstances("activeInstanceGroup")).thenReturn(Collections.emptyList()); when(agentManagementService.getAgentInstances("phasedOutInstanceGroup")).thenReturn(Collections.emptyList()); when(agentManagementService.scaleUp(eq("activeInstanceGroup"), anyInt())).thenReturn(Completable.complete()); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(1)).scaleUp(eq("activeInstanceGroup"), eq(5)); verify(agentManagementService, times(0)).scaleUp(eq("phasedOutInstanceGroup"), anyInt()); }
@Test public void testScaleDownForNonPrimaryInstanceType() { AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "m4.16xlarge", 0, 12, 20); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); List<AgentInstance> agentInstances = createAgents(12, "instanceGroup1", false); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(agentInstances); testScheduler.advanceTimeBy(6, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(0)).updateAgentInstanceAttributes(any(), any()); }
@Test public void testScaleDownForIdleAgents() { when(configuration.getFlexMinIdle()).thenReturn(0); when(configuration.getFlexMaxIdle()).thenReturn(0); AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 20, 20); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); List<AgentInstance> agentInstances = createAgents(20, "instanceGroup1", false); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(agentInstances); List<Task> tasks = createTasks(17, "jobId"); for (int i = 0; i < tasks.size(); i++) { Task task = tasks.get(i); when(task.getTaskContext()).thenReturn(singletonMap(TaskAttributes.TASK_ATTRIBUTES_AGENT_INSTANCE_ID, "instanceGroup1" + i)); } when(v3JobOperations.getTasks()).thenReturn(tasks); testScheduler.advanceTimeBy(11, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(3)).updateAgentInstanceAttributes(any(), any()); }
@Test public void testScaleDownForIdleAgentsPrefersPhasedOutInstanceGroups() { when(configuration.getFlexMinIdle()).thenReturn(0); when(configuration.getFlexMaxIdle()).thenReturn(5); AgentInstanceGroup activeInstanceGroup = createPartition("activeInstanceGroup", InstanceGroupLifecycleState.Active, "r4.16xlarge", 0, 5, 10); AgentInstanceGroup phasedOutInstanceGroup = createPartition("phasedOutInstanceGroup", InstanceGroupLifecycleState.PhasedOut, "r4.16xlarge", 0, 5, 10); List<AgentInstanceGroup> instanceGroups = Lists.newArrayList(phasedOutInstanceGroup, activeInstanceGroup); when(agentManagementService.getInstanceGroups()).thenReturn(instanceGroups); List<AgentInstance> activeInstances = createAgents(5, "activeInstanceGroup", false); when(agentManagementService.getAgentInstances("activeInstanceGroup")).thenReturn(activeInstances); List<AgentInstance> phasedOutInstances = createAgents(5, "phasedOutInstanceGroup", false); when(agentManagementService.getAgentInstances("phasedOutInstanceGroup")).thenReturn(phasedOutInstances); when(v3JobOperations.getTasks()).thenReturn(Collections.emptyList()); testScheduler.advanceTimeBy(11, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(5)).updateAgentInstanceAttributes(matches("phasedOutInstanceGroup.*"), any()); verify(agentManagementService, times(0)).updateAgentInstanceAttributes(matches("activeInstanceGroup.*"), any()); }
@Test public void testScaleDownForIdleAgentsDoesNotGoPastInstanceGroupMin() { when(configuration.getFlexMinIdle()).thenReturn(0); when(configuration.getAgentInstanceRemovableTimeoutMs()).thenReturn(9999999999L); AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "r4.16xlarge", 13, 20, 20); when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup)); List<AgentInstance> agentInstances = CollectionsExt.merge( createAgents(18, "instanceGroup1", false), createAgents(2, "instanceGroup1", true) ); when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(agentInstances); when(v3JobOperations.getTasks()).thenReturn(emptyList()); testScheduler.advanceTimeBy(11, TimeUnit.MINUTES); ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration, agentManagementService, v3JobOperations, schedulingService, testScheduler); clusterAgentAutoScaler.doAgentScaling().await(); verify(agentManagementService, times(5)).updateAgentInstanceAttributes(any(), any()); }