diff options
| -rw-r--r-- | examples/custom_strategy.py | 48 | ||||
| -rw-r--r-- | krr.py | 4 | ||||
| -rw-r--r-- | robusta_krr/__init__.py | 4 | ||||
| -rw-r--r-- | robusta_krr/main.py | 179 |
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() @@ -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() |