From 16e16bdb534f93154de844f6b72a2d40d03e6767 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9F=D0=B0=D0=B2=D0=B5=D0=BB=20=D0=96=D1=83=D0=BA=D0=BE?= =?UTF-8?q?=D0=B2?= <33721692+LeaveMyYard@users.noreply.github.com> Date: Thu, 9 Mar 2023 09:42:01 +0200 Subject: More refactoring, add loading of jobs data --- robusta_krr/core/abstract/formatters.py | 2 +- robusta_krr/core/abstract/strategies.py | 2 +- robusta_krr/core/integrations/kubernetes.py | 133 ++++++++++++++++++++++++++++ robusta_krr/core/integrations/prometheus.py | 27 ++++++ robusta_krr/core/models/kubernetes.py | 130 --------------------------- robusta_krr/core/models/result.py | 70 +++++++++++++++ robusta_krr/core/prometheus.py | 27 ------ robusta_krr/core/result.py | 70 --------------- robusta_krr/core/runner.py | 6 +- robusta_krr/formatters/json.py | 2 +- robusta_krr/formatters/table.py | 2 +- robusta_krr/formatters/yaml.py | 2 +- 12 files changed, 238 insertions(+), 235 deletions(-) create mode 100644 robusta_krr/core/integrations/kubernetes.py create mode 100644 robusta_krr/core/integrations/prometheus.py delete mode 100644 robusta_krr/core/models/kubernetes.py create mode 100644 robusta_krr/core/models/result.py delete mode 100644 robusta_krr/core/prometheus.py delete mode 100644 robusta_krr/core/result.py diff --git a/robusta_krr/core/abstract/formatters.py b/robusta_krr/core/abstract/formatters.py index 8487f7f..8316923 100644 --- a/robusta_krr/core/abstract/formatters.py +++ b/robusta_krr/core/abstract/formatters.py @@ -5,7 +5,7 @@ import os from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from robusta_krr.core.result import Result + from robusta_krr.core.models.result import Result DEFAULT_FORMATTERS_PATH = os.path.join(os.path.dirname(__file__), "formatters") diff --git a/robusta_krr/core/abstract/strategies.py b/robusta_krr/core/abstract/strategies.py index ebc23b4..51edf31 100644 --- a/robusta_krr/core/abstract/strategies.py +++ b/robusta_krr/core/abstract/strategies.py @@ -6,7 +6,7 @@ from typing import Generic, TypeVar import pydantic as pd -from robusta_krr.core.result import K8sObjectData, ResourceType +from robusta_krr.core.models.result import K8sObjectData, ResourceType class ResourceRecommendation(pd.BaseModel): diff --git a/robusta_krr/core/integrations/kubernetes.py b/robusta_krr/core/integrations/kubernetes.py new file mode 100644 index 0000000..10f85fd --- /dev/null +++ b/robusta_krr/core/integrations/kubernetes.py @@ -0,0 +1,133 @@ +import asyncio +import itertools + +from kubernetes import client, config +from kubernetes.client.models import ( + V1PodList, + V1DeploymentList, + V1StatefulSetList, + V1JobList, + V1DaemonSetList, + V1Deployment, + V1Container, + V1DaemonSet, + V1StatefulSet, +) + +from robusta_krr.core.models.objects import K8sObjectData +from robusta_krr.core.models.result import ResourceAllocations +from robusta_krr.utils.configurable import Configurable + + +class ClusterLoader(Configurable): + def __init__(self, cluster: str, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.cluster = cluster + self.apps = client.AppsV1Api(api_client=config.new_client_from_config(context=cluster)) + self.batch = client.BatchV1Api(api_client=config.new_client_from_config(context=cluster)) + + async def list_scannable_objects(self) -> list[K8sObjectData]: + """List all scannable objects. + + Returns: + A list of scannable objects. + """ + + self.debug("Listing scannable objects") + + try: + objects_tuple = await asyncio.gather( + self._list_deployments(), + self._list_all_statefulsets(), + self._list_all_daemon_set(), + self._list_all_jobs(), + ) + except Exception as e: + self.error(f"Error trying to list pods in cluster {self.cluster}: {e}") + return [] + + return list(itertools.chain(*objects_tuple)) + + def __build_obj(self, item: V1Deployment | V1DaemonSet | V1StatefulSet, container: V1Container) -> K8sObjectData: + return K8sObjectData( + cluster=self.cluster, + namespace=item.metadata.namespace, + name=item.metadata.name, + kind=item.__class__.__name__[2:], + container=container.name, + allocations=ResourceAllocations.from_container(container), + ) + + async def _list_deployments(self) -> list[K8sObjectData]: + ret: V1DeploymentList = await asyncio.to_thread(self.apps.list_deployment_for_all_namespaces, watch=False) + + return [ + self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers + ] + + async def _list_all_statefulsets(self) -> list[K8sObjectData]: + ret: V1StatefulSetList = await asyncio.to_thread(self.apps.list_stateful_set_for_all_namespaces, watch=False) + + return [ + self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers + ] + + async def _list_all_daemon_set(self) -> list[K8sObjectData]: + ret: V1DaemonSetList = await asyncio.to_thread(self.apps.list_daemon_set_for_all_namespaces, watch=False) + + return [ + self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers + ] + + async def _list_all_jobs(self) -> list[K8sObjectData]: + """Not working yet.""" + + ret: V1JobList = await asyncio.to_thread(self.batch.list_job_for_all_namespaces, watch=False) + + return [ + self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers + ] + + async def _list_pods(self) -> list[K8sObjectData]: + """For future use, not supported yet.""" + + ret: V1PodList = await asyncio.to_thread(self.apps.list_pod_for_all_namespaces, watch=False) + + return [self.__build_obj(item, container) for item in ret.items for container in item.spec.containers] + + +class KubernetesLoader(Configurable): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.debug("Initializing Kubernetes client") + config.load_kube_config() + + self._kubernetes_object_allocation_cache: dict[K8sObjectData, ResourceAllocations] = {} + + async def list_clusters(self) -> list[str]: + """List all clusters. + + Returns: + A list of clusters. + """ + + self.debug("Listing clusters") + + contexts, _ = await asyncio.to_thread(config.list_kube_config_contexts) + + return [context["name"] for context in contexts] + + async def list_scannable_objects(self, clusters: list[str]) -> list[K8sObjectData]: + """List all scannable objects. + + Returns: + A list of scannable objects. + """ + + self.debug("Listing scannable objects") + + cluster_loaders = [ClusterLoader(cluster=cluster, config=self.config) for cluster in clusters] + objects = await asyncio.gather(*[cluster_loader.list_scannable_objects() for cluster_loader in cluster_loaders]) + return list(itertools.chain(*objects)) diff --git a/robusta_krr/core/integrations/prometheus.py b/robusta_krr/core/integrations/prometheus.py new file mode 100644 index 0000000..64e2f08 --- /dev/null +++ b/robusta_krr/core/integrations/prometheus.py @@ -0,0 +1,27 @@ +import asyncio +import datetime +import random + +from robusta_krr.core.models.objects import K8sObjectData +from robusta_krr.core.models.result import ResourceType +from robusta_krr.core.abstract.strategies import HistoryData +from robusta_krr.utils.configurable import Configurable + + +class PrometheusLoader(Configurable): + async def gather_data( + self, + object: K8sObjectData, + resource: ResourceType, + period: datetime.timedelta, + *, + timeframe: datetime.timedelta = datetime.timedelta(minutes=1), + ) -> HistoryData: + # TODO: This is mock function. Implement this later using the Prometheus API + self.debug(f"Gathering data for {object} and {resource} for the last {period}") + await asyncio.sleep(1.5) # Simulate a slow API call + points = int(period / timeframe) + return { + "container_1": [random.randrange(30, 300) for _ in range(points)], + "container_2": [random.randrange(70, 500) for _ in range(points)], + } diff --git a/robusta_krr/core/models/kubernetes.py b/robusta_krr/core/models/kubernetes.py deleted file mode 100644 index 2c2cbc2..0000000 --- a/robusta_krr/core/models/kubernetes.py +++ /dev/null @@ -1,130 +0,0 @@ -import asyncio -import itertools - -from kubernetes import client, config -from kubernetes.client.models import ( - V1PodList, - V1DeploymentList, - V1StatefulSetList, - V1Deployment, - V1Container, - V1DaemonSet, - V1StatefulSet, -) - -from robusta_krr.core.models.objects import K8sObjectData -from robusta_krr.core.result import ResourceAllocations -from robusta_krr.utils.configurable import Configurable - - -class ClusterLoader(Configurable): - def __init__(self, cluster: str, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.cluster = cluster - self.v1 = client.AppsV1Api(api_client=config.new_client_from_config(context=cluster)) - - async def list_scannable_objects(self) -> list[K8sObjectData]: - """List all scannable objects. - - Returns: - A list of scannable objects. - """ - - self.debug("Listing scannable objects") - - try: - objects_tuple = await asyncio.gather( - self._list_deployments(), - self._list_all_statefulsets(), - self._list_all_daemon_set(), - # self._list_all_jobs(), # TODO: Add support for Jobs - ) - except Exception as e: - self.error(f"Error trying to list pods in cluster {self.cluster}: {e}") - return [] - - return list(itertools.chain(*objects_tuple)) - - def __build_obj(self, item: V1Deployment | V1DaemonSet | V1StatefulSet, container: V1Container) -> K8sObjectData: - return K8sObjectData( - cluster=self.cluster, - namespace=item.metadata.namespace, - name=item.metadata.name, - kind=item.__class__.__name__[2:], - container=container.name, - allocations=ResourceAllocations.from_container(container), - ) - - async def _list_deployments(self) -> list[K8sObjectData]: - ret: V1DeploymentList = await asyncio.to_thread(self.v1.list_deployment_for_all_namespaces, watch=False) - - return [ - self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers - ] - - async def _list_all_statefulsets(self) -> list[K8sObjectData]: - ret: V1StatefulSetList = await asyncio.to_thread(self.v1.list_stateful_set_for_all_namespaces, watch=False) - - return [ - self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers - ] - - async def _list_all_daemon_set(self) -> list[K8sObjectData]: - ret: V1StatefulSetList = await asyncio.to_thread(self.v1.list_daemon_set_for_all_namespaces, watch=False) - - return [ - self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers - ] - - async def _list_all_jobs(self) -> list[K8sObjectData]: - """Not working yet.""" - - ret: V1StatefulSetList = await asyncio.to_thread(self.v1.list_, watch=False) - - return [ - self.__build_obj(item, container) for item in ret.items for container in item.spec.template.spec.containers - ] - - async def _list_pods(self) -> list[K8sObjectData]: - """For future use, not supported yet.""" - - ret: V1PodList = await asyncio.to_thread(self.v1.list_pod_for_all_namespaces, watch=False) - - return [self.__build_obj(item, container) for item in ret.items for container in item.spec.containers] - - -class KubernetesLoader(Configurable): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.debug("Initializing Kubernetes client") - config.load_kube_config() - - self._kubernetes_object_allocation_cache: dict[K8sObjectData, ResourceAllocations] = {} - - async def list_clusters(self) -> list[str]: - """List all clusters. - - Returns: - A list of clusters. - """ - - self.debug("Listing clusters") - - contexts, _ = await asyncio.to_thread(config.list_kube_config_contexts) - - return [context["name"] for context in contexts] - - async def list_scannable_objects(self, clusters: list[str]) -> list[K8sObjectData]: - """List all scannable objects. - - Returns: - A list of scannable objects. - """ - - self.debug("Listing scannable objects") - - cluster_loaders = [ClusterLoader(cluster=cluster, config=self.config) for cluster in clusters] - objects = await asyncio.gather(*[cluster_loader.list_scannable_objects() for cluster_loader in cluster_loaders]) - return list(itertools.chain(*objects)) diff --git a/robusta_krr/core/models/result.py b/robusta_krr/core/models/result.py new file mode 100644 index 0000000..89ba72b --- /dev/null +++ b/robusta_krr/core/models/result.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import itertools +from typing import Any + +import pydantic as pd + +from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.models.objects import K8sObjectData +from robusta_krr.core.models.allocations import ResourceAllocations, ResourceType + + +class ResourceScan(pd.BaseModel): + object: K8sObjectData + recommended: ResourceAllocations + + +class Result(pd.BaseModel): + scans: list[ResourceScan] + score: float = 0.0 + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.score = self.__calculate_score() + + def format(self, formatter: type[BaseFormatter] | str, **kwargs: Any) -> Any: + """Format the result. + + Args: + formatter: The formatter to use. + + Returns: + The formatted result. + """ + + FormatterType = BaseFormatter.find(formatter) if isinstance(formatter, str) else formatter + _formatter = FormatterType(**kwargs) + return _formatter.format(self) + + @staticmethod + def __percentage_difference(current: float | str | None, recommended: float | str | None) -> float: + """Get the percentage difference between two numbers. + + Args: + current: The current value. + recommended: The recommended value. + + Returns: + The percentage difference. + """ + + return 1 + + def __calculate_score(self) -> float: + """Get the score of the result. + + Returns: + The score of the result. + """ + + total_diff = 0.0 + for scan, resource_type in itertools.product(self.scans, ResourceType): + total_diff += self.__percentage_difference( + scan.object.allocations.requests[resource_type], scan.recommended.requests[resource_type] + ) + total_diff += self.__percentage_difference( + scan.object.allocations.limits[resource_type], scan.recommended.limits[resource_type] + ) + + return max(0, round(100 - total_diff / len(self.scans) / len(ResourceType) / 50, 2)) # 50 is just a constant diff --git a/robusta_krr/core/prometheus.py b/robusta_krr/core/prometheus.py deleted file mode 100644 index d81aad4..0000000 --- a/robusta_krr/core/prometheus.py +++ /dev/null @@ -1,27 +0,0 @@ -import asyncio -import datetime -import random - -from robusta_krr.core.models.objects import K8sObjectData -from robusta_krr.core.result import ResourceType -from robusta_krr.core.abstract.strategies import HistoryData -from robusta_krr.utils.configurable import Configurable - - -class PrometheusLoader(Configurable): - async def gather_data( - self, - object: K8sObjectData, - resource: ResourceType, - period: datetime.timedelta, - *, - timeframe: datetime.timedelta = datetime.timedelta(minutes=1), - ) -> HistoryData: - # TODO: This is mock function. Implement this later using the Prometheus API - self.debug(f"Gathering data for {object} and {resource} for the last {period}") - await asyncio.sleep(1.5) # Simulate a slow API call - points = int(period / timeframe) - return { - "container_1": [random.randrange(30, 300) for _ in range(points)], - "container_2": [random.randrange(70, 500) for _ in range(points)], - } diff --git a/robusta_krr/core/result.py b/robusta_krr/core/result.py deleted file mode 100644 index 89ba72b..0000000 --- a/robusta_krr/core/result.py +++ /dev/null @@ -1,70 +0,0 @@ -from __future__ import annotations - -import itertools -from typing import Any - -import pydantic as pd - -from robusta_krr.core.abstract.formatters import BaseFormatter -from robusta_krr.core.models.objects import K8sObjectData -from robusta_krr.core.models.allocations import ResourceAllocations, ResourceType - - -class ResourceScan(pd.BaseModel): - object: K8sObjectData - recommended: ResourceAllocations - - -class Result(pd.BaseModel): - scans: list[ResourceScan] - score: float = 0.0 - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.score = self.__calculate_score() - - def format(self, formatter: type[BaseFormatter] | str, **kwargs: Any) -> Any: - """Format the result. - - Args: - formatter: The formatter to use. - - Returns: - The formatted result. - """ - - FormatterType = BaseFormatter.find(formatter) if isinstance(formatter, str) else formatter - _formatter = FormatterType(**kwargs) - return _formatter.format(self) - - @staticmethod - def __percentage_difference(current: float | str | None, recommended: float | str | None) -> float: - """Get the percentage difference between two numbers. - - Args: - current: The current value. - recommended: The recommended value. - - Returns: - The percentage difference. - """ - - return 1 - - def __calculate_score(self) -> float: - """Get the score of the result. - - Returns: - The score of the result. - """ - - total_diff = 0.0 - for scan, resource_type in itertools.product(self.scans, ResourceType): - total_diff += self.__percentage_difference( - scan.object.allocations.requests[resource_type], scan.recommended.requests[resource_type] - ) - total_diff += self.__percentage_difference( - scan.object.allocations.limits[resource_type], scan.recommended.limits[resource_type] - ) - - return max(0, round(100 - total_diff / len(self.scans) / len(ResourceType) / 50, 2)) # 50 is just a constant diff --git a/robusta_krr/core/runner.py b/robusta_krr/core/runner.py index 1029750..eb70f16 100644 --- a/robusta_krr/core/runner.py +++ b/robusta_krr/core/runner.py @@ -2,10 +2,10 @@ import asyncio import itertools from robusta_krr.core.models.config import Config -from robusta_krr.core.models.kubernetes import KubernetesLoader +from robusta_krr.core.integrations.kubernetes import KubernetesLoader from robusta_krr.core.models.objects import K8sObjectData -from robusta_krr.core.prometheus import PrometheusLoader -from robusta_krr.core.result import ResourceAllocations, ResourceScan, ResourceType, Result +from robusta_krr.core.integrations.prometheus import PrometheusLoader +from robusta_krr.core.models.result import ResourceAllocations, ResourceScan, ResourceType, Result from robusta_krr.core.abstract.strategies import ResourceRecommendation from robusta_krr.utils.configurable import Configurable from robusta_krr.utils.version import get_version diff --git a/robusta_krr/formatters/json.py b/robusta_krr/formatters/json.py index 650c292..2c3a51e 100644 --- a/robusta_krr/formatters/json.py +++ b/robusta_krr/formatters/json.py @@ -1,7 +1,7 @@ from __future__ import annotations from robusta_krr.core.abstract.formatters import BaseFormatter -from robusta_krr.core.result import Result +from robusta_krr.core.models.result import Result class JSONFormatter(BaseFormatter): diff --git a/robusta_krr/formatters/table.py b/robusta_krr/formatters/table.py index 1f1439b..2fcd16f 100644 --- a/robusta_krr/formatters/table.py +++ b/robusta_krr/formatters/table.py @@ -3,7 +3,7 @@ from __future__ import annotations import itertools from robusta_krr.core.abstract.formatters import BaseFormatter -from robusta_krr.core.result import Result, ResourceType +from robusta_krr.core.models.result import Result, ResourceType from robusta_krr.utils import resource_units from rich.table import Table diff --git a/robusta_krr/formatters/yaml.py b/robusta_krr/formatters/yaml.py index a493fa1..24be2d6 100644 --- a/robusta_krr/formatters/yaml.py +++ b/robusta_krr/formatters/yaml.py @@ -1,7 +1,7 @@ from __future__ import annotations from robusta_krr.core.abstract.formatters import BaseFormatter -from robusta_krr.core.result import Result +from robusta_krr.core.models.result import Result import yaml import json -- cgit v1.2.3