Autoscale Table Synces Current Scaleset Settings (#2255)

* Autoscale Table Syncs Current Scaleset Settings.

* Fixing formatting.

* Fixing lint.

* Remove white space.

* Adding default check.

* Formatting.

* Fix.

* Fixing import.

* Fixing ref.

* Trying regular save.

* using self.

* Trying to fix lint.

* Fixing args.

* Fixed.

* Adding function call.

* Reformat.

* Formatting.

* trying again.

* FIxing profile call.

* Fixing type.

* Fixing type.

* Fixing format.

* Fixing typing.

* Fixing.

* Fixing function call

* Getting rid of length check.

* Getting rid of length check.

* Ready for review.

* Formatting.

* Removing bad change.

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

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]})

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():