Fix custom strategies

examples/custom_strategy.py

@@ -1,11 +1,17 @@
 # This is an example on how to create your own custom strategy
-
 import pydantic as pd
+from typing import List, Optional
+
+import asyncio
+import typer
+from pydantic import ValidationError
 
-import robusta_krr
 from robusta_krr.api.models import K8sObjectData, MetricsPodData, ResourceRecommendation, ResourceType, RunResult
+from robusta_krr.core.models.config import Config, option_kubeconfig, option_clusters, option_all_clusters, option_namespaces, option_resources, option_selector, option_prometheus_url, option_prometheus_auth_header, option_prometheus_other_headers, option_prometheus_ssl_enabled, option_prometheus_cluster_label, option_prometheus_label, option_eks_managed_prom, option_eks_managed_prom_profile_name, option_eks_access_key, option_eks_secret_key, option_eks_service_name, option_eks_managed_prom_region, option_coralogix_token, option_openshift, option_max_workers, option_format, option_verbose, option_quiet, option_log_to_stderr, option_width, option_file_output, option_slack_output, option_cpu_min_value, option_memory_min_value
 from robusta_krr.api.strategies import BaseStrategy, StrategySettings
 from robusta_krr.core.integrations.prometheus.metrics import MaxMemoryLoader, PercentileCPULoader
+from robusta_krr.main import app, logger
+from robusta_krr.core.runner import Runner
 
 
 # Providing description to the settings will make it available in the CLI help
@@ -14,12 +20,16 @@ class CustomStrategySettings(StrategySettings):
     param_2: float = pd.Field(105_000, gt=0, description="Second example parameter")
 
 
-class CustomStrategy(BaseStrategy[CustomStrategySettings]):
+class CustomStrategy(BaseStrategy):
     """
     A custom strategy that uses the provided parameters for CPU and memory.
     Made only in order to demonstrate how to create a custom strategy.
     """
 
+    def __init__(self, settings: CustomStrategySettings):
+        super().__init__(settings=settings)
+        self.settings = settings
+
     display_name = "custom"  # The name of the strategy
     rich_console = True  # Whether to use rich console for the CLI
     metrics = [PercentileCPULoader(90), MaxMemoryLoader]  # The metrics to use for the strategy
@@ -31,7 +41,101 @@ def run(self, history_data: MetricsPodData, object_data: K8sObjectData) -> RunRe
         }
 
 
+# add a new command - much of this is boilerplate
+@app.command(rich_help_panel="Strategies")
+def custom(
+    kubeconfig: Optional[str] = option_kubeconfig,
+    clusters: List[str] = option_clusters,
+    all_clusters: bool = option_all_clusters,
+    namespaces: List[str] = option_namespaces,
+    resources: List[str] = option_resources,
+    selector: Optional[str] = option_selector,
+    prometheus_url: Optional[str] = option_prometheus_url,
+    prometheus_auth_header: Optional[str] = option_prometheus_auth_header,
+    prometheus_other_headers: Optional[List[str]]= option_prometheus_other_headers,
+    prometheus_ssl_enabled: bool = option_prometheus_ssl_enabled,
+    prometheus_cluster_label: Optional[str] = option_prometheus_cluster_label,
+    prometheus_label: str = option_prometheus_label,
+    eks_managed_prom: bool = option_eks_managed_prom,
+    eks_managed_prom_profile_name: Optional[str] = option_eks_managed_prom_profile_name,
+    eks_access_key: Optional[str] = option_eks_access_key,
+    eks_secret_key: Optional[str] = option_eks_secret_key,
+    eks_service_name: Optional[str] = option_eks_service_name,
+    eks_managed_prom_region: Optional[str] = option_eks_managed_prom_region,
+    coralogix_token: Optional[str] = option_coralogix_token,
+    openshift: bool = option_openshift,
+    max_workers: int = option_max_workers,
+    format: str = option_format,
+    verbose: bool = option_verbose,
+    quiet: bool = option_quiet,
+    log_to_stderr: bool = option_log_to_stderr,
+    width: Optional[int] = option_width,
+    file_output: Optional[str] = option_file_output,
+    slack_output: Optional[str] = option_slack_output,
+    cpu_min_value: int = option_cpu_min_value,
+    memory_min_value: int = option_memory_min_value,
+
+    # define parameters that are specific to this strategy
+    param1: float = typer.Option(
+        None,
+        "--param1",
+        help="Explanation of param1",
+        rich_help_panel="Strategy Settings",
+    ),
+    param2: float = typer.Option(
+        None,
+        "--param2",
+        help="Explanation of param2",
+        rich_help_panel="Strategy Settings",
+    )
+) -> None:
+    """A custom strategy for KRR"""
+    try:
+        config = Config(
+            kubeconfig=kubeconfig,
+            clusters="*" if all_clusters else clusters,
+            namespaces="*" if "*" in namespaces else namespaces,
+            resources="*" if "*" in resources else resources,
+            selector=selector,
+            prometheus_url=prometheus_url,
+            prometheus_auth_header=prometheus_auth_header,
+            prometheus_other_headers=prometheus_other_headers,
+            prometheus_ssl_enabled=prometheus_ssl_enabled,
+            prometheus_cluster_label=prometheus_cluster_label,
+            prometheus_label=prometheus_label,
+            eks_managed_prom=eks_managed_prom,
+            eks_managed_prom_region=eks_managed_prom_region,
+            eks_managed_prom_profile_name=eks_managed_prom_profile_name,
+            eks_access_key=eks_access_key,
+            eks_secret_key=eks_secret_key,
+            eks_service_name=eks_service_name,
+            coralogix_token=coralogix_token,
+            openshift=openshift,
+            max_workers=max_workers,
+            format=format,
+            verbose=verbose,
+            cpu_min_value=cpu_min_value,
+            memory_min_value=memory_min_value,
+            quiet=quiet,
+            log_to_stderr=log_to_stderr,
+            width=width,
+            file_output=file_output,
+            slack_output=slack_output,
+        )
+        Config.set_config(config)
+        strategy_settings = CustomStrategySettings(
+            param_1=param1,
+            param_2=param2
+        )
+    except ValidationError:
+        logger.exception("Error occured while parsing arguments")
+    else:
+        strategy = CustomStrategy(strategy_settings)
+        runner = Runner(strategy)
+        asyncio.run(runner.run())
+
+
 # Running this file will register the strategy and make it available to the CLI
-# Run it as `python ./custom_strategy.py my_strategy`
+# Run it as `python ./custom_strategy.py custom --param1 2.0 --param2 1.0`
 if __name__ == "__main__":
-    robusta_krr.run()
+    app()

robusta_krr/core/abstract/strategies.py

@@ -107,13 +107,12 @@ def __str__(self) -> str:
         return self.display_name.title()
 
     @property
-    @abc.abstractmethod
     def description(self) -> Optional[str]:
         """
         Generate a description for the strategy.
         You can use Rich's markdown syntax to format the description.
         """
-        pass
+        return ""
 
     # 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.

robusta_krr/core/models/config.py

@@ -2,16 +2,17 @@
 
 import logging
 import sys
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union, List
 
 import pydantic as pd
+import typer
 from kubernetes import config
 from kubernetes.config.config_exception import ConfigException
+from pydantic.fields import ModelField
 from rich.console import Console
 from rich.logging import RichHandler
 
 from robusta_krr.core.abstract import formatters
-from robusta_krr.core.abstract.strategies import BaseStrategy
 from robusta_krr.core.models.objects import KindLiteral
 
 logger = logging.getLogger("krr")
@@ -52,7 +53,6 @@ class Config(pd.BaseSettings):
 
     # Logging Settings
     format: str
-    strategy: str
     log_to_stderr: bool
     width: Optional[int] = pd.Field(None, ge=1)
 
@@ -164,3 +164,199 @@ def __getattr__(self, name: str):
 
 _config: Optional[Config] = None
 settings = _Settings()
+
+
+# Helper for defining command line options
+def pydantic_field_to_typer_option(option_names: List[str], rich_help_panel: str, field: ModelField) -> typer.Option:
+    """
+    Create a typer option from a pydantic field.
+    We use this to generate cli options with default values and help text, without repeating ourselves.
+    """
+    return typer.Option(
+        field.default,
+        *option_names,
+        help=field.field_info.description,
+        rich_help_panel=rich_help_panel,
+    )
+
+# Common command line options
+# For new CLI options, use "--dashes-like-this" and "--not_undersores_like_this"
+# For backwards compatibility, some CLI options support both styles  
+option_kubeconfig: Optional[str] = typer.Option(
+    None,
+    "--kubeconfig",
+    "-k",
+    help="Path to kubeconfig file. If not provided, will attempt to find it.",
+    rich_help_panel="Kubernetes Settings",
+)
+option_clusters: List[str] = typer.Option(
+    None,
+    "--context",
+    "--cluster",
+    "-c",
+    help="List of clusters to run on. By default, will run on the current cluster. Use --all-clusters to run on all clusters.",
+    rich_help_panel="Kubernetes Settings",
+)
+option_all_clusters: bool = typer.Option(
+    False,
+    "--all-clusters",
+    help="Run on all clusters. Overrides --context.",
+    rich_help_panel="Kubernetes Settings",
+)
+option_namespaces: List[str] = typer.Option(
+    None,
+    "--namespace",
+    "-n",
+    help="List of namespaces to run on. By default, will run on all namespaces except 'kube-system'.",
+    rich_help_panel="Kubernetes Settings",
+)
+option_resources: List[str] = typer.Option(
+    None,
+    "--resource",
+    "-r",
+    help="List of resources to run on (Deployment, StatefulSet, DaemonSet, Job, Rollout). By default, will run on all resources. Case insensitive.",
+    rich_help_panel="Kubernetes Settings",
+)
+option_selector: Optional[str] = typer.Option(
+    None,
+    "--selector",
+    "-s",
+    help="Selector (label query) to filter on, supports '=', '==', and '!='.(e.g. -s key1=value1,key2=value2). Matching objects must satisfy all of the specified label constraints.",
+    rich_help_panel="Kubernetes Settings",
+)
+option_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",
+)
+option_prometheus_auth_header: Optional[str] = typer.Option(
+    None,
+    "--prometheus-auth-header",
+    help="Prometheus authentication header.",
+    rich_help_panel="Prometheus Settings",
+)
+option_prometheus_other_headers: Optional[List[str]] = typer.Option(
+    None,
+    "--prometheus-headers",
+    "-H",
+    help="Additional headers to add to Prometheus requests. Format as 'key: value', for example 'X-MyHeader: 123'. Trailing whitespaces will be stripped.",
+    rich_help_panel="Prometheus Settings",
+)
+option_prometheus_ssl_enabled: bool = typer.Option(
+    False,
+    "--prometheus-ssl-enabled",
+    help="Enable SSL for Prometheus requests.",
+    rich_help_panel="Prometheus Settings",
+)
+option_prometheus_cluster_label: Optional[str] = typer.Option(
+    None,
+    "--prometheus-cluster-label",
+    "-l",
+    help="The label in prometheus for your cluster.(Only relevant for centralized prometheus)",
+    rich_help_panel="Prometheus Settings",
+)
+option_prometheus_label: str = typer.Option(
+    None,
+    "--prometheus-label",
+    help="The label in prometheus used to differentiate clusters. (Only relevant for centralized prometheus)",
+    rich_help_panel="Prometheus Settings",
+)
+option_eks_managed_prom: bool = typer.Option(
+    False,
+    "--eks-managed-prom",
+    help="Adds additional signitures for eks prometheus connection.",
+    rich_help_panel="Prometheus EKS Settings",
+)
+option_eks_managed_prom_profile_name: Optional[str] = typer.Option(
+    None,
+    "--eks-profile-name",
+    help="Sets the profile name for eks prometheus connection.",
+    rich_help_panel="Prometheus EKS Settings",
+)
+option_eks_access_key: Optional[str] = typer.Option(
+    None,
+    "--eks-access-key",
+    help="Sets the access key for eks prometheus connection.",
+    rich_help_panel="Prometheus EKS Settings",
+)
+option_eks_secret_key: Optional[str] = typer.Option(
+    None,
+    "--eks-secret-key",
+    help="Sets the secret key for eks prometheus connection.",
+    rich_help_panel="Prometheus EKS Settings",
+)
+option_eks_service_name: Optional[str] = typer.Option(
+    "aps",
+    "--eks-service-name",
+    help="Sets the service name for eks prometheus connection.",
+    rich_help_panel="Prometheus EKS Settings",
+)
+option_eks_managed_prom_region: Optional[str] = typer.Option(
+    None,
+    "--eks-managed-prom-region",
+    help="Sets the region for eks prometheus connection.",
+    rich_help_panel="Prometheus EKS Settings",
+)
+option_coralogix_token: Optional[str] = typer.Option(
+    None,
+    "--coralogix-token",
+    help="Adds the token needed to query Coralogix managed prometheus.",
+    rich_help_panel="Prometheus Coralogix Settings",
+)
+option_openshift: bool = typer.Option(
+    False,
+    "--openshift",
+    help="Used when running by Robusta inside an OpenShift cluster.",
+    rich_help_panel="Prometheus Openshift Settings",
+    hidden=True,
+)
+option_max_workers: int = typer.Option(
+    10,
+    "--max-workers",
+    "-w",
+    help="Max workers to use for async requests.",
+    rich_help_panel="Threading Settings",
+)
+option_format: str = typer.Option(
+    "table",
+    "--formatter",
+    "-f",
+    help=f"Output formatter ({', '.join(formatters.list_available())})",
+    rich_help_panel="Logging Settings",
+)
+option_verbose: bool = typer.Option(
+    False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings"
+)
+option_quiet: bool = typer.Option(
+    False, "--quiet", "-q", help="Enable quiet mode", rich_help_panel="Logging Settings"
+)
+option_log_to_stderr: bool = typer.Option(
+    False, "--logtostderr", help="Pass logs to stderr", rich_help_panel="Logging Settings"
+)
+option_width: Optional[int] = typer.Option(
+    None,
+    "--width",
+    help="Width of the output. Will use console width by default.",
+    rich_help_panel="Logging Settings",
+)
+option_file_output: Optional[str] = typer.Option(
+    None, "--fileoutput", help="Print the output to a file", rich_help_panel="Output Settings"
+)
+option_slack_output: Optional[str] = typer.Option(
+    None,
+    "--slackoutput",
+    help="Send to output to a slack channel, must have SLACK_BOT_TOKEN",
+    rich_help_panel="Output Settings",
+)
+option_cpu_min_value: int = pydantic_field_to_typer_option(
+    ["--cpu-min", "--cpu_min"],
+    "Strategy Settings",
+    Config.__fields__["cpu_min_value"]
+)
+option_memory_min_value: int = pydantic_field_to_typer_option(
+    ["--mem-min", "--mem_min"],
+    "Strategy Settings",
+    Config.__fields__["memory_min_value"]
+)