DREIMP-10951: Autoscale VitessCell pods with HPA

paasta_tools/kubernetes_tools.py

@@ -871,6 +871,11 @@ def get_autoscaling_provider_spec(
         )
         return None
 
+    def get_autoscaling_target(self, name: str) -> V2beta2CrossVersionObjectReference:
+        return V2beta2CrossVersionObjectReference(
+            api_version="apps/v1", kind="Deployment", name=name
+        )
+
     def get_autoscaling_metric_spec(
         self,
         name: str,
@@ -926,9 +931,7 @@ def get_autoscaling_metric_spec(
                 max_replicas=max_replicas,
                 min_replicas=min_replicas,
                 metrics=metrics,
-                scale_target_ref=V2beta2CrossVersionObjectReference(
-                    api_version="apps/v1", kind="Deployment", name=name
-                ),
+                scale_target_ref=self.get_autoscaling_target(name),
             ),
         )
 

paasta_tools/setup_kubernetes_cr.py

@@ -50,6 +50,7 @@
 from paasta_tools.utils import load_all_configs
 from paasta_tools.utils import load_system_paasta_config
 from paasta_tools.vitesscell_tools import load_vitess_cell_instance_configs
+from paasta_tools.vitesscell_tools import update_vitess_cell_related_api_objects
 from paasta_tools.vitesscell_tools import VITESSCELL_KUBERNETES_NAMESPACE
 from paasta_tools.vitesscluster_tools import load_vitess_cluster_instance_configs
 from paasta_tools.vitesscluster_tools import VITESSCLUSTER_KUBERNETES_NAMESPACE
@@ -67,6 +68,11 @@
 }
 
 
+INSTANCE_TYPE_TO_RELATED_OBJECTS_UPDATER = {
+    "vitesscell": update_vitess_cell_related_api_objects,
+}
+
+
 INSTANCE_TYPE_TO_NAMESPACE_LOADER = {
     "vitesscluster": VITESSCLUSTER_KUBERNETES_NAMESPACE,
     "vitesscell": VITESSCELL_KUBERNETES_NAMESPACE,
@@ -444,6 +450,15 @@ def reconcile_kubernetes_resource(
                 )
             else:
                 log.info(f"{desired_resource} is up to date, no action taken")
+
+            if crd.file_prefix in INSTANCE_TYPE_TO_RELATED_OBJECTS_UPDATER:
+                INSTANCE_TYPE_TO_RELATED_OBJECTS_UPDATER[crd.file_prefix](
+                    service=service,
+                    instance=inst,
+                    cluster=cluster,
+                    kube_client=kube_client,
+                    soa_dir=DEFAULT_SOA_DIR,
+                )
         except Exception as e:
             log.error(str(e))
             succeeded = False

paasta_tools/vitesscell_tools.py

@@ -8,7 +8,9 @@
 from typing import Union
 
 import service_configuration_lib
+from kubernetes.client import V2beta2CrossVersionObjectReference
 
+from paasta_tools.kubernetes_tools import KubeClient
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfigDict
 from paasta_tools.kubernetes_tools import sanitised_cr_name
 from paasta_tools.utils import BranchDictV2
@@ -59,6 +61,7 @@ class GatewayConfigDict(TypedDict, total=False):
     extraFlags: Dict[str, str]
     extraLabels: Dict[str, str]
     replicas: int
+    scale_selector: str
     resources: Dict[str, Any]
     annotations: Mapping[str, Any]
 
@@ -115,6 +118,7 @@ def get_cell_config(
             },
             extraLabels=labels,
             replicas=replicas,
+            scale_selector=",".join([f"{k}={v}" for k, v in labels.items()]),
             resources={
                 "requests": requests,
                 "limits": requests,
@@ -175,6 +179,69 @@ def get_global_lock_server(self) -> Dict[str, str]:
             "rootPath": TOPO_GLOBAL_ROOT,
         }
 
+    def get_autoscaling_target(self, name: str) -> V2beta2CrossVersionObjectReference:
+        return V2beta2CrossVersionObjectReference(
+            api_version="planetscale.com/v2", kind="VitessCell", name=name
+        )
+
+    def get_min_instances(self) -> Optional[int]:
+        vtgate_resources = self.config_dict.get("vtgate_resources")
+        return vtgate_resources.get("min_instances", 1)
+
+    def get_max_instances(self) -> Optional[int]:
+        vtgate_resources = self.config_dict.get("vtgate_resources")
+        return vtgate_resources.get("max_instances")
+
+    def update_related_api_objects(
+        self,
+        kube_client: KubeClient,
+    ):
+        name = sanitised_cr_name(self.service, self.instance)
+
+        vtgate_resources = self.config_dict.get("vtgate_resources")
+        min_instances = vtgate_resources.get("min_instances", 1)
+        max_instances = vtgate_resources.get("max_instances")
+        should_exist = min_instances and max_instances
+
+        exists = (
+            len(
+                kube_client.autoscaling.list_namespaced_horizontal_pod_autoscaler(
+                    field_selector=f"metadata.name={name}",
+                    namespace=self.get_namespace(),
+                    limit=1,
+                ).items
+            )
+            > 0
+        )
+
+        if should_exist:
+            hpa = self.get_autoscaling_metric_spec(
+                name=name,
+                cluster=self.get_cluster(),
+                kube_client=kube_client,
+                namespace=self.get_namespace(),
+            )
+            if not hpa:
+                return
+
+            if exists:
+                kube_client.autoscaling.replace_namespaced_horizontal_pod_autoscaler(
+                    name=name,
+                    namespace=self.get_namespace(),
+                    body=hpa,
+                )
+            else:
+                log.info(f"Creating HPA for {name} in {self.get_namespace()}")
+                kube_client.autoscaling.create_namespaced_horizontal_pod_autoscaler(
+                    namespace=self.get_namespace(),
+                    body=hpa,
+                )
+        elif exists:
+            kube_client.autoscaling.delete_namespaced_horizontal_pod_autoscaler(
+                name=name,
+                namespace=self.get_namespace(),
+            )
+
     def get_vitess_cell_config(self) -> VitessCellConfigDict:
         cell = self.config_dict.get("cell")
         all_cells = self.config_dict.get("cells")
@@ -278,6 +345,18 @@ def load_vitess_cell_instance_configs(
     return vitess_cell_instance_configs
 
 
+def update_vitess_cell_related_api_objects(
+    service: str,
+    instance: str,
+    cluster: str,
+    kube_client: KubeClient,
+    soa_dir: str = DEFAULT_SOA_DIR,
+) -> None:
+    load_vitess_cell_instance_config(
+        service, instance, cluster, soa_dir=soa_dir
+    ).update_related_api_objects(kube_client)
+
+
 # TODO: read this from CRD in service configs
 def cr_id(service: str, instance: str) -> Mapping[str, str]:
     return dict(

paasta_tools/vitesscluster_tools.py

@@ -73,6 +73,8 @@ class RequestsDict(TypedDict, total=False):
 
 class ResourceConfigDict(TypedDict, total=False):
     replicas: int
+    min_instances: Optional[int]
+    max_instances: Optional[int]
     requests: Dict[str, RequestsDict]
     limits: Dict[str, RequestsDict]