Reworked processing of units

10 files changed, 150 insertions, 97 deletions

diff --git a/robusta_krr/core/abstract/strategies.py b/robusta_krr/core/abstract/strategies.py
index 51edf31..e5f6813 100644
--- a/robusta_krr/core/abstract/strategies.py
+++ b/robusta_krr/core/abstract/strategies.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import abc
 import datetime
+from decimal import Decimal
 from typing import Generic, TypeVar
 
 import pydantic as pd
@@ -10,8 +11,8 @@ from robusta_krr.core.models.result import K8sObjectData, ResourceType
 
 
 class ResourceRecommendation(pd.BaseModel):
-    request: float
-    limit: float
+    request: Decimal | None
+    limit: Decimal | None
 
 
 class StrategySettings(pd.BaseModel):
@@ -25,7 +26,8 @@ class StrategySettings(pd.BaseModel):
 
 
 _StrategySettings = TypeVar("_StrategySettings", bound=StrategySettings)
-HistoryData = dict[str, list[float]]
+HistoryData = dict[ResourceType, list[int]]
+RunResult = dict[ResourceType, ResourceRecommendation]
 
 
 class BaseStrategy(abc.ABC, Generic[_StrategySettings]):
@@ -39,9 +41,7 @@ class BaseStrategy(abc.ABC, Generic[_StrategySettings]):
         return self.__display_name__.title()
 
     @abc.abstractmethod
-    def run(
-        self, history_data: HistoryData, object_data: K8sObjectData, resource_type: ResourceType
-    ) -> ResourceRecommendation:
+    def run(self, history_data: HistoryData, object_data: K8sObjectData) -> RunResult:
         """Run the strategy to calculate the recommendation"""
 
     @staticmethod
diff --git a/robusta_krr/core/integrations/kubernetes.py b/robusta_krr/core/integrations/kubernetes.py
index add85fb..10b0c2e 100644
--- a/robusta_krr/core/integrations/kubernetes.py
+++ b/robusta_krr/core/integrations/kubernetes.py
@@ -3,15 +3,15 @@ import itertools
 
 from kubernetes import client, config
 from kubernetes.client.models import (
-    V1PodList,
-    V1DeploymentList,
-    V1StatefulSetList,
-    V1JobList,
-    V1DaemonSetList,
-    V1Deployment,
     V1Container,
     V1DaemonSet,
+    V1DaemonSetList,
+    V1Deployment,
+    V1DeploymentList,
+    V1JobList,
+    V1PodList,
     V1StatefulSet,
+    V1StatefulSetList,
 )
 
 from robusta_krr.core.models.objects import K8sObjectData
@@ -34,7 +34,7 @@ class ClusterLoader(Configurable):
             A list of scannable objects.
         """
 
-        self.debug("Listing scannable objects")
+        self.debug(f"Listing scannable objects in {self.cluster}")
 
         try:
             objects_tuple = await asyncio.gather(
@@ -45,6 +45,7 @@ class ClusterLoader(Configurable):
             )
         except Exception as e:
             self.error(f"Error trying to list pods in cluster {self.cluster}: {e}")
+            self.debug_exception()
             return []
 
         return list(itertools.chain(*objects_tuple))
@@ -98,12 +99,8 @@ class ClusterLoader(Configurable):
 class KubernetesLoader(Configurable):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-
-        self.debug("Initializing Kubernetes client")
         config.load_kube_config()
 
-        self._kubernetes_object_allocation_cache: dict[K8sObjectData, ResourceAllocations] = {}
-
     async def list_clusters(self) -> list[str]:
         """List all clusters.
 
@@ -111,8 +108,6 @@ class KubernetesLoader(Configurable):
             A list of clusters.
         """
 
-        self.debug("Listing clusters")
-
         contexts, _ = await asyncio.to_thread(config.list_kube_config_contexts)
 
         return [context["name"] for context in contexts]
@@ -124,8 +119,6 @@ class KubernetesLoader(Configurable):
             A list of scannable objects.
         """
 
-        self.debug("Listing scannable objects")
-
         cluster_loaders = [ClusterLoader(cluster=cluster, config=self.config) for cluster in clusters]
         objects = await asyncio.gather(*[cluster_loader.list_scannable_objects() for cluster_loader in cluster_loaders])
         return list(itertools.chain(*objects))
diff --git a/robusta_krr/core/integrations/prometheus.py b/robusta_krr/core/integrations/prometheus.py
index 64e2f08..5c78838 100644
--- a/robusta_krr/core/integrations/prometheus.py
+++ b/robusta_krr/core/integrations/prometheus.py
@@ -1,10 +1,8 @@
-import asyncio
 import datetime
 import random
 
 from robusta_krr.core.models.objects import K8sObjectData
 from robusta_krr.core.models.result import ResourceType
-from robusta_krr.core.abstract.strategies import HistoryData
 from robusta_krr.utils.configurable import Configurable
 
 
@@ -16,12 +14,14 @@ class PrometheusLoader(Configurable):
         period: datetime.timedelta,
         *,
         timeframe: datetime.timedelta = datetime.timedelta(minutes=1),
-    ) -> HistoryData:
+    ) -> list[int]:
         # TODO: This is mock function. Implement this later using the Prometheus API
         self.debug(f"Gathering data for {object} and {resource} for the last {period}")
-        await asyncio.sleep(1.5)  # Simulate a slow API call
         points = int(period / timeframe)
-        return {
-            "container_1": [random.randrange(30, 300) for _ in range(points)],
-            "container_2": [random.randrange(70, 500) for _ in range(points)],
-        }
+
+        if resource == ResourceType.CPU:
+            return [random.randrange(1, 3000) for _ in range(points)]
+        elif resource == ResourceType.Memory:
+            return [random.randrange(70_000_000, 5_000_000_000) for _ in range(points)]
+        else:
+            raise ValueError(f"Unknown resource type: {resource}")
diff --git a/robusta_krr/core/models/allocations.py b/robusta_krr/core/models/allocations.py
index 12268ae..c4ec506 100644
--- a/robusta_krr/core/models/allocations.py
+++ b/robusta_krr/core/models/allocations.py
@@ -1,21 +1,35 @@
 from __future__ import annotations
 
 import enum
+from decimal import Decimal
 from typing import Self
 
+import pydantic as pd
 from kubernetes.client.models import V1Container
 
-import pydantic as pd
+from robusta_krr.utils import resource_units
 
 
 class ResourceType(str, enum.Enum):
+    """The type of resource.
+
+    Just add new types here and they will be automatically supported.
+    """
+
     CPU = "cpu"
     Memory = "memory"
 
 
 class ResourceAllocations(pd.BaseModel):
-    requests: dict[ResourceType, str | None]
-    limits: dict[ResourceType, str | None]
+    requests: dict[ResourceType, Decimal | None]
+    limits: dict[ResourceType, Decimal | None]
+
+    @pd.validator("requests", "limits", pre=True)
+    def validate_requests(cls, value: dict[ResourceType, str | Decimal | None]) -> dict[ResourceType, Decimal | None]:
+        return {
+            resource_type: resource_units.parse(resource_value) if isinstance(resource_value, str) else resource_value
+            for resource_type, resource_value in value.items()
+        }
 
     @classmethod
     def from_container(cls, container: V1Container) -> Self:
diff --git a/robusta_krr/core/models/result.py b/robusta_krr/core/models/result.py
index 89ba72b..4d690ca 100644
--- a/robusta_krr/core/models/result.py
+++ b/robusta_krr/core/models/result.py
@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import itertools
+from decimal import Decimal
 from typing import Any
 
 import pydantic as pd
 
 from robusta_krr.core.abstract.formatters import BaseFormatter
-from robusta_krr.core.models.objects import K8sObjectData
 from robusta_krr.core.models.allocations import ResourceAllocations, ResourceType
+from robusta_krr.core.models.objects import K8sObjectData
 
 
 class ResourceScan(pd.BaseModel):
@@ -38,7 +39,7 @@ class Result(pd.BaseModel):
         return _formatter.format(self)
 
     @staticmethod
-    def __percentage_difference(current: float | str | None, recommended: float | str | None) -> float:
+    def __percentage_difference(current: Decimal | None, recommended: Decimal | None) -> float:
         """Get the percentage difference between two numbers.
 
         Args:
@@ -67,4 +68,7 @@ class Result(pd.BaseModel):
                 scan.object.allocations.limits[resource_type], scan.recommended.limits[resource_type]
             )
 
+        if len(self.scans) == 0:
+            return 0.0
+
         return max(0, round(100 - total_diff / len(self.scans) / len(ResourceType) / 50, 2))  # 50 is just a constant
diff --git a/robusta_krr/core/runner.py b/robusta_krr/core/runner.py
index eb70f16..b96ba53 100644
--- a/robusta_krr/core/runner.py
+++ b/robusta_krr/core/runner.py
@@ -1,15 +1,14 @@
 import asyncio
-import itertools
 
-from robusta_krr.core.models.config import Config
+from robusta_krr.core.abstract.strategies import RunResult
 from robusta_krr.core.integrations.kubernetes import KubernetesLoader
-from robusta_krr.core.models.objects import K8sObjectData
 from robusta_krr.core.integrations.prometheus import PrometheusLoader
+from robusta_krr.core.models.config import Config
+from robusta_krr.core.models.objects import K8sObjectData
 from robusta_krr.core.models.result import ResourceAllocations, ResourceScan, ResourceType, Result
-from robusta_krr.core.abstract.strategies import ResourceRecommendation
 from robusta_krr.utils.configurable import Configurable
-from robusta_krr.utils.version import get_version
 from robusta_krr.utils.logo import ASCII_LOGO
+from robusta_krr.utils.version import get_version
 
 
 class Runner(Configurable):
@@ -31,35 +30,34 @@ class Runner(Configurable):
         self.echo("\n", no_prefix=True)
         self.console.print(formatted)
 
-    async def _calculate_object_recommendations(
-        self, object: K8sObjectData, resource: ResourceType
-    ) -> ResourceRecommendation:
-        data = await self._prometheus_loader.gather_data(
-            object,
-            resource,
-            self._strategy.settings.history_timedelta,
+    async def _calculate_object_recommendations(self, object: K8sObjectData) -> RunResult:
+        data_tuple = await asyncio.gather(
+            *[
+                self._prometheus_loader.gather_data(
+                    object,
+                    resource,
+                    self._strategy.settings.history_timedelta,
+                )
+                for resource in ResourceType
+            ]
         )
+        data = dict(zip(ResourceType, data_tuple))
 
         # NOTE: We run this in a threadpool as the strategy calculation might be CPU intensive
-        # TODO: Maybe we should do it in a processpool instead?
         # But keep in mind that numpy calcluations will not block the GIL
-        return await asyncio.to_thread(self._strategy.run, data, object, resource)
+        return await asyncio.to_thread(self._strategy.run, data, object)
 
     async def _gather_objects_recommendations(self, objects: list[K8sObjectData]) -> list[ResourceAllocations]:
-        recommendations: list[ResourceRecommendation] = await asyncio.gather(
-            *[
-                self._calculate_object_recommendations(object, resource)
-                for object, resource in itertools.product(objects, ResourceType)
-            ]
+        recommendations: list[RunResult] = await asyncio.gather(
+            *[self._calculate_object_recommendations(object) for object in objects]
         )
-        recommendations_dict = dict(zip(itertools.product(objects, ResourceType), recommendations))
 
         return [
             ResourceAllocations(
-                requests={resource: recommendations_dict[object, resource].request for resource in ResourceType},
-                limits={resource: recommendations_dict[object, resource].limit for resource in ResourceType},
+                requests={resource: recommendation[resource].request for resource in ResourceType},
+                limits={resource: recommendation[resource].limit for resource in ResourceType},
             )
-            for object in objects
+            for recommendation in recommendations
         ]
 
     async def _collect_result(self) -> Result:
diff --git a/robusta_krr/formatters/table.py b/robusta_krr/formatters/table.py
index 2fcd16f..92711c4 100644
--- a/robusta_krr/formatters/table.py
+++ b/robusta_krr/formatters/table.py
@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 import itertools
+from decimal import Decimal
+
+from rich.table import Table
 
 from robusta_krr.core.abstract.formatters import BaseFormatter
-from robusta_krr.core.models.result import Result, ResourceType
+from robusta_krr.core.models.result import ResourceScan, ResourceType, Result
 from robusta_krr.utils import resource_units
 
-from rich.table import Table
+NONE_LITERAL = "none"
+PRESCISION = 4
 
 
 class TableFormatter(BaseFormatter):
@@ -14,6 +18,16 @@ class TableFormatter(BaseFormatter):
 
     __display_name__ = "table"
 
+    def _format_united_decimal(self, value: Decimal | None, prescision: int | None = None) -> str:
+        return resource_units.format(value, prescision=prescision) if value is not None else NONE_LITERAL
+
+    def _format_request_str(self, item: ResourceScan, resource: ResourceType, selector: str) -> str:
+        return (
+            self._format_united_decimal(getattr(item.object.allocations, selector)[resource])
+            + " -> "
+            + self._format_united_decimal(getattr(item.recommended, selector)[resource], prescision=PRESCISION)
+        )
+
     def format(self, result: Result) -> Table:
         """Format the result as text.
 
@@ -52,9 +66,7 @@ class TableFormatter(BaseFormatter):
                     item.object.kind if full_info_row else "",
                     item.object.container,
                     *[
-                        f"{getattr(item.object.allocations, selector)[resource]}"
-                        + "->"
-                        + f"{resource_units.format(getattr(item.recommended, selector)[resource])}"
+                        self._format_request_str(item, resource, selector)
                         for resource in ResourceType
                         for selector in ["requests", "limits"]
                     ],
diff --git a/robusta_krr/strategies/simple.py b/robusta_krr/strategies/simple.py
index 4d46151..3f1ad60 100644
--- a/robusta_krr/strategies/simple.py
+++ b/robusta_krr/strategies/simple.py
@@ -1,3 +1,5 @@
+from decimal import Decimal
+
 import pydantic as pd
 
 from robusta_krr.core.abstract.strategies import (
@@ -6,6 +8,7 @@ from robusta_krr.core.abstract.strategies import (
     K8sObjectData,
     ResourceRecommendation,
     ResourceType,
+    RunResult,
     StrategySettings,
 )
 
@@ -22,15 +25,21 @@ class SimpleStrategySettings(StrategySettings):
 class SimpleStrategy(BaseStrategy[SimpleStrategySettings]):
     __display_name__ = "simple"
 
-    def run(
-        self, history_data: HistoryData, object_data: K8sObjectData, resource_type: ResourceType
-    ) -> ResourceRecommendation:
-        points_flatten = [point for points in history_data.values() for point in points]
-        return ResourceRecommendation(
-            request=self._calculate_percentile(points_flatten, self.settings.request_percentile),
-            limit=self._calculate_percentile(points_flatten, self.settings.limit_percentile),
-        )
+    def run(self, history_data: HistoryData, object_data: K8sObjectData) -> RunResult:
+        cpu_usage = self._calculate_percentile(history_data[ResourceType.CPU], self.settings.request_percentile)
+        memory_usage = self._calculate_percentile(history_data[ResourceType.Memory], self.settings.request_percentile)
+
+        return {
+            ResourceType.CPU: ResourceRecommendation(
+                request=Decimal(cpu_usage) / 1000,
+                limit=None,
+            ),
+            ResourceType.Memory: ResourceRecommendation(
+                request=memory_usage,
+                limit=memory_usage,
+            ),
+        }
 
-    def _calculate_percentile(self, data: list[float], percentile: float) -> float:
+    def _calculate_percentile(self, data: list[int], percentile: float) -> int:
         data = sorted(data)
         return data[int(len(data) * percentile)]
diff --git a/robusta_krr/utils/configurable.py b/robusta_krr/utils/configurable.py
index 650f0ed..e4c80f6 100644
--- a/robusta_krr/utils/configurable.py
+++ b/robusta_krr/utils/configurable.py
@@ -1,9 +1,8 @@
-from rich.console import Console
-
 from typing import Literal
 
-from robusta_krr.core.models.config import Config
+from rich.console import Console
 
+from robusta_krr.core.models.config import Config
 
 console = Console()
 
@@ -18,6 +17,14 @@ class Configurable:
         self.config = config
         self.console = console
 
+    @property
+    def debug_active(self) -> bool:
+        return self.config.verbose and not self.config.quiet
+
+    @property
+    def echo_active(self) -> bool:
+        return not self.config.quiet
+
     @staticmethod
     def __add_prefix(text: str, prefix: str, /, no_prefix: bool) -> str:
         return f"{prefix} {text}" if not no_prefix else text
@@ -32,7 +39,7 @@ class Configurable:
 
         color = {"INFO": "green", "WARNING": "yellow", "ERROR": "red"}[type]
 
-        if not self.config.quiet:
+        if self.echo_active:
             self.console.print(
                 self.__add_prefix(message, f"[bold {color}][{type}][/bold {color}]", no_prefix=no_prefix)
             )
@@ -42,9 +49,17 @@ class Configurable:
         Echoes a message to the user if verbose mode is enabled
         """
 
-        if self.config.verbose and not self.config.quiet:
+        if self.debug_active:
             self.console.print(self.__add_prefix(message, "[bold green][DEBUG][/bold green]", no_prefix=False))
 
+    def debug_exception(self) -> None:
+        """
+        Echoes the exception traceback to the user if verbose mode is enabled
+        """
+
+        if self.debug_active:
+            self.console.print_exception()
+
     def info(self, message: str = "") -> None:
         """
         Echoes an info message to the user
diff --git a/robusta_krr/utils/resource_units.py b/robusta_krr/utils/resource_units.py
index d52de35..7d51034 100644
--- a/robusta_krr/utils/resource_units.py
+++ b/robusta_krr/utils/resource_units.py
@@ -1,19 +1,19 @@
 from decimal import Decimal
 
 UNITS = {
-    "m": 1e-3,
-    "Ki": 1024,
-    "Mi": 1024**2,
-    "Gi": 1024**3,
-    "Ti": 1024**4,
-    "Pi": 1024**5,
-    "Ei": 1024**6,
-    "k": 1e3,
-    "M": 1e6,
-    "G": 1e9,
-    "T": 1e12,
-    "P": 1e15,
-    "E": 1e18,
+    "m": Decimal("1e-3"),
+    "Ki": Decimal(1024),
+    "Mi": Decimal(1024**2),
+    "Gi": Decimal(1024**3),
+    "Ti": Decimal(1024**4),
+    "Pi": Decimal(1024**5),
+    "Ei": Decimal(1024**6),
+    "k": Decimal(1e3),
+    "M": Decimal(1e6),
+    "G": Decimal(1e9),
+    "T": Decimal(1e12),
+    "P": Decimal(1e15),
+    "E": Decimal(1e18),
 }
 
 
@@ -21,16 +21,24 @@ def parse(x: str) -> Decimal:
     """Converts a string to an integer with respect of units."""
     for unit, multiplier in UNITS.items():
         if x.endswith(unit):
-            return Decimal(x[: -len(unit)]) * Decimal(multiplier)
+            return Decimal(x[: -len(unit)]) * multiplier
     return Decimal(x)
 
 
-def format(x: float | None) -> str | None:
+def format(x: Decimal, prescision: int | None = None) -> str:
     """Converts an integer to a string with respect of units."""
-    if x is None:
-        return None
 
-    for unit, multiplier in UNITS.items():
-        if Decimal(x) % Decimal(multiplier) == 0:
-            return f"{Decimal(x) / Decimal(multiplier)}{unit}"
-    return str(x)[:-2]
+    if prescision is not None:
+        # Use inly the first prescision digits, starting from the biggest one
+        # Example? 123456 -> 123000
+        assert prescision >= 0
+
+        exponent: int
+        sign, digits, exponent = x.as_tuple()  # type: ignore
+        x = Decimal((sign, list(digits[:prescision]) + [0] * (len(digits) - prescision), exponent))
+
+    for unit, multiplier in reversed(UNITS.items()):
+        if x % multiplier == 0:
+            v = int(x / multiplier)
+            return f"{v}{unit}"
+    return str(x)