Fix CPU calculation

6 files changed, 63 insertions, 12 deletions

diff --git a/robusta_krr/core/abstract/strategies.py b/robusta_krr/core/abstract/strategies.py
index a9a084d..d91f814 100644
--- a/robusta_krr/core/abstract/strategies.py
+++ b/robusta_krr/core/abstract/strategies.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import abc
 import datetime
 from textwrap import dedent
-from typing import Annotated, Generic, Literal, Optional, TypeVar, get_args, TYPE_CHECKING
+from typing import Annotated, Generic, Literal, Optional, TypeVar, get_args, TYPE_CHECKING, Sequence
 
 import numpy as np
 import pydantic as pd
@@ -65,7 +65,7 @@ ArrayNx2 = Annotated[NDArray[np.float64], Literal["N", 2]]
 
 
 PodsTimeData = dict[str, ArrayNx2]  # Mapping: pod -> [(time, value)]
-MetricsPodData = dict[type["BaseMetric"], PodsTimeData]
+MetricsPodData = dict[str, PodsTimeData]
 
 RunResult = dict[ResourceType, ResourceRecommendation]
 
@@ -95,8 +95,12 @@ class BaseStrategy(abc.ABC, Generic[_StrategySettings]):
 
     display_name: str
     rich_console: bool = False
-    # TODO: this should be BaseMetricLoader, but currently we only support Prometheus
-    metrics: list[type[PrometheusMetric]] = []
+
+    # TODO: this should be BaseMetric, but currently we only support Prometheus
+    @property
+    @abc.abstractmethod
+    def metrics(self) -> Sequence[type[PrometheusMetric]]:
+        pass
 
     def __init__(self, settings: _StrategySettings):
         self.settings = settings
@@ -106,7 +110,7 @@ class BaseStrategy(abc.ABC, Generic[_StrategySettings]):
 
     @property
     def _display_name(self) -> str:
-        return getattr(self, "__display_name__", self.__class__.__name__.lower().removeprefix("strategy"))
+        return getattr(self, "display_name", self.__class__.__name__.lower().removeprefix("strategy"))
 
     @property
     def description(self) -> Optional[str]:
diff --git a/robusta_krr/core/integrations/prometheus/loader.py b/robusta_krr/core/integrations/prometheus/loader.py
index 0daf3a6..a49c01c 100644
--- a/robusta_krr/core/integrations/prometheus/loader.py
+++ b/robusta_krr/core/integrations/prometheus/loader.py
@@ -100,6 +100,6 @@ class PrometheusMetricsLoader(Configurable):
         await self.loader.add_historic_pods(object, period)
 
         return {
-            MetricLoader: await self.loader.gather_data(object, MetricLoader, period, step)
+            MetricLoader.__name__: await self.loader.gather_data(object, MetricLoader, period, step)
             for MetricLoader in strategy.metrics
         }
diff --git a/robusta_krr/core/integrations/prometheus/metrics/__init__.py b/robusta_krr/core/integrations/prometheus/metrics/__init__.py
index 5b9a45b..daee396 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/__init__.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/__init__.py
@@ -1,3 +1,3 @@
-from .cpu import CPULoader, MaxCPULoader
-from .memory import MemoryLoader, MaxMemoryLoader
+from .cpu import CPULoader, MaxCPULoader, PercentileCPULoader
+from .memory import MemoryLoader, MaxMemoryLoader, PercentileMemoryLoader
 from .base import PrometheusMetric
diff --git a/robusta_krr/core/integrations/prometheus/metrics/cpu.py b/robusta_krr/core/integrations/prometheus/metrics/cpu.py
index 41e1137..0239b38 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/cpu.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/cpu.py
@@ -37,3 +37,25 @@ class MaxCPULoader(QueryMetric, FilterMetric):
                 )
             ) by (container, pod, job)
         """
+
+
+def PercentileCPULoader(percentile: float) -> type[QueryMetric]:
+    class PercentileCPULoader(QueryMetric, FilterMetric):
+        def get_query(self, object: K8sObjectData, resolution: str) -> str:
+            pods_selector = "|".join(pod.name for pod in object.pods)
+            cluster_label = self.get_prometheus_cluster_label()
+            return f"""
+                quantile_over_time(
+                    {round(percentile / 100, 2)},
+                    irate(
+                        container_cpu_usage_seconds_total{{
+                            namespace="{object.namespace}",
+                            pod=~"{pods_selector}",
+                            container="{object.container}"
+                            {cluster_label}
+                        }}[1m]
+                    )[{resolution}]
+                )
+            """
+
+    return PercentileCPULoader
diff --git a/robusta_krr/core/integrations/prometheus/metrics/memory.py b/robusta_krr/core/integrations/prometheus/metrics/memory.py
index 32df5f4..5fc6f73 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/memory.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/memory.py
@@ -33,3 +33,23 @@ class MaxMemoryLoader(QueryMetric, FilterMetric):
                 }}[{resolution}]
             )
         """
+
+
+def PercentileMemoryLoader(percentile: float) -> type[QueryMetric]:
+    class PercentileMemoryLoader(QueryMetric, FilterMetric):
+        def get_query(self, object: K8sObjectData, resolution: str) -> str:
+            pods_selector = "|".join(pod.name for pod in object.pods)
+            cluster_label = self.get_prometheus_cluster_label()
+            return f"""
+                quantile_over_time(
+                    {round(percentile / 100, 2)},
+                    container_memory_working_set_bytes{{
+                        namespace="{object.namespace}",
+                        pod=~"{pods_selector}",
+                        container="{object.container}"
+                        {cluster_label}
+                    }}[{resolution}]
+                )
+            """
+
+    return PercentileMemoryLoader
diff --git a/robusta_krr/strategies/simple.py b/robusta_krr/strategies/simple.py
index d3b20be..5db2292 100644
--- a/robusta_krr/strategies/simple.py
+++ b/robusta_krr/strategies/simple.py
@@ -1,3 +1,4 @@
+from typing import Sequence
 import numpy as np
 import pydantic as pd
 
@@ -11,7 +12,8 @@ from robusta_krr.core.abstract.strategies import (
     RunResult,
     StrategySettings,
 )
-from robusta_krr.core.integrations.prometheus.metrics import MaxCPULoader, MaxMemoryLoader
+from robusta_krr.core.abstract.metrics import BaseMetric
+from robusta_krr.core.integrations.prometheus.metrics import PercentileCPULoader, MaxMemoryLoader, PrometheusMetric
 
 
 class SimpleStrategySettings(StrategySettings):
@@ -52,12 +54,15 @@ class SimpleStrategy(BaseStrategy[SimpleStrategySettings]):
 
     display_name = "simple"
     rich_console = True
-    metrics = [MaxCPULoader, MaxMemoryLoader]
+
+    @property
+    def metrics(self) -> list[type[PrometheusMetric]]:
+        return [PercentileCPULoader(self.settings.cpu_percentile), MaxMemoryLoader]
 
     def __calculate_cpu_proposal(
         self, history_data: MetricsPodData, object_data: K8sObjectData
     ) -> ResourceRecommendation:
-        data = history_data[MaxCPULoader]
+        data = history_data["PercentileCPULoader"]
 
         if len(data) == 0:
             return ResourceRecommendation.undefined(info="No data")
@@ -71,7 +76,7 @@ class SimpleStrategy(BaseStrategy[SimpleStrategySettings]):
     def __calculate_memory_proposal(
         self, history_data: MetricsPodData, object_data: K8sObjectData
     ) -> ResourceRecommendation:
-        data = history_data[MaxMemoryLoader]
+        data = history_data["MaxMemoryLoader"]
 
         if len(data) == 0:
             return ResourceRecommendation.undefined(info="No data")