Add the example on using KRR as a library

1 files changed, 48 insertions, 0 deletions

diff --git a/examples/custom_strategy.py b/examples/custom_strategy.py
new file mode 100644
index 0000000..e426477
--- /dev/null
+++ b/examples/custom_strategy.py
@@ -0,0 +1,48 @@
+# This is an example on how to create your own custom strategy
+
+from decimal import Decimal
+
+import pydantic as pd
+
+from robusta_krr import run
+from robusta_krr.core.abstract.strategies import (
+    BaseStrategy,
+    HistoryData,
+    K8sObjectData,
+    ResourceRecommendation,
+    ResourceType,
+    RunResult,
+    StrategySettings,
+)
+
+
+class CustomStrategySettings(StrategySettings):
+    cpu_percentile: Decimal = pd.Field(99, gt=0, description="The percentile to use for the request recommendation.")
+    memory_percentile: Decimal = pd.Field(
+        105, gt=0, description="The percentile to use for the request recommendation."
+    )
+
+
+class CustomStrategy(BaseStrategy[CustomStrategySettings]):
+    __display_name__ = "custom"
+
+    def run(self, history_data: HistoryData, object_data: K8sObjectData) -> RunResult:
+        cpu_usage = self._calculate_percentile(history_data[ResourceType.CPU], self.settings.cpu_percentile)
+        memory_usage = self._calculate_percentile(history_data[ResourceType.Memory], self.settings.memory_percentile)
+
+        return {
+            ResourceType.CPU: ResourceRecommendation(request=cpu_usage, limit=None),
+            ResourceType.Memory: ResourceRecommendation(request=memory_usage, limit=memory_usage),
+        }
+
+    def _calculate_percentile(self, data: dict[str, list[Decimal]], percentile: Decimal) -> Decimal:
+        data_ = [value for values in data.values() for value in values]
+        if len(data_) == 0:
+            return Decimal("NaN")
+
+        return max(data_) * percentile / 100
+
+
+# Running this file will register the strategy and make it available to the CLI
+if __name__ == "__main__":
+    run()