Add the example on using KRR as a library

4 files changed, 144 insertions, 91 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()
diff --git a/krr.py b/krr.py
index c423e9f..3c7ac1e 100644
--- a/krr.py
+++ b/krr.py
@@ -1,4 +1,4 @@
-from robusta_krr import app
+from robusta_krr import run
 
 if __name__ == "__main__":
-    app()
+    run()
diff --git a/robusta_krr/__init__.py b/robusta_krr/__init__.py
index e4c7aa5..9c4f0db 100644
--- a/robusta_krr/__init__.py
+++ b/robusta_krr/__init__.py
@@ -1 +1,3 @@
-from .main import app
+from .main import run
+
+__all__ = ["run"]
diff --git a/robusta_krr/main.py b/robusta_krr/main.py
index 0453d26..ae43f27 100644
--- a/robusta_krr/main.py
+++ b/robusta_krr/main.py
@@ -36,95 +36,98 @@ def __process_type(_T: type) -> str:
         return "str"  # It the type is unknown, just use str and let pydantic handle it
 
 
-for strategy_name, strategy_type in BaseStrategy.get_all().items():  # type: ignore
-    FUNC_TEMPLATE = textwrap.dedent(
-        """
-        @app.command(rich_help_panel="Strategies")
-        def {func_name}(
-            ctx: typer.Context,
-            clusters: List[str] = typer.Option(
-                None,
-                "--cluster",
-                "-c",
-                help="List of clusters to run on. By default, will run on the current cluster. Use '*' to run on all clusters.",
-                rich_help_panel="Kubernetes Settings"
-            ),
-            namespaces: List[str] = typer.Option(
-                None,
-                "--namespace",
-                "-n",
-                help="List of namespaces to run on. By default, will run on all namespaces.",
-                rich_help_panel="Kubernetes Settings"
-            ),
-            prometheus_url: Optional[str] = typer.Option(
-                None,
-                "--prometheus-url",
-                "-p",
-                help="Prometheus URL. If not provided, will attempt to find it in kubernetes cluster",
-                rich_help_panel="Prometheus Settings",
-            ),
-            prometheus_auth_header: Optional[str] = typer.Option(
-                None,
-                "--prometheus-auth-header",
-                help="Prometheus authentication header.",
-                rich_help_panel="Prometheus Settings",
-            ),
-            prometheus_ssl_enabled: bool = typer.Option(
-                False,
-                "--prometheus-ssl-enabled",
-                help="Enable SSL for Prometheus requests.",
-                rich_help_panel="Prometheus Settings",
-            ),
-            format: str = typer.Option("table", "--formatter", "-f", help="Output formatter ({formatters})", rich_help_panel="Logging Settings"),
-            verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings"),
-            quiet: bool = typer.Option(False, "--quiet", "-q", help="Enable quiet mode", rich_help_panel="Logging Settings"),
-            {strategy_settings},
-        ) -> None:
-            '''Run KRR using the `{func_name}` strategy'''
-
-            config = Config(
-                clusters="*" if "*" in clusters else clusters,
-                namespaces="*" if "*" in namespaces else namespaces,
-                prometheus_url=prometheus_url,
-                prometheus_auth_header=prometheus_auth_header,
-                prometheus_ssl_enabled=prometheus_ssl_enabled,
-                format=format,
-                verbose=verbose,
-                quiet=quiet,
-                strategy="{func_name}",
-                other_args=ctx.args,
-            )
-            runner = Runner(config)
-            asyncio.run(runner.run())
-        """
-    )
-
-    exec(
-        FUNC_TEMPLATE.format(
-            func_name=strategy_name,
-            strategy_name=strategy_type.__name__,
-            strategy_settings=",\n".join(
-                f'{field_name}: {__process_type(field_meta.type_)} = typer.Option({field_meta.default!r}, "--{field_name}", help="{field_meta.field_info.description}", rich_help_panel="Strategy Settings")'
-                for field_name, field_meta in strategy_type.get_settings_type().__fields__.items()
+def run() -> None:
+    for strategy_name, strategy_type in BaseStrategy.get_all().items():  # type: ignore
+        FUNC_TEMPLATE = textwrap.dedent(
+            """
+            @app.command(rich_help_panel="Strategies")
+            def {func_name}(
+                ctx: typer.Context,
+                clusters: List[str] = typer.Option(
+                    None,
+                    "--cluster",
+                    "-c",
+                    help="List of clusters to run on. By default, will run on the current cluster. Use '*' to run on all clusters.",
+                    rich_help_panel="Kubernetes Settings"
+                ),
+                namespaces: List[str] = typer.Option(
+                    None,
+                    "--namespace",
+                    "-n",
+                    help="List of namespaces to run on. By default, will run on all namespaces.",
+                    rich_help_panel="Kubernetes Settings"
+                ),
+                prometheus_url: Optional[str] = typer.Option(
+                    None,
+                    "--prometheus-url",
+                    "-p",
+                    help="Prometheus URL. If not provided, will attempt to find it in kubernetes cluster",
+                    rich_help_panel="Prometheus Settings",
+                ),
+                prometheus_auth_header: Optional[str] = typer.Option(
+                    None,
+                    "--prometheus-auth-header",
+                    help="Prometheus authentication header.",
+                    rich_help_panel="Prometheus Settings",
+                ),
+                prometheus_ssl_enabled: bool = typer.Option(
+                    False,
+                    "--prometheus-ssl-enabled",
+                    help="Enable SSL for Prometheus requests.",
+                    rich_help_panel="Prometheus Settings",
+                ),
+                format: str = typer.Option("table", "--formatter", "-f", help="Output formatter ({formatters})", rich_help_panel="Logging Settings"),
+                verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings"),
+                quiet: bool = typer.Option(False, "--quiet", "-q", help="Enable quiet mode", rich_help_panel="Logging Settings"),
+                {strategy_settings},
+            ) -> None:
+                '''Run KRR using the `{func_name}` strategy'''
+
+                config = Config(
+                    clusters="*" if "*" in clusters else clusters,
+                    namespaces="*" if "*" in namespaces else namespaces,
+                    prometheus_url=prometheus_url,
+                    prometheus_auth_header=prometheus_auth_header,
+                    prometheus_ssl_enabled=prometheus_ssl_enabled,
+                    format=format,
+                    verbose=verbose,
+                    quiet=quiet,
+                    strategy="{func_name}",
+                    other_args=ctx.args,
+                )
+                runner = Runner(config)
+                asyncio.run(runner.run())
+            """
+        )
+
+        exec(
+            FUNC_TEMPLATE.format(
+                func_name=strategy_name,
+                strategy_name=strategy_type.__name__,
+                strategy_settings=",\n".join(
+                    f'{field_name}: {__process_type(field_meta.type_)} = typer.Option({field_meta.default!r}, "--{field_name}", help="{field_meta.field_info.description}", rich_help_panel="Strategy Settings")'
+                    for field_name, field_meta in strategy_type.get_settings_type().__fields__.items()
+                ),
+                formatters=", ".join(BaseFormatter.get_all()),
             ),
-            formatters=", ".join(BaseFormatter.get_all()),
-        ),
-        globals()
-        | {strategy.__name__: strategy for strategy in AnyStrategy.get_all().values()}  # Defined strategies
-        | {
-            "Runner": Runner,
-            "Config": Config,
-            "List": List,
-            "Optional": Optional,
-            "Union": Union,
-            "Literal": Literal,
-            "asyncio": asyncio,
-            "typer": typer,
-            # Required imports, here to make the linter happy (it doesn't know that exec will use them)
-        },
-        locals(),
-    )
+            globals()
+            | {strategy.__name__: strategy for strategy in AnyStrategy.get_all().values()}  # Defined strategies
+            | {
+                "Runner": Runner,
+                "Config": Config,
+                "List": List,
+                "Optional": Optional,
+                "Union": Union,
+                "Literal": Literal,
+                "asyncio": asyncio,
+                "typer": typer,
+                # Required imports, here to make the linter happy (it doesn't know that exec will use them)
+            },
+            locals(),
+        )
+
+    app()
 
 
 if __name__ == "__main__":
-    app()
+    run()