@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()); }
@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); }
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); } }
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); }
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(); }
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();  }