diff --git a/src/api-service/__app__/onefuzzlib/azure/auto_scale.py b/src/api-service/__app__/onefuzzlib/azure/auto_scale.py index dc90891dd..da6ab44bd 100644 --- a/src/api-service/__app__/onefuzzlib/azure/auto_scale.py +++ b/src/api-service/__app__/onefuzzlib/azure/auto_scale.py @@ -230,6 +230,35 @@ def create_auto_scale_profile( ) +def get_auto_scale_profile(scaleset_id: UUID) -> AutoscaleProfile: + logging.info("Getting scaleset %s for existing auto scale resources" % scaleset_id) + client = get_monitor_client() + resource_group = get_base_resource_group() + + auto_scale_resource = None + + try: + auto_scale_collections = client.autoscale_settings.list_by_resource_group( + resource_group + ) + for auto_scale in auto_scale_collections: + if str(auto_scale.target_resource_uri).endswith(str(scaleset_id)): + auto_scale_resource = auto_scale + auto_scale_profiles = auto_scale_resource.profiles + if len(auto_scale_profiles) != 1: + logging.info( + "Found more than one autoscaling profile for scaleset %s" + % scaleset_id + ) + return auto_scale_profiles[0] + + except (ResourceNotFoundError, CloudError): + return Error( + code=ErrorCode.INVALID_CONFIGURATION, + errors=["Failed to query scaleset %s autoscale resource" % scaleset_id], + ) + + def default_auto_scale_profile(queue_uri: str, scaleset_size: int) -> AutoscaleProfile: return create_auto_scale_profile( queue_uri, 1, scaleset_size, scaleset_size, 1, 10, 1, 5 diff --git a/src/api-service/__app__/onefuzzlib/workers/scalesets.py b/src/api-service/__app__/onefuzzlib/workers/scalesets.py index 4ab0ece17..56ae4b64a 100644 --- a/src/api-service/__app__/onefuzzlib/workers/scalesets.py +++ b/src/api-service/__app__/onefuzzlib/workers/scalesets.py @@ -35,6 +35,7 @@ from ..azure.auto_scale import ( add_auto_scale_to_vmss, create_auto_scale_profile, default_auto_scale_profile, + get_auto_scale_profile, get_auto_scale_settings, shutdown_scaleset_rule, update_auto_scale, @@ -800,6 +801,15 @@ class Scaleset(BASE_SCALESET, ORMMixin): self.scaleset_id, ) self.delete() + autoscale_entry = AutoScale.get_settings_for_scaleset(self.scaleset_id) + if not autoscale_entry: + logging.info( + "Could not find any auto scale settings for scaleset %s" + % self.scaleset_id + ) + return None + logging.info("Deleting autoscale entry for scaleset %s" % self.scaleset_id) + autoscale_entry.delete() else: self.save() @@ -936,6 +946,75 @@ class Scaleset(BASE_SCALESET, ORMMixin): ) ) + def sync_auto_scale_settings(self) -> Optional[Error]: + from .pools import Pool + + # No need to update tables when in shutdown state + if self.state == ScalesetState.shutdown: + return None + + logging.info( + "Trying to sync auto scale settings for scaleset %s" % self.scaleset_id + ) + + auto_scale_profile = get_auto_scale_profile(self.scaleset_id) + if auto_scale_profile is None: + auto_scale_profile_failed = Error( + code=ErrorCode.UNABLE_TO_FIND, + errors=[ + "Failed to get auto scale profile for scaleset %s" + % self.scaleset_id + ], + ) + logging.error(auto_scale_profile_failed) + return auto_scale_profile_failed + + minimum = auto_scale_profile.capacity.minimum + maximum = auto_scale_profile.capacity.maximum + default = auto_scale_profile.capacity.default + + scale_out_amount = 1 + scale_out_cooldown = 10 + scale_in_amount = 1 + scale_in_cooldown = 15 + + for rule in auto_scale_profile.rules: + scale_action = rule.scale_action + if scale_action.direction == "Increase": + scale_out_amount = scale_action.value + logging.info("Number of seconds: %d" % scale_out_cooldown) + scale_out_cooldown = ( + int((scale_action.cooldown).total_seconds() / 60) % 60 + ) + logging.info("Number of seconds: %d" % scale_out_cooldown) + elif scale_action.direction == "Decrease": + scale_in_amount = scale_action.value + scale_in_cooldown = ( + int((scale_action.cooldown).total_seconds() / 60) % 60 + ) + else: + pass + + pool = Pool.get_by_name(self.pool_name) + if isinstance(pool, Error): + logging.error( + "Failed to get pool by name: %s error: %s" % (self.pool_name, pool) + ) + return pool + + logging.info("Updating auto scale entry: %s" % self.scaleset_id) + AutoScale.update( + scaleset_id=self.scaleset_id, + min=minimum, + max=maximum, + default=default, + scale_out_amount=scale_out_amount, + scale_out_cooldown=scale_out_cooldown, + scale_in_amount=scale_in_amount, + scale_in_cooldown=scale_in_cooldown, + ) + return None + def try_to_enable_auto_scaling(self) -> Optional[Error]: from .pools import Pool @@ -1005,6 +1084,44 @@ class AutoScale(BASE_AUTOSCALE, ORMMixin): entry.save() return entry + @classmethod + def update( + cls, + *, + scaleset_id: UUID, + min: int, + max: int, + default: int, + scale_out_amount: int, + scale_out_cooldown: int, + scale_in_amount: int, + scale_in_cooldown: int, + ) -> None: + + autoscale = cls.search(query={"scaleset_id": [scaleset_id]}) + if not autoscale: + logging.info( + "Could not find any auto scale settings for scaleset %s" % scaleset_id + ) + return None + if len(autoscale) != 1: + logging.info( + "Found more than one autoscaling setting for scaleset %s" % scaleset_id + ) + autoscale[0].scaleset_id = scaleset_id + autoscale[0].min = min + autoscale[0].max = max + autoscale[0].default = default + autoscale[0].scale_out_amount = scale_out_amount + autoscale[0].scale_out_cooldown = scale_out_cooldown + autoscale[0].scale_in_amount = scale_in_amount + autoscale[0].scale_in_cooldown = scale_in_cooldown + + autoscale[0].save() + + def delete(self) -> None: + super().delete() + @classmethod def get_settings_for_scaleset(cls, scaleset_id: UUID) -> Union["AutoScale", None]: autoscale = cls.search(query={"scaleset_id": [scaleset_id]}) diff --git a/src/api-service/__app__/timer_workers/__init__.py b/src/api-service/__app__/timer_workers/__init__.py index 9178f2838..9f8813838 100644 --- a/src/api-service/__app__/timer_workers/__init__.py +++ b/src/api-service/__app__/timer_workers/__init__.py @@ -19,6 +19,7 @@ def process_scaleset(scaleset: Scaleset) -> None: logging.debug("checking scaleset for updates: %s", scaleset.scaleset_id) scaleset.update_configs() + scaleset.sync_auto_scale_settings() # if the scaleset is touched during cleanup, don't continue to process it if scaleset.cleanup_nodes():