Adding auto scale via cli (#1717)

* Initial implementation for adding auto scale via cli

* Remove unused argument

* Remove unused import

* I had a 👻extra line👻

src/api-service/__app__/onefuzzlib/azure/auto_scale.py

@@ -104,6 +104,7 @@ def create_auto_scale_resource_for(
         "location": location,
         "profiles": [profile],
         "target_resource_uri": scaleset_uri,
+        "enabled": True,
     }
 
     try:
@@ -125,7 +126,16 @@ def create_auto_scale_resource_for(
         )
 
 
-def create_auto_scale_profile(min: int, max: int, queue_uri: str) -> AutoscaleProfile:
+def create_auto_scale_profile(
+    queue_uri: str,
+    min: int,
+    max: int,
+    default: int,
+    scale_out_amount: int,
+    scale_out_cooldown: int,
+    scale_in_amount: int,
+    scale_in_cooldown: int,
+) -> AutoscaleProfile:
     return AutoscaleProfile(
         name=str(uuid.uuid4()),
         capacity=ScaleCapacity(minimum=min, maximum=max, default=max),
@@ -150,8 +160,8 @@ def create_auto_scale_profile(min: int, max: int, queue_uri: str) -> AutoscalePr
                 scale_action=ScaleAction(
                     direction=ScaleDirection.INCREASE,
                     type=ScaleType.CHANGE_COUNT,
-                    value=2,
-                    cooldown=timedelta(minutes=10),
+                    value=scale_out_amount,
+                    cooldown=timedelta(minutes=scale_out_cooldown),
                 ),
             ),
             # Scale in
@@ -174,14 +184,20 @@ def create_auto_scale_profile(min: int, max: int, queue_uri: str) -> AutoscalePr
                 scale_action=ScaleAction(
                     direction=ScaleDirection.DECREASE,
                     type=ScaleType.CHANGE_COUNT,
-                    value=1,
-                    cooldown=timedelta(minutes=15),
+                    value=scale_in_amount,
+                    cooldown=timedelta(minutes=scale_in_cooldown),
                 ),
             ),
         ],
     )
 
 
+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, 15
+    )
+
+
 def setup_auto_scale_diagnostics(
     auto_scale_resource_uri: str,
     auto_scale_resource_name: str,

src/api-service/__app__/onefuzzlib/workers/scalesets.py

@@ -22,6 +22,7 @@ from onefuzztypes.events import (
     EventScalesetResizeScheduled,
     EventScalesetStateUpdated,
 )
+from onefuzztypes.models import AutoScale as BASE_AUTOSCALE
 from onefuzztypes.models import Error
 from onefuzztypes.models import Scaleset as BASE_SCALESET
 from onefuzztypes.models import ScalesetNodeState
@@ -29,7 +30,11 @@ from onefuzztypes.primitives import PoolName, Region
 
 from ..__version__ import __version__
 from ..azure.auth import build_auth
-from ..azure.auto_scale import add_auto_scale_to_vmss, create_auto_scale_profile
+from ..azure.auto_scale import (
+    add_auto_scale_to_vmss,
+    create_auto_scale_profile,
+    default_auto_scale_profile,
+)
 from ..azure.image import get_os
 from ..azure.network import Network
 from ..azure.queue import get_resource_id
@@ -104,6 +109,7 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             tags=tags,
         )
         entry.save()
+
         send_event(
             EventScalesetCreated(
                 scaleset_id=entry.scaleset_id,
@@ -874,8 +880,67 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             logging.error(capacity_failed)
             return capacity_failed
 
+        auto_scale_config = AutoScale.get_settings_for_scaleset(self.scaleset_id)
+        if auto_scale_config is None:
+            auto_scale_profile = default_auto_scale_profile(pool_queue_uri, capacity)
+        else:
+            logging.error("Using existing auto scale settings from database")
             auto_scale_profile = create_auto_scale_profile(
-            capacity, capacity, pool_queue_uri
+                pool_queue_uri,
+                auto_scale_config.min,
+                auto_scale_config.max,
+                auto_scale_config.default,
+                auto_scale_config.scale_out_amount,
+                auto_scale_config.scale_out_cooldown,
+                auto_scale_config.scale_in_amount,
+                auto_scale_config.scale_in_cooldown,
             )
         logging.info("Added auto scale resource to scaleset: %s" % self.scaleset_id)
         return add_auto_scale_to_vmss(self.scaleset_id, auto_scale_profile)
+
+
+class AutoScale(BASE_AUTOSCALE, ORMMixin):
+    @classmethod
+    def create(
+        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,
+    ) -> "AutoScale":
+        entry = cls(
+            scaleset_id=scaleset_id,
+            min=min,
+            max=max,
+            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,
+        )
+        entry.save()
+        return entry
+
+    @classmethod
+    def get_settings_for_scaleset(cls, scaleset_id: UUID) -> Union["AutoScale", 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
+            )
+
+        return autoscale[0]
+
+    @classmethod
+    def key_fields(cls) -> Tuple[str, None]:
+        return ("scaleset_id", None)

src/api-service/__app__/scaleset/__init__.py

@@ -20,7 +20,7 @@ from ..onefuzzlib.config import InstanceConfig
 from ..onefuzzlib.endpoint_authorization import call_if_user, check_can_manage_pools
 from ..onefuzzlib.request import not_ok, ok, parse_request
 from ..onefuzzlib.workers.pools import Pool
-from ..onefuzzlib.workers.scalesets import Scaleset
+from ..onefuzzlib.workers.scalesets import AutoScale, Scaleset
 
 
 def get(req: func.HttpRequest) -> func.HttpResponse:
@@ -104,6 +104,19 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
         ephemeral_os_disks=request.ephemeral_os_disks,
         tags=tags,
     )
+
+    if request.auto_scale:
+        AutoScale.create(
+            scaleset_id=scaleset.scaleset_id,
+            min=request.auto_scale.min,
+            max=request.auto_scale.max,
+            default=request.auto_scale.default,
+            scale_out_amount=request.auto_scale.scale_out_amount,
+            scale_out_cooldown=request.auto_scale.scale_out_cooldown,
+            scale_in_amount=request.auto_scale.scale_in_amount,
+            scale_in_cooldown=request.auto_scale.scale_in_cooldown,
+        )
+
     # don't return auths during create, only 'get' with include_auth
     scaleset.auth = None
     return ok(scaleset)

src/cli/onefuzz/api.py

@@ -1357,6 +1357,11 @@ class Scaleset(Endpoint):
         spot_instances: bool = False,
         ephemeral_os_disks: bool = False,
         tags: Optional[Dict[str, str]] = None,
+        min_instances: Optional[int] = 1,
+        scale_out_amount: Optional[int] = 1,
+        scale_out_cooldown: Optional[int] = 10,
+        scale_in_amount: Optional[int] = 1,
+        scale_in_cooldown: Optional[int] = 15,
     ) -> models.Scaleset:
         self.logger.debug("create scaleset")
 
@@ -1372,6 +1377,16 @@ class Scaleset(Endpoint):
             else:
                 raise NotImplementedError
 
+        auto_scale = requests.AutoScaleOptions(
+            min=min_instances,
+            max=size,
+            default=size,
+            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 self._req_model(
             "POST",
             models.Scaleset,
@@ -1384,6 +1399,7 @@ class Scaleset(Endpoint):
                 spot_instances=spot_instances,
                 ephemeral_os_disks=ephemeral_os_disks,
                 tags=tags,
+                auto_scale=auto_scale,
             ),
         )
 

src/pytypes/onefuzztypes/models.py

@@ -661,6 +661,17 @@ class Scaleset(BaseModel):
     tags: Dict[str, str] = Field(default_factory=lambda: {})
 
 
+class AutoScale(BaseModel):
+    scaleset_id: UUID
+    min: int = Field(ge=1)
+    max: int = Field(ge=1)
+    default: int = Field(ge=1)
+    scale_out_amount: int = Field(ge=1)
+    scale_out_cooldown: int = Field(ge=1)
+    scale_in_amount: int = Field(ge=1)
+    scale_in_cooldown: int = Field(ge=1)
+
+
 class NotificationConfig(BaseModel):
     config: NotificationTemplate
 

src/pytypes/onefuzztypes/requests.py

@@ -169,6 +169,16 @@ class ScalesetUpdate(BaseRequest):
     size: Optional[int] = Field(ge=1)
 
 
+class AutoScaleOptions(BaseModel):
+    min: int = Field(ge=1)
+    max: int = Field(ge=1)
+    default: int = Field(ge=1)
+    scale_out_amount: int = Field(ge=1)
+    scale_out_cooldown: int = Field(ge=1)
+    scale_in_amount: int = Field(ge=1)
+    scale_in_cooldown: int = Field(ge=1)
+
+
 class ScalesetCreate(BaseRequest):
     pool_name: PoolName
     vm_sku: str
@@ -178,6 +188,7 @@ class ScalesetCreate(BaseRequest):
     spot_instances: bool
     ephemeral_os_disks: bool = Field(default=False)
     tags: Dict[str, str]
+    auto_scale: Optional[AutoScaleOptions]
 
 
 class ContainerGet(BaseRequest):