Replace decimals on numpy arrays, keep timestapms

1 files changed, 13 insertions, 13 deletions

diff --git a/robusta_krr/strategies/simple.py b/robusta_krr/strategies/simple.py
index 5dba2db..456123c 100644
--- a/robusta_krr/strategies/simple.py
+++ b/robusta_krr/strategies/simple.py
@@ -1,6 +1,6 @@
-from decimal import Decimal
-
 import pydantic as pd
+import numpy as np
+from numpy.typing import NDArray
 
 from robusta_krr.core.abstract.strategies import (
     BaseStrategy,
@@ -14,26 +14,26 @@ from robusta_krr.core.abstract.strategies import (
 
 
 class SimpleStrategySettings(StrategySettings):
-    cpu_percentile: Decimal = pd.Field(
+    cpu_percentile: float = pd.Field(
         99, gt=0, le=100, description="The percentile to use for the CPU recommendation."
     )
-    memory_buffer_percentage: Decimal = pd.Field(
+    memory_buffer_percentage: float = pd.Field(
         5, gt=0, description="The percentage of added buffer to the peak memory usage for memory recommendation."
     )
 
-    def calculate_memory_proposal(self, data: dict[str, list[Decimal]]) -> Decimal:
-        data_ = [value for values in data.values() for value in values]
+    def calculate_memory_proposal(self, data: dict[str, NDArray[np.float64]]) -> float:
+        data_ = [np.max(values[:, 1]) for values in data.values()]
         if len(data_) == 0:
-            return Decimal("NaN")
+            return float("NaN")
 
-        return max(data_) * Decimal(1 + self.memory_buffer_percentage / 100)
+        return float(max(data_) * (1 + self.memory_buffer_percentage / 100))
 
-    def calculate_cpu_proposal(self, data: dict[str, list[Decimal]]) -> Decimal:
-        data_ = [value for values in data.values() for value in values]
-        if len(data_) == 0:
-            return Decimal("NaN")
+    def calculate_cpu_proposal(self, data: dict[str, NDArray[np.float64]]) -> float:
+        if len(data) == 0:
+            return float("NaN")
 
-        return data_[int((len(data_) - 1) * self.cpu_percentile / 100)]
+        data_ = np.concatenate([values[:, 1] for values in data.values()]) if len(data) > 1 else list(data.values())[0]
+        return float(np.percentile(data_, self.cpu_percentile / 100))
 
 
 class SimpleStrategy(BaseStrategy[SimpleStrategySettings]):