feat: Allow remote cluster sessions to function on lock down cluster (#1193)

* Introduce K8s Subclasses for K8sResourceQuota & K8sPriorityClass
  * Properly handle quota.id
  * Move QuotaRepo to crc.db, and conversion methods to models.Quota

* K8sClients
  * Always use the namespace from the ClusterConnection
  * Stabilise some tests by ensuring generated values do not overlap

* K8sWatcher: prevent spurious logs in case of failing connection to remote clusters

* Set AmaltheaSession.spec.url correctly

Author: sambuc

bases/renku_data_services/data_api/dependencies.py

@@ -28,7 +28,7 @@
 from renku_data_services.connected_services.oauth_http import DefaultOAuthHttpClientFactory, OAuthHttpClientFactory
 from renku_data_services.crc import models as crc_models
 from renku_data_services.crc.constants import DEFAULT_RUNTIME_PLATFORM
-from renku_data_services.crc.db import ClusterRepository, ResourcePoolRepository, UserRepository
+from renku_data_services.crc.db import ClusterRepository, QuotaRepository, ResourcePoolRepository, UserRepository
 from renku_data_services.data_api.config import Config
 from renku_data_services.data_connectors.db import (
     DataConnectorRepository,
@@ -37,11 +37,11 @@
 from renku_data_services.git.gitlab import DummyGitlabAPI, EmptyGitlabAPI, GitlabAPI
 from renku_data_services.k8s.clients import (
     K8sClusterClientsPool,
+    K8sPriorityClassClient,
     K8sResourceQuotaClient,
-    K8sSchedulingClient,
 )
 from renku_data_services.k8s.config import KubeConfigEnv
-from renku_data_services.k8s.db import K8sDbCache, QuotaRepository
+from renku_data_services.k8s.db import K8sDbCache
 from renku_data_services.message_queue.db import ReprovisioningRepository
 from renku_data_services.metrics.core import StagingMetricsService
 from renku_data_services.metrics.db import MetricsRepository
@@ -236,9 +236,7 @@ def from_env(cls) -> DependencyManager:
                 kinds_to_cache=[AMALTHEA_SESSION_GVK, JUPYTER_SESSION_GVK, BUILD_RUN_GVK, TASK_RUN_GVK],
             ),
         )
-        quota_repo = QuotaRepository(
-            K8sResourceQuotaClient(client), K8sSchedulingClient(client), namespace=config.k8s_namespace
-        )
+        quota_repo = QuotaRepository(K8sResourceQuotaClient(client), K8sPriorityClassClient(client))
 
         if config.dummy_stores:
             authenticator = DummyAuthenticator()

bases/renku_data_services/k8s_cache/dependencies.py

@@ -2,9 +2,9 @@
 
 from dataclasses import dataclass, field
 
-from renku_data_services.crc.db import ClusterRepository, ResourcePoolRepository
-from renku_data_services.k8s.clients import DummyCoreClient, DummySchedulingClient
-from renku_data_services.k8s.db import K8sDbCache, QuotaRepository
+from renku_data_services.crc.db import ClusterRepository, QuotaRepository, ResourcePoolRepository
+from renku_data_services.k8s.clients import DummyPriorityClassClient, DummyResourceQuotaClient
+from renku_data_services.k8s.db import K8sDbCache
 from renku_data_services.k8s_cache.config import Config
 from renku_data_services.metrics.core import StagingMetricsService
 from renku_data_services.metrics.db import MetricsRepository
@@ -63,9 +63,7 @@ def quota_repo(self) -> QuotaRepository:
             # NOTE: We only need the QuotaRepository to instantiate the ResourcePoolRepository which is used to get
             # the resource class and pool information for metrics. We don't need quota information for metrics at all
             # so we use the dummy client for quotas here as we don't actually access k8s, just the db.
-            self._quota_repo = QuotaRepository(
-                DummyCoreClient(), DummySchedulingClient(), namespace=self.config.k8s.renku_namespace
-            )
+            self._quota_repo = QuotaRepository(DummyResourceQuotaClient(), DummyPriorityClassClient())
         return self._quota_repo
 
     @classmethod

components/renku_data_services/crc/core.py

@@ -12,11 +12,19 @@
 
 def validate_quota(body: apispec.QuotaWithOptionalId) -> models.UnsavedQuota:
     """Validate a quota object."""
-    return models.UnsavedQuota(
-        cpu=body.cpu,
-        memory=body.memory,
-        gpu=body.gpu,
-    )
+    if body.id is None:
+        return models.UnsavedQuota(
+            cpu=body.cpu,
+            memory=body.memory,
+            gpu=body.gpu,
+        )
+    else:
+        return models.Quota(
+            cpu=body.cpu,
+            memory=body.memory,
+            gpu=body.gpu,
+            id=body.id,
+        )
 
 
 def validate_quota_put_patch(body: apispec.QuotaWithId | apispec.QuotaPatch) -> models.QuotaPatch:

components/renku_data_services/crc/db.py

@@ -6,11 +6,14 @@
 it all in one place.
 """
 
+from __future__ import annotations
+
 from asyncio import gather
 from collections.abc import AsyncGenerator, Callable, Collection, Coroutine, Sequence
 from dataclasses import asdict, dataclass, field
 from functools import wraps
 from typing import Any, Concatenate, Optional, ParamSpec, TypeVar
+from uuid import uuid4
 
 from sqlalchemy import NullPool, delete, false, select, true
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
@@ -27,8 +30,9 @@
 from renku_data_services.crc.core import validate_resource_class_update, validate_resource_pool_update
 from renku_data_services.crc.models import ClusterPatch, ClusterSettings, SavedClusterSettings, SessionProtocol
 from renku_data_services.crc.orm import ClusterORM
-from renku_data_services.k8s.constants import DEFAULT_K8S_CLUSTER
-from renku_data_services.k8s.db import QuotaRepository
+from renku_data_services.k8s.client_interfaces import PriorityClassClient, ResourceQuotaClient
+from renku_data_services.k8s.constants import DEFAULT_K8S_CLUSTER, ClusterId
+from renku_data_services.k8s.models import DeletePropagationPolicy, K8sPriorityClass
 from renku_data_services.users.db import UserRepo
 
 
@@ -290,7 +294,6 @@ async def insert_resource_pool(
         self, api_user: base_models.APIUser, new_resource_pool: models.UnsavedResourcePool
     ) -> models.ResourcePool:
         """Insert resource pool into database."""
-
         cluster = None
         if new_resource_pool.cluster_id:
             cluster = await self.__cluster_repo.select(cluster_id=new_resource_pool.cluster_id)
@@ -1108,3 +1111,65 @@ async def delete(self, api_user: base_models.APIUser, cluster_id: ULID) -> None:
             cluster = r.one_or_none()
             if cluster is not None:
                 await session.delete(cluster)
+
+
+@dataclass
+class QuotaRepository:
+    """Adapter for CRUD operations on resource quotas and priority classes in k8s."""
+
+    rq_client: ResourceQuotaClient
+    pc_client: PriorityClassClient
+    _label_name: str = field(init=False, default="app")
+    _label_value: str = field(init=False, default="renku")
+
+    async def get_quota(self, name: str | None, cluster_id: ClusterId) -> models.Quota | None:
+        """Get a specific quota by name."""
+        if not name:
+            return None
+        try:
+            res_quota = await self.rq_client.read_resource_quota(name=name, cluster_id=cluster_id)
+        except errors.MissingResourceError:
+            return None
+        return models.Quota.from_k8s_resource_quota(res_quota)
+
+    async def create_quota(self, new_quota: models.UnsavedQuota, cluster_id: ClusterId) -> models.Quota:
+        """Create a resource quota and priority class."""
+        quota = models.Quota(
+            cpu=new_quota.cpu,
+            memory=new_quota.memory,
+            gpu=new_quota.gpu,
+            gpu_kind=new_quota.gpu_kind,
+            id=new_quota.id if isinstance(new_quota, models.Quota) else str(uuid4()),
+        )
+        labels = {self._label_name: self._label_value}
+
+        # Check if we have a priority class with the given name, if not, create one it.
+        pc = await self.pc_client.read_priority_class(K8sPriorityClass.meta(quota.id, cluster_id))
+        if pc is None:
+            await self.pc_client.create_priority_class(
+                K8sPriorityClass.new(
+                    name=quota.id,
+                    cluster=cluster_id,
+                    global_default=False,
+                    value=100,
+                    preemption_policy="Never",
+                    description="Renku resource quota priority class",
+                    labels=labels,
+                ),
+            )
+
+        res = await self.rq_client.create_resource_quota(quota.to_patch(labels), cluster_id)
+        return models.Quota.from_k8s_resource_quota(res)
+
+    async def delete_quota(self, name: str, cluster_id: ClusterId) -> None:
+        """Delete a resource quota and priority class."""
+        await self.pc_client.delete_priority_class(
+            meta=K8sPriorityClass.meta(name, cluster_id), propagation_policy=DeletePropagationPolicy.foreground
+        )
+        await self.rq_client.delete_resource_quota(name=name, cluster_id=cluster_id)
+
+    async def update_quota(self, quota: models.Quota, cluster_id: ClusterId) -> models.Quota:
+        """Update a specific resource quota."""
+        patch = quota.to_patch({self._label_name: self._label_value})
+        patched_quota = await self.rq_client.patch_resource_quota(quota.id, patch, cluster_id)
+        return models.Quota.from_k8s_resource_quota(patched_quota)

components/renku_data_services/crc/models.py

@@ -7,9 +7,12 @@
 from enum import StrEnum
 from typing import Any, Optional, Protocol, Self
 
+from kubernetes.utils import parse_quantity
+
 from renku_data_services import errors
 from renku_data_services.base_models import ResetType
 from renku_data_services.k8s.constants import DEFAULT_K8S_CLUSTER, ClusterId
+from renku_data_services.k8s.models import K8sPatch, K8sResourceQuota
 
 
 class ResourcesProtocol(Protocol):
@@ -152,15 +155,65 @@ def is_resource_class_compatible(self, rc: ResourceClass | UnsavedResourceClass)
 
 
 @dataclass(frozen=True, eq=True, kw_only=True)
-class Quota(ResourcesCompareMixin):
+class Quota(UnsavedQuota):
     """Quota model."""
 
-    cpu: float
-    memory: int
-    gpu: int
-    gpu_kind: GpuKind = GpuKind.NVIDIA
     id: str
 
+    @classmethod
+    def from_k8s_resource_quota(cls, quota: K8sResourceQuota) -> Quota:
+        """Convert a K8s Resource Quota."""
+
+        def require_key(manifest: dict, key: str, prefix: str = "") -> Any:
+            try:
+                return manifest.get(key)
+            except Exception:
+                raise errors.ValidationError(
+                    message=f"Kubernetes resource quota with missing {prefix}{key} is not supported"
+                ) from None
+
+        spec = require_key(quota.manifest.to_dict(), "spec", "")
+        hard = require_key(spec, "hard", "spec.")
+
+        gpu = 0
+        gpu_kind = GpuKind.NVIDIA
+        for igpu_kind in GpuKind:
+            key = f"requests.{igpu_kind}/gpu"
+            if key in hard:
+                gpu = int(parse_quantity(hard.get(key)))
+                gpu_kind = igpu_kind
+                break
+
+        memory_raw = require_key(hard, "requests.memory", "hard.")
+        cpu_raw = require_key(hard, "requests.cpu", "hard.")
+        return cls(
+            cpu=float(parse_quantity(cpu_raw)) / 1_000_000_000,
+            memory=round(parse_quantity(memory_raw) / 1_000_000_000),
+            gpu=gpu,
+            gpu_kind=gpu_kind,
+            id=quota.name,
+        )
+
+    def to_patch(self, labels: dict[str, str]) -> K8sPatch:
+        """Convert to a manifest."""
+
+        return {
+            "metadata": {
+                "name": self.id,
+                "labels": labels,
+            },
+            "spec": {
+                "hard": {
+                    "requests.cpu": str(round(self.cpu * 1_000_000_000)),
+                    "requests.memory": str(self.memory * 1_000_000_000),
+                    f"requests.{self.gpu_kind}/gpu": self.gpu,
+                },
+                "scopeSelector": {
+                    "matchExpressions": [{"operator": "In", "scopeName": "PriorityClass", "values": [self.id]}]
+                },
+            },
+        }
+
     def is_resource_class_compatible(self, rc: ResourceClass | UnsavedResourceClass) -> bool:
         """Determine if a resource class is compatible with the quota."""
         return rc <= self

components/renku_data_services/k8s/client_interfaces.py

@@ -3,17 +3,18 @@
 from __future__ import annotations
 
 from collections.abc import AsyncIterable
-from typing import Any, Protocol, overload
-
-from kubernetes.client import V1PriorityClass, V1ResourceQuota
+from typing import Protocol
 
 from renku_data_services.k8s.constants import ClusterId
 from renku_data_services.k8s.models import (
-    ClusterScopedK8sObject,
     DeletePropagationPolicy,
     K8sObject,
     K8sObjectFilter,
     K8sObjectMeta,
+    K8sPatch,
+    K8sPatches,
+    K8sPriorityClass,
+    K8sResourceQuota,
     K8sSecret,
 )
 
@@ -22,28 +23,24 @@ class ResourceQuotaClient(Protocol):
     """Methods to manipulate ResourceQuota kubernetes resources."""
 
     def list_resource_quota(
-        self, namespace: str, label_selector: dict[str, str], cluster_id: ClusterId
-    ) -> AsyncIterable[V1ResourceQuota]:
+        self, label_selector: dict[str, str], cluster_id: ClusterId
+    ) -> AsyncIterable[K8sResourceQuota]:
         """List resource quotas."""
         ...
 
-    async def read_resource_quota(self, name: str, namespace: str, cluster_id: ClusterId) -> V1ResourceQuota:
+    async def read_resource_quota(self, name: str, cluster_id: ClusterId) -> K8sResourceQuota:
         """Get a resource quota."""
         ...
 
-    async def create_resource_quota(
-        self, namespace: str, body: V1ResourceQuota, cluster_id: ClusterId
-    ) -> V1ResourceQuota:
+    async def create_resource_quota(self, quota: K8sPatch, cluster_id: ClusterId) -> K8sResourceQuota:
         """Create a resource quota."""
         ...
 
-    async def delete_resource_quota(self, name: str, namespace: str, cluster_id: ClusterId) -> None:
+    async def delete_resource_quota(self, name: str, cluster_id: ClusterId) -> None:
         """Delete a resource quota."""
         ...
 
-    async def patch_resource_quota(
-        self, name: str, namespace: str, body: V1ResourceQuota, cluster_id: ClusterId
-    ) -> V1ResourceQuota:
+    async def patch_resource_quota(self, name: str, patch: K8sPatches, cluster_id: ClusterId) -> K8sResourceQuota:
         """Update a resource quota."""
         ...
 
@@ -59,7 +56,7 @@ async def create_secret(self, secret: K8sSecret) -> K8sSecret:
         """Create a secret."""
         ...
 
-    async def patch_secret(self, secret: K8sObjectMeta, patch: dict[str, Any] | list[dict[str, Any]]) -> K8sSecret:
+    async def patch_secret(self, secret: K8sObjectMeta, patch: K8sPatches) -> K8sSecret:
         """Patch an existing secret."""
         ...
 
@@ -71,18 +68,17 @@ async def delete_secret(self, secret: K8sObjectMeta) -> None:
 class PriorityClassClient(Protocol):
     """Methods to manipulate kubernetes Priority Class resources."""
 
-    async def create_priority_class(self, body: V1PriorityClass, cluster_id: ClusterId) -> V1PriorityClass:
+    async def create_priority_class(self, priority_class: K8sPriorityClass) -> K8sPriorityClass:
         """Create a priority class."""
         ...
 
-    async def read_priority_class(self, name: str, cluster_id: ClusterId) -> V1PriorityClass | None:
+    async def read_priority_class(self, meta: K8sObjectMeta) -> K8sPriorityClass | None:
         """Retrieve a priority class."""
         ...
 
     async def delete_priority_class(
         self,
-        name: str,
-        cluster_id: ClusterId,
+        meta: K8sObjectMeta,
         propagation_policy: DeletePropagationPolicy = DeletePropagationPolicy.foreground,
     ) -> None:
         """Delete a priority class."""
@@ -92,17 +88,11 @@ async def delete_priority_class(
 class K8sClient(Protocol):
     """Methods to manipulate resources on a Kubernetes cluster."""
 
-    @overload
-    async def create(self, obj: K8sObject, refresh: bool) -> K8sObject: ...
-    @overload
-    async def create(self, obj: ClusterScopedK8sObject, refresh: bool) -> ClusterScopedK8sObject: ...
-    async def create(
-        self, obj: K8sObject | ClusterScopedK8sObject, refresh: bool
-    ) -> K8sObject | ClusterScopedK8sObject:
+    async def create(self, obj: K8sObject, refresh: bool) -> K8sObject:
         """Create the k8s object."""
         ...
 
-    async def patch(self, meta: K8sObjectMeta, patch: dict[str, Any] | list[dict[str, Any]]) -> K8sObject:
+    async def patch(self, meta: K8sObjectMeta, patch: K8sPatches) -> K8sObject:
         """Patch a k8s object.
 
         If the patch is a list we assume that we have a rfc6902 json patch like