Merge pull request #99 from robusta-dev/hpa-v1-fallback

Fallback to HPAv1 if HPAv2 not available

robusta_krr/core/integrations/kubernetes.py

@@ -18,6 +18,7 @@
     V1PodList,
     V1StatefulSet,
     V1StatefulSetList,
+    V1HorizontalPodAutoscalerList,
     V2HorizontalPodAutoscaler,
     V2HorizontalPodAutoscalerList,
 )
@@ -30,6 +31,7 @@
 from .rollout import RolloutAppsV1Api
 
 AnyKubernetesAPIObject = Union[V1Deployment, V1DaemonSet, V1StatefulSet, V1Pod, V1Job]
+HPAKey = tuple[str, str, str]
 
 
 class ClusterLoader(Configurable):
@@ -48,7 +50,8 @@ def __init__(self, cluster: Optional[str], *args, **kwargs):
         self.rollout = RolloutAppsV1Api(api_client=self.api_client)
         self.batch = client.BatchV1Api(api_client=self.api_client)
         self.core = client.CoreV1Api(api_client=self.api_client)
-        self.autoscaling = client.AutoscalingV2Api(api_client=self.api_client)
+        self.autoscaling_v1 = client.AutoscalingV1Api(api_client=self.api_client)
+        self.autoscaling_v2 = client.AutoscalingV2Api(api_client=self.api_client)
 
     async def list_scannable_objects(self) -> list[K8sObjectData]:
         """List all scannable objects.
@@ -133,7 +136,7 @@ async def __build_obj(self, item: AnyKubernetesAPIObject, container: V1Container
             container=container.name,
             allocations=ResourceAllocations.from_container(container),
             pods=await self.__list_pods(item),
-            hpa=self.__hpa_list.get((kind, name)),
+            hpa=self.__hpa_list.get((namespace, kind, name)),
         )
 
     async def _list_deployments(self) -> list[K8sObjectData]:
@@ -233,16 +236,35 @@ async def _list_pods(self) -> list[K8sObjectData]:
             *[self.__build_obj(item, container) for item in ret.items for container in item.spec.containers]
         )
 
-    async def __list_hpa(self) -> dict[tuple[str, str], HPAData]:
-        """List all HPA objects in the cluster.
+    async def __list_hpa_v1(self) -> dict[HPAKey, HPAData]:
+        loop = asyncio.get_running_loop()
 
-        Returns:
-            dict[tuple[str, str], HPAData]: A dictionary of HPA objects, indexed by scaleTargetRef (kind, name).
-        """
+        res: V1HorizontalPodAutoscalerList = await loop.run_in_executor(
+            self.executor, lambda: self.autoscaling_v1.list_horizontal_pod_autoscaler_for_all_namespaces(watch=False)
+        )
 
+        return {
+            (
+                hpa.metadata.namespace,
+                hpa.spec.scale_target_ref.kind,
+                hpa.spec.scale_target_ref.name,
+            ): HPAData(
+                min_replicas=hpa.spec.min_replicas,
+                max_replicas=hpa.spec.max_replicas,
+                current_replicas=hpa.status.current_replicas,
+                desired_replicas=hpa.status.desired_replicas,
+                target_cpu_utilization_percentage=hpa.spec.target_cpu_utilization_percentage,
+                target_memory_utilization_percentage=None,
+            )
+            for hpa in res.items
+        }
+
+    async def __list_hpa_v2(self) -> dict[HPAKey, HPAData]:
         loop = asyncio.get_running_loop()
+
         res: V2HorizontalPodAutoscalerList = await loop.run_in_executor(
-            self.executor, lambda: self.autoscaling.list_horizontal_pod_autoscaler_for_all_namespaces(watch=False)
+            self.executor,
+            lambda: self.autoscaling_v2.list_horizontal_pod_autoscaler_for_all_namespaces(watch=False),
         )
 
         def __get_metric(hpa: V2HorizontalPodAutoscaler, metric_name: str) -> Optional[float]:
@@ -256,7 +278,11 @@ def __get_metric(hpa: V2HorizontalPodAutoscaler, metric_name: str) -> Optional[f
             )
 
         return {
-            (hpa.spec.scale_target_ref.kind, hpa.spec.scale_target_ref.name): HPAData(
+            (
+                hpa.metadata.namespace,
+                hpa.spec.scale_target_ref.kind,
+                hpa.spec.scale_target_ref.name,
+            ): HPAData(
                 min_replicas=hpa.spec.min_replicas,
                 max_replicas=hpa.spec.max_replicas,
                 current_replicas=hpa.status.current_replicas,
@@ -267,6 +293,25 @@ def __get_metric(hpa: V2HorizontalPodAutoscaler, metric_name: str) -> Optional[f
             for hpa in res.items
         }
 
+    # TODO: What should we do in case of other metrics bound to the HPA?
+    async def __list_hpa(self) -> dict[HPAKey, HPAData]:
+        """List all HPA objects in the cluster.
+
+        Returns:
+            dict[tuple[str, str], HPAData]: A dictionary of HPA objects, indexed by scaleTargetRef (kind, name).
+        """
+
+        try:
+            # Try to use V2 API first
+            return await self.__list_hpa_v2()
+        except ApiException as e:
+            if e.status != 404:
+                # If the error is other than not found, then re-raise it.
+                raise
+
+            # If V2 API does not exist, fall back to V1
+            return await self.__list_hpa_v1()
+
 
 class KubernetesLoader(Configurable):
     async def list_clusters(self) -> Optional[list[str]]: