Add multiple workload loaders, refactor kubeapi workload loader

robusta-dev · Apr 22, 2024 · 4fedd82 · 4fedd82
1 parent 5695c66
commit 4fedd82
Show file tree

Hide file tree

Showing 30 changed files with 739 additions and 491 deletions.
diff --git a/examples/custom_strategy.py b/examples/custom_strategy.py
@@ -3,7 +3,7 @@
 import pydantic as pd
 
 import robusta_krr
-from robusta_krr.api.models import K8sObjectData, MetricsPodData, ResourceRecommendation, ResourceType, RunResult
+from robusta_krr.api.models import K8sWorkload, MetricsPodData, ResourceRecommendation, ResourceType, RunResult
 from robusta_krr.api.strategies import BaseStrategy, StrategySettings
 from robusta_krr.core.integrations.prometheus.metrics import MaxMemoryLoader, PercentileCPULoader
 
@@ -24,7 +24,7 @@ class CustomStrategy(BaseStrategy[CustomStrategySettings]):
     rich_console = True  # Whether to use rich console for the CLI
     metrics = [PercentileCPULoader(90), MaxMemoryLoader]  # The metrics to use for the strategy
 
-    def run(self, history_data: MetricsPodData, object_data: K8sObjectData) -> RunResult:
+    def run(self, history_data: MetricsPodData, object_data: K8sWorkload) -> RunResult:
         return {
             ResourceType.CPU: ResourceRecommendation(request=self.settings.param_1, limit=None),
             ResourceType.Memory: ResourceRecommendation(request=self.settings.param_2, limit=self.settings.param_2),

diff --git a/robusta_krr/api/models.py b/robusta_krr/api/models.py
@@ -1,14 +1,14 @@
 from robusta_krr.core.abstract.strategies import MetricsPodData, PodsTimeData, ResourceRecommendation, RunResult
 from robusta_krr.core.models.allocations import RecommendationValue, ResourceAllocations, ResourceType
-from robusta_krr.core.models.objects import K8sObjectData, PodData
+from robusta_krr.core.models.objects import K8sWorkload, PodData
 from robusta_krr.core.models.result import ResourceScan, Result
 from robusta_krr.core.models.severity import Severity, register_severity_calculator
 
 __all__ = [
     "ResourceType",
     "ResourceAllocations",
     "RecommendationValue",
-    "K8sObjectData",
+    "K8sWorkload",
     "PodData",
     "Result",
     "Severity",

diff --git a/robusta_krr/core/abstract/metrics.py b/robusta_krr/core/abstract/metrics.py
@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 
 from robusta_krr.core.abstract.strategies import PodsTimeData
-from robusta_krr.core.models.objects import K8sObjectData
+from robusta_krr.core.models.objects import K8sWorkload
 
 
 class BaseMetric(ABC):
@@ -16,6 +16,6 @@ class BaseMetric(ABC):
 
     @abstractmethod
     async def load_data(
-        self, object: K8sObjectData, period: datetime.timedelta, step: datetime.timedelta
+        self, object: K8sWorkload, period: datetime.timedelta, step: datetime.timedelta
     ) -> PodsTimeData:
         ...
diff --git a/robusta_krr/core/abstract/strategies.py b/robusta_krr/core/abstract/strategies.py
@@ -9,7 +9,7 @@
 import pydantic as pd
 from numpy.typing import NDArray
 
-from robusta_krr.core.models.result import K8sObjectData, ResourceType
+from robusta_krr.core.models.result import K8sWorkload, ResourceType
 
 if TYPE_CHECKING:
     from robusta_krr.core.abstract.metrics import BaseMetric  # noqa: F401
@@ -133,7 +133,7 @@ def description(self) -> Optional[str]:
     # Abstract method that needs to be implemented by subclass.
     # This method is intended to calculate resource recommendation based on history data and kubernetes object data.
     @abc.abstractmethod
-    def run(self, history_data: MetricsPodData, object_data: K8sObjectData) -> RunResult:
+    def run(self, history_data: MetricsPodData, object_data: K8sWorkload) -> RunResult:
         pass
 
     # This method is intended to return a strategy by its name.
@@ -167,6 +167,6 @@ def get_settings_type(cls) -> type[StrategySettings]:
     "StrategySettings",
     "PodsTimeData",
     "MetricsPodData",
-    "K8sObjectData",
+    "K8sWorkload",
     "ResourceType",
 ]