feat: Enable resource requests tracking (#1207)

* Record certain data in a `resource_request_log table`
* Snapshots are stored periodically from data_tasks
* Resource request tracking is disabled by default, requires `ENABLE_RESOURCE_REQUEST_TRACKING=true` to be in the env

Author: eikek

Makefile

@@ -53,7 +53,8 @@ API_SPECS := \
     components/renku_data_services/data_connectors/apispec.py \
     components/renku_data_services/search/apispec.py \
     components/renku_data_services/notifications/apispec.py \
-    components/renku_data_services/capacity_reservation/apispec.py
+    components/renku_data_services/capacity_reservation/apispec.py \
+    components/renku_data_services/resource_usage/apispec.py
 
 schemas: ${API_SPECS}  ## Generate pydantic classes from apispec yaml files
 	@echo "generated classes based on ApiSpec"

bases/renku_data_services/data_api/app.py

@@ -31,6 +31,7 @@
 from renku_data_services.platform.blueprints import PlatformConfigBP, PlatformUrlRedirectBP
 from renku_data_services.project.blueprints import ProjectsBP, ProjectSessionSecretBP
 from renku_data_services.repositories.blueprints import RepositoriesBP
+from renku_data_services.resource_usage.blueprints import ResourceUsageBP
 from renku_data_services.search.blueprints import SearchBP
 from renku_data_services.search.reprovision import SearchReprovision
 from renku_data_services.search.solr_user_query import UsernameResolve
@@ -273,6 +274,13 @@ def register_all_handlers(app: Sanic, dm: DependencyManager) -> Sanic:
         occurrence_repo=dm.occurrence_repo,
         authenticator=dm.authenticator,
     )
+    resource_usage = ResourceUsageBP(
+        name="resource_usage",
+        url_prefix=url_prefix,
+        rr_repo=dm.resource_requests_repo,
+        rr_svc=dm.resource_usage_service,
+        authenticator=dm.authenticator,
+    )
     app.blueprint(
         [
             resource_pools.blueprint(),
@@ -302,6 +310,7 @@ def register_all_handlers(app: Sanic, dm: DependencyManager) -> Sanic:
             platform_redirects.blueprint(),
             notifications.blueprint(),
             capacity_reservation.blueprint(),
+            resource_usage.blueprint(),
         ]
     )
     if builds is not None:

bases/renku_data_services/data_api/dependencies.py

@@ -37,6 +37,7 @@
     DataConnectorSecretRepository,
 )
 from renku_data_services.git.gitlab import DummyGitlabAPI, EmptyGitlabAPI, GitlabAPI
+from renku_data_services.k8s.client_interfaces import K8sClient
 from renku_data_services.k8s.clients import (
     K8sClusterClientsPool,
     K8sPriorityClassClient,
@@ -62,6 +63,8 @@
     ProjectSessionSecretRepository,
 )
 from renku_data_services.repositories.db import GitRepositoriesRepository
+from renku_data_services.resource_usage.core import ResourceUsageService
+from renku_data_services.resource_usage.db import ResourceRequestsRepo
 from renku_data_services.search import query_manual
 from renku_data_services.search.db import SearchUpdatesRepo
 from renku_data_services.search.reprovision import SearchReprovision
@@ -118,6 +121,7 @@ class DependencyManager:
 
     config: Config
 
+    k8s_client: K8sClient
     user_store: base_models.UserStore
     authenticator: base_models.Authenticator
     gitlab_authenticator: base_models.Authenticator
@@ -158,6 +162,8 @@ class DependencyManager:
     oauth_http_client_factory: OAuthHttpClientFactory
     capacity_reservation_repo: CapacityReservationRepository
     occurrence_repo: OccurrenceRepository
+    resource_requests_repo: ResourceRequestsRepo
+    resource_usage_service: ResourceUsageService
 
     spec: dict[str, Any] = field(init=False, repr=False, default_factory=dict)
     app_name: str = "renku_data_services"
@@ -187,6 +193,7 @@ def load_apispec() -> dict[str, Any]:
             renku_data_services.search.__file__,
             renku_data_services.notifications.__file__,
             renku_data_services.capacity_reservation.__file__,
+            renku_data_services.resource_usage.__file__,
         ]
 
         api_specs = []
@@ -417,8 +424,13 @@ def from_env(cls) -> DependencyManager:
         occurrence_repo = OccurrenceRepository(
             session_maker=config.db.async_session_maker,
         )
+        resource_requests_repo = ResourceRequestsRepo(
+            session_maker=config.db.async_session_maker,
+        )
+        resource_usage_service = ResourceUsageService(repo=resource_requests_repo)
         return cls(
             config,
+            k8s_client=client,
             authenticator=authenticator,
             gitlab_authenticator=gitlab_authenticator,
             gitlab_client=gitlab_client,
@@ -459,4 +471,6 @@ def from_env(cls) -> DependencyManager:
             oauth_http_client_factory=oauth_http_client_factory,
             capacity_reservation_repo=capacity_reservation_repo,
             occurrence_repo=occurrence_repo,
+            resource_requests_repo=resource_requests_repo,
+            resource_usage_service=resource_usage_service,
         )

bases/renku_data_services/data_tasks/config.py

@@ -42,6 +42,7 @@ class Config:
     posthog: PosthogConfig
     authz: AuthzConfig
     keycloak: KeycloakConfig | None
+    k8s_config_root: str
     dummy_stores: bool
     max_retry_wait_seconds: int
     main_log_interval_seconds: int
@@ -50,7 +51,7 @@ class Config:
     x_short_task_period_s: int
     short_task_period_s: int
     long_task_period_s: int
-    k8s_config_root: str
+    enable_resource_request_tracking: bool
 
     @classmethod
     def from_env(cls) -> Config:
@@ -71,6 +72,7 @@ def from_env(cls) -> Config:
 
         k8s_config_root = os.environ.get("K8S_CONFIG_ROOT", "/secrets/kube_configs")
 
+        enable_resource_request_tracking = os.environ.get("ENABLE_RESOURCE_REQUEST_TRACKING", "false").lower() == "true"
         authz = AuthzConfig.from_env()
 
         keycloak = None if dummy_stores else KeycloakConfig.from_env()
@@ -82,11 +84,12 @@ def from_env(cls) -> Config:
             posthog=posthog_config,
             authz=authz,
             keycloak=keycloak,
+            k8s_config_root=k8s_config_root,
             tcp_host=tcp_host,
             tcp_port=tcp_port,
             x_short_task_period_s=x_short_task_period,
             short_task_period_s=short_task_period,
             long_task_period_s=long_task_period,
-            k8s_config_root=k8s_config_root,
             dummy_stores=dummy_stores,
+            enable_resource_request_tracking=enable_resource_request_tracking,
         )

bases/renku_data_services/data_tasks/dependencies.py

@@ -2,21 +2,28 @@
 
 from dataclasses import dataclass
 
+from renku_data_services.app_config import logging
 from renku_data_services.authz.authz import Authz
 from renku_data_services.capacity_reservation.db import CapacityReservationRepository, OccurrenceRepository
 from renku_data_services.capacity_reservation.k8s_client import CapacityReservationK8sClient
 from renku_data_services.capacity_reservation.tasks import CapacityReservationTasks
 from renku_data_services.crc.db import ClusterRepository
 from renku_data_services.data_tasks.config import Config
 from renku_data_services.k8s.clients import K8sClusterClientsPool
-from renku_data_services.k8s.config import KubeConfigEnv
+from renku_data_services.k8s.config import KubeConfigEnv, get_clusters
 from renku_data_services.k8s.db import K8sDbCache
 from renku_data_services.metrics.core import StagingMetricsService
 from renku_data_services.metrics.db import MetricsRepository
 from renku_data_services.namespace.db import GroupRepository
-from renku_data_services.notebooks.config import get_clusters
 from renku_data_services.notebooks.constants import AMALTHEA_SESSION_GVK
 from renku_data_services.project.db import ProjectRepository
+from renku_data_services.resource_usage.core import (
+    DefaultResourcesRequestRecorder,
+    NoopResourcesRequestRecorder,
+    ResourceRequestsFetch,
+    ResourcesRequestRecorder,
+)
+from renku_data_services.resource_usage.db import ResourceRequestsRepo
 from renku_data_services.search.db import SearchUpdatesRepo
 from renku_data_services.session.db import SessionRepository
 from renku_data_services.session.tasks import SessionTasks
@@ -25,6 +32,8 @@
 from renku_data_services.users.kc_api import IKeycloakAPI, KeycloakAPI
 from renku_data_services.users.models import UnsavedUserInfo
 
+logger = logging.getLogger(__file__)
+
 
 @dataclass
 class DependencyManager:
@@ -40,6 +49,7 @@ class DependencyManager:
     kc_api: IKeycloakAPI
     session_tasks: SessionTasks
     capacity_reservation_tasks: CapacityReservationTasks
+    resource_requests_recorder: ResourcesRequestRecorder
 
     @classmethod
     def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
@@ -83,10 +93,11 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
         session_tasks = SessionTasks(session_environment_repo=session_environment_repo)
         cluster_repo = ClusterRepository(session_maker=cfg.db.async_session_maker)
         k8s_db_cache = K8sDbCache(cfg.db.async_session_maker)
+        default_kubeconfig = KubeConfigEnv()
         k8s_client = K8sClusterClientsPool(
             lambda: get_clusters(
                 kube_conf_root_dir=cfg.k8s_config_root,
-                default_kubeconfig=KubeConfigEnv(),
+                default_kubeconfig=default_kubeconfig,
                 cluster_repo=cluster_repo,
                 cache=k8s_db_cache,
                 kinds_to_cache=[AMALTHEA_SESSION_GVK],
@@ -100,6 +111,16 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
             ),
             k8s_client=cr_k8s_client,
         )
+
+        resource_requests_recorder: ResourcesRequestRecorder
+        if cfg.enable_resource_request_tracking:
+            resource_requests_recorder = DefaultResourcesRequestRecorder(
+                repo=ResourceRequestsRepo(cfg.db.async_session_maker), fetch=ResourceRequestsFetch(k8s_client)
+            )
+        else:
+            logger.warning("Resource request tracking is disabled!")
+            resource_requests_recorder = NoopResourcesRequestRecorder()
+
         kc_api: IKeycloakAPI
         if cfg.dummy_stores:
             dummy_users = [
@@ -127,4 +148,5 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
             kc_api=kc_api,
             session_tasks=session_tasks,
             capacity_reservation_tasks=capacity_reservation_tasks,
+            resource_requests_recorder=resource_requests_recorder,
         )

bases/renku_data_services/data_tasks/task_defs.py

@@ -1,6 +1,7 @@
 """The task definitions in form of coroutines."""
 
 import asyncio
+from datetime import timedelta
 
 from authzed.api.v1 import (
     Consistency,
@@ -433,6 +434,14 @@ async def cleanup_orphaned_capacity_reservations(dm: DependencyManager) -> None:
             await asyncio.sleep(dm.config.x_short_task_period_s)
 
 
+async def record_resource_requests(dm: DependencyManager) -> None:
+    """Periodically record all resource requests."""
+    interval_seconds = 600
+    while True:
+        await dm.resource_requests_recorder.record_resource_requests(timedelta(seconds=interval_seconds))
+        await asyncio.sleep(interval_seconds)
+
+
 def all_tasks(dm: DependencyManager) -> TaskDefininions:
     """A dict of task factories to be managed in main."""
     # Impl. note: We pass the entire config to the coroutines, because
@@ -457,5 +466,6 @@ def all_tasks(dm: DependencyManager) -> TaskDefininions:
             "activate_capacity_reservations": lambda: activate_capacity_reservations(dm),
             "monitor_capacity_reservations": lambda: monitor_capacity_reservations(dm),
             "cleanup_orphaned_capacity_reservations": lambda: cleanup_orphaned_capacity_reservations(dm),
+            "record_resource_requests": lambda: record_resource_requests(dm),
         }
     )

components/renku_data_services/migrations/env.py

@@ -13,6 +13,7 @@
 from renku_data_services.notifications.orm import BaseORM as notifications
 from renku_data_services.platform.orm import BaseORM as platform
 from renku_data_services.project.orm import BaseORM as project
+from renku_data_services.resource_usage.orm import BaseORM as resource_usage
 from renku_data_services.search.orm import BaseORM as search
 from renku_data_services.secrets.orm import BaseORM as secrets
 from renku_data_services.session.orm import BaseORM as sessions
@@ -37,6 +38,7 @@
     sessions.metadata,
     storage.metadata,
     users.metadata,
+    resource_usage.metadata,
 ]
 
 run_migrations(all_metadata)

components/renku_data_services/migrations/utils.py

@@ -30,7 +30,7 @@ def include_object(
     compare_to: SchemaItem | None,
 ) -> bool:
     """Prevents from alembic migrating the alembic_version tables."""
-    return type_ != "table" or name != "alembic_version"
+    return type_ != "table" or (name != "alembic_version" and name != "resource_requests_view")
 
 
 def combine_version_tables(conn: Connection, metadata_schema: str | None) -> None: