feat: task to check session quota and send alerts (#1252)

Author: mohammad-alisafaee

bases/renku_data_services/data_tasks/config.py

@@ -52,6 +52,10 @@ class Config:
     short_task_period_s: int
     long_task_period_s: int
     enable_resource_request_tracking: bool
+    session_quota_alert_check_interval_s: int
+    session_quota_alert_remaining_threshold_p: int
+    session_quota_alert_critical_m: int
+    alertmanager_webhook_role: str
 
     @classmethod
     def from_env(cls) -> Config:
@@ -69,6 +73,9 @@ 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))
+        session_quota_alert_remaining_threshold = int(os.environ.get("SESSION_QUOTA_ALERT_REMAINING_THRESHOLD_P", 20))
+        session_quota_alert_critical = int(os.environ.get("SESSION_QUOTA_ALERT_CRITICAL_M", 10))
 
         k8s_config_root = os.environ.get("K8S_CONFIG_ROOT", "/secrets/kube_configs")
 
@@ -92,4 +99,8 @@ 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,
+            session_quota_alert_remaining_threshold_p=session_quota_alert_remaining_threshold,
+            session_quota_alert_critical_m=session_quota_alert_critical,
+            alertmanager_webhook_role=os.environ.get("ALERTMANAGER_WEBHOOK_ROLE", "alertmanager-webhook"),
         )

bases/renku_data_services/data_tasks/dependencies.py

@@ -16,12 +16,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 +52,10 @@ class DependencyManager:
     session_tasks: SessionTasks
     capacity_reservation_tasks: CapacityReservationTasks
     resource_requests_recorder: ResourcesRequestRecorder
+    k8s_client: K8sClusterClientsPool
+    notifications_repo: NotificationsRepository
+    resource_usage_service: ResourceUsageService
+    resource_requests_repo: ResourceRequestsRepo
 
     @classmethod
     def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
@@ -112,15 +118,23 @@ 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,
+            alertmanager_webhook_role=cfg.alertmanager_webhook_role,
+        )
+
         kc_api: IKeycloakAPI
         if cfg.dummy_stores:
             dummy_users = [
@@ -149,4 +163,8 @@ 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,
+            resource_requests_repo=resource_requests_repo,
         )

bases/renku_data_services/data_tasks/task_defs.py

@@ -27,8 +27,12 @@
 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
+from renku_data_services.utils.core import get_nonzero_minimum
 
 logger = logging.getLogger(__name__)
 
@@ -122,8 +126,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 +446,121 @@ async def record_resource_requests(dm: DependencyManager) -> None:
         await asyncio.sleep(interval_seconds)
 
 
+def _extract_session_quota_metadata(session: K8sObject) -> tuple[str, str, int, int | None] | None:
+    """Extract session name, user id, resource pool id and resource class 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)
+
+    resource_class_id_raw = annotations.get("renku.io/resource_class_id")
+    resource_class_id = int(resource_class_id_raw) if resource_class_id_raw else None
+
+    return session.name, user_id, resource_pool_id, resource_class_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."""
+    admin_user = InternalServiceAdmin(id=ServiceAdminId.capacity_reservation)
+    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 not metadata:
+                continue
+
+            session_name, user_id, resource_pool_id, resource_class_id = metadata
+
+            usage = await dm.resource_usage_service.get_running_week(resource_pool_id=resource_pool_id, user_id=user_id)
+            if not usage:
+                continue
+
+            total_quota = get_nonzero_minimum(usage.pool_limits.user_limit.value, usage.pool_limits.total_limit.value)
+            if total_quota <= 0:
+                continue
+
+            usage_p = (usage.user_usage.cost.value / total_quota) * 100
+            usage_threshold_p = 100 - dm.config.session_quota_alert_remaining_threshold_p
+            if usage_p <= usage_threshold_p:
+                continue
+
+            # NOTE: Without the resource_class_id, we cannot calculate the remaining time
+            if resource_class_id:
+                if usage.user_usage.cost.value >= total_quota:
+                    message = f"Your session {session_name} in resource pool {resource_pool_id} has exhausted its quota"
+                    log_message = (
+                        f"Session {session_name} for user {user_id} has exhausted its quota in resource pool "
+                        f"{resource_pool_id}"
+                    )
+                    hibernation_alert = UnsavedAlert(
+                        user_id=user_id,
+                        event_type="session_quota_exhausted",
+                        session_name=session_name,
+                        title="Session paused due to quota exhaustion",
+                        message=message,
+                    )
+                    await dm.notifications_repo.create_or_update_alert(user=admin_user, alert=hibernation_alert)
+                    logger.info(log_message)
+                    continue
+
+                resource_class_cost = await dm.resource_requests_repo.find_resource_class_costs(
+                    resource_pool_id, resource_class_id
+                )
+
+                if resource_class_cost and resource_class_cost.cost.value > 0:
+                    remaining_credits = total_quota - usage.user_usage.cost.value
+                    remaining_minutes = (remaining_credits / resource_class_cost.cost.value) * 60
+                    if remaining_minutes < dm.config.session_quota_alert_critical_m:
+                        critical_alert = UnsavedAlert(
+                            user_id=user_id,
+                            event_type="session_quota_critically_low",
+                            session_name=session_name,
+                            title="Session quota expiring soon",
+                            message=f"Your session in resource pool {resource_pool_id} will run out of quota in "
+                            + f"approximately {remaining_minutes:.0f} minutes.",
+                        )
+                        await dm.notifications_repo.create_or_update_alert(user=admin_user, alert=critical_alert)
+                        logger.info(f"Created critical quota alert for user {user_id}, session {session_name}")
+                        continue
+
+            alert = UnsavedAlert(
+                user_id=user_id,
+                event_type="session_quota_low",
+                session_name=session_name,
+                title="Session quota running low",
+                message=f"You have used {usage_p:.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 +586,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),
         }
     )

components/renku_data_services/crc/db.py

@@ -264,14 +264,10 @@ async def filter_resource_pools(
             rp: schemas.ResourcePoolORM
             for rp in res.scalars().all():
                 quota = await self.quotas_repo.get_quota(rp.quota, rp.get_cluster_id())
-
-                # TODO: Do we need to calculate usage/remaining for admins!?
                 credits_used = None
                 if self.resource_usage_service and api_user.is_authenticated:
                     usage_summary = await self.resource_usage_service.usage_of_running_week(rp.id, api_user.id)
                     if not usage_summary.is_empty():
-                        # TODO: Should we use the cost_raw and runtime_hours here instead?
-                        # TODO: Should we calculate usage for individual entries?
                         credits_used = usage_summary.cost
 
                 rp_model = rp.dump(quota, criteria, credits_used)

components/renku_data_services/notifications/db.py

@@ -103,7 +103,7 @@ async def update_alert(self, user: base_models.APIUser, alert_id: ULID, patch: m
             return alert_orm.dump()
 
     def __update_alert(self, alert: schemas.AlertORM, update: models.AlertPatch) -> None:
-        if update.resolved is True:
+        if update.resolved:
             alert.resolved_date = datetime.now(UTC)
 
     async def get_alerts_by_properties(

components/renku_data_services/resource_usage/core.py

@@ -175,7 +175,7 @@ async def usage_of_running_week(
         """Return the resource usage for the given pool of the currently running week.
 
         The week start is Monday 0:00 UTC. Resource usage is returned in 'credits'. When a user_id
-        is given, the results represent the usage of only that user. Otherwise the overall pool usage
+        is given, the results represent the usage of only that user. Otherwise, the overall pool usage
         is returned. The running week is calculated from the `current_time` argument, which is the current
         time if not specified.
         """

components/renku_data_services/utils/core.py

@@ -236,3 +236,8 @@ async def get_openbis_pat(
                 "from creating a personal access token."
             )
         return json2["result"][0]["permId"], valid_to
+
+
+def get_nonzero_minimum(a: int, b: int) -> int:
+    """Return the minimum nonzero value, where both inputs are non-negative."""
+    return a if b == 0 else b if a == 0 else min(a, b)

test/bases/renku_data_services/data_api/test_repositories.py

@@ -111,7 +111,7 @@ async def test_get_repository_without_connection(
 
 @pytest.mark.asyncio
 async def test_get_one_repository(oauth2_test_client: SanicASGITestClient, user_headers, create_oauth2_connection):
-    connection = await create_oauth2_connection("provider_1")
+    connection = await create_oauth2_connection("provider_2")
     repository_url = "https://example.org/username/my_repo.git"
 
     _, res = await oauth2_test_client.get(
@@ -122,7 +122,7 @@ async def test_get_one_repository(oauth2_test_client: SanicASGITestClient, user_
     assert res.json is not None
     result = res.json
     assert result.get("connection", {}).get("id") == connection["id"]
-    assert result.get("provider", {}).get("id") == "provider_1"
+    assert result.get("provider", {}).get("id") == "provider_2"
     assert result.get("metadata") is not None
     repository_metadata = result["metadata"]
     assert repository_metadata.get("git_url") == repository_url