feat: task to check session quota and send alerts

Author: mohammad-alisafaee

bases/renku_data_services/data_tasks/config.py

@@ -52,6 +52,7 @@ class Config:
     short_task_period_s: int
     long_task_period_s: int
     enable_resource_request_tracking: bool
+    session_quota_alert_check_interval_s: int
 
     @classmethod
     def from_env(cls) -> Config:
@@ -69,6 +70,7 @@ def from_env(cls) -> Config:
         x_short_task_period = int(os.environ.get("X_SHORT_TASK_PERIOD_S", 30))
         short_task_period = int(os.environ.get("SHORT_TASK_PERIOD_S", 2 * 60))
         long_task_period = int(os.environ.get("LONG_TASK_PERIOD_S", 3 * 60 * 60))
+        session_quota_alert_check_interval = int(os.environ.get("SESSION_QUOTA_ALERT_CHECK_INTERVAL_S", 5 * 60))
 
         k8s_config_root = os.environ.get("K8S_CONFIG_ROOT", "/secrets/kube_configs")
 
@@ -92,4 +94,5 @@ def from_env(cls) -> Config:
             long_task_period_s=long_task_period,
             dummy_stores=dummy_stores,
             enable_resource_request_tracking=enable_resource_request_tracking,
+            session_quota_alert_check_interval_s=session_quota_alert_check_interval,
         )

bases/renku_data_services/data_tasks/dependencies.py

@@ -1,5 +1,6 @@
 """Dependency management for data tasks."""
 
+import os
 from dataclasses import dataclass
 
 from renku_data_services.app_config import logging
@@ -16,12 +17,14 @@
 from renku_data_services.metrics.db import MetricsRepository
 from renku_data_services.namespace.db import GroupRepository
 from renku_data_services.notebooks.constants import AMALTHEA_SESSION_GVK
+from renku_data_services.notifications.db import NotificationsRepository
 from renku_data_services.project.db import ProjectRepository
 from renku_data_services.resource_usage.core import (
     DefaultResourcesRequestRecorder,
     NoopResourcesRequestRecorder,
     ResourceRequestsFetch,
     ResourcesRequestRecorder,
+    ResourceUsageService,
 )
 from renku_data_services.resource_usage.db import ResourceRequestsRepo
 from renku_data_services.search.db import SearchUpdatesRepo
@@ -50,6 +53,9 @@ class DependencyManager:
     session_tasks: SessionTasks
     capacity_reservation_tasks: CapacityReservationTasks
     resource_requests_recorder: ResourcesRequestRecorder
+    k8s_client: K8sClusterClientsPool
+    notifications_repo: NotificationsRepository
+    resource_usage_service: ResourceUsageService
 
     @classmethod
     def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
@@ -112,15 +118,24 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
             k8s_client=cr_k8s_client,
         )
 
+        resource_requests_repo = ResourceRequestsRepo(cfg.db.async_session_maker)
+        resource_usage_service = ResourceUsageService(repo=resource_requests_repo)
+
         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)
+                repo=resource_requests_repo, fetch=ResourceRequestsFetch(k8s_client)
             )
         else:
             logger.warning("Resource request tracking is disabled!")
             resource_requests_recorder = NoopResourcesRequestRecorder()
 
+        notifications_repo = NotificationsRepository(
+            session_maker=cfg.db.async_session_maker,
+            # TODO: Make sure that this is the correct name
+            alertmanager_webhook_role=os.environ.get("ALERTMANAGER_WEBHOOK_ROLE", "alertmanager-webhook"),
+        )
+
         kc_api: IKeycloakAPI
         if cfg.dummy_stores:
             dummy_users = [
@@ -149,4 +164,7 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
             session_tasks=session_tasks,
             capacity_reservation_tasks=capacity_reservation_tasks,
             resource_requests_recorder=resource_requests_recorder,
+            k8s_client=k8s_client,
+            notifications_repo=notifications_repo,
+            resource_usage_service=resource_usage_service,
         )

bases/renku_data_services/data_tasks/task_defs.py

@@ -27,7 +27,10 @@
 from renku_data_services.base_models.metrics import MetricsEvent
 from renku_data_services.data_tasks.dependencies import DependencyManager
 from renku_data_services.data_tasks.taskman import TaskDefininions
+from renku_data_services.k8s.models import K8sObject, K8sObjectFilter
 from renku_data_services.namespace.models import NamespaceKind
+from renku_data_services.notebooks.constants import AMALTHEA_SESSION_GVK
+from renku_data_services.notifications.models import UnsavedAlert
 from renku_data_services.solr.solr_client import DefaultSolrClient
 
 logger = logging.getLogger(__name__)
@@ -122,8 +125,8 @@ async def sync_user_namespaces(dm: DependencyManager) -> None:
             await dm.authz.client.WriteRelationships(authz_change.apply)
             num_authz += 1
         except Exception as err:
-            # NOTE: We do not rollback the authz changes here because it is OK if something is in Authz DB
-            # but not in the message queue but not vice-versa.
+            # NOTE: We do not roll back the authz changes here because it is OK if something is in Authz DB
+            # but not in the message queue but not vice versa.
             logger.error(f"Failed to sync user namespace {user_namespace} because {err}")
             await tx.rollback()
         else:
@@ -442,6 +445,84 @@ async def record_resource_requests(dm: DependencyManager) -> None:
         await asyncio.sleep(interval_seconds)
 
 
+def _extract_session_quota_metadata(session: K8sObject) -> tuple[str, str, int] | None:
+    """Extract session name, user id and resource pool id from an AmaltheaSession object."""
+    manifest = session.manifest
+
+    state = str(manifest.get("status", {}).get("state", "")).lower()
+    if state != "running":
+        return None
+
+    labels = manifest.get("metadata", {}).get("labels", {})
+    user_id = labels.get("renku.io/safe-username")
+    if not user_id:
+        return None
+
+    annotations = manifest.get("metadata", {}).get("annotations", {})
+    resource_pool_id_raw = annotations.get("renku.io/resource_pool_id")
+    if resource_pool_id_raw is None:
+        return None
+    resource_pool_id = int(resource_pool_id_raw)
+
+    return session.name, user_id, resource_pool_id
+
+
+async def check_session_quota_and_send_alerts(dm: DependencyManager) -> None:
+    """Check all active sessions and send alerts if the remaining user quota is below threshold."""
+    # TODO: Use a proper admin user here
+    admin_user = InternalServiceAdmin(id=ServiceAdminId.k8s_watcher)
+    session_filter = K8sObjectFilter(gvk=AMALTHEA_SESSION_GVK)
+
+    async for session in dm.k8s_client.list(session_filter):
+        try:
+            metadata = _extract_session_quota_metadata(session)
+            if metadata is None:
+                continue
+
+            session_name, user_id, resource_pool_id = metadata
+            usage = await dm.resource_usage_service.get_running_week(resource_pool_id=resource_pool_id, user_id=user_id)
+            if usage is None:
+                continue
+
+            available_quota = usage.pool_limits.user_limit.value or usage.pool_limits.total_limit.value
+            if available_quota <= 0:
+                continue
+
+            usage_percentage = (usage.user_usage.cost.value / available_quota) * 100
+
+            if usage_percentage < 75.0:
+                # TODO: Do we need to do any cleanup of old alerts when the quota is back to normal!?
+                continue
+
+            # TODO: Does this send an alert every time the check runs and the quota is low!?
+
+            alert = UnsavedAlert(
+                user_id=user_id,
+                # TODO: Use proper event type here
+                event_type="session_quota_low_remaining",
+                session_name=session_name,
+                title="Session quota running low",
+                message=f"You have used {usage_percentage:.1f}% of your quota in resource pool {resource_pool_id}.",
+            )
+            await dm.notifications_repo.create_or_update_alert(user=admin_user, alert=alert)
+
+            logger.info(f"Created quota alert for user {user_id}, session {session_name}")
+        except Exception as e:
+            logger.warning(f"Failed to check quota for pod: {e}")
+            continue
+
+
+async def monitor_session_quota_and_send_alerts(dm: DependencyManager) -> None:
+    """Periodically check session quotas and send alerts when the remaining quota is low."""
+    while True:
+        try:
+            await check_session_quota_and_send_alerts(dm)
+        except (asyncio.CancelledError, KeyboardInterrupt) as e:
+            logger.warning(f"Exiting: {e}")
+        else:
+            await asyncio.sleep(dm.config.session_quota_alert_check_interval_s)
+
+
 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
@@ -467,5 +548,6 @@ def all_tasks(dm: DependencyManager) -> TaskDefininions:
             "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),
+            "monitor_session_quota_and_send_alerts": lambda: monitor_session_quota_and_send_alerts(dm),
         }
     )