feat: move shipwright clients over to use multi cluster cached client (#822)

* feat: move shipwright clients over to use multi cluster cached client

* address comments

* fix tests

Author: Panaetius

bases/renku_data_services/k8s_cache/main.py

@@ -8,7 +8,13 @@
 from renku_data_services.k8s.models import Cluster, ClusterId
 from renku_data_services.k8s_cache.config import Config
 from renku_data_services.k8s_watcher import K8sWatcher, k8s_object_handler
-from renku_data_services.notebooks.constants import AMALTHEA_SESSION_KIND, JUPYTER_SESSION_KIND
+from renku_data_services.notebooks.constants import (
+    AMALTHEA_SESSION_KIND,
+    AMALTHEA_SESSION_VERSION,
+    JUPYTER_SESSION_KIND,
+    JUPYTER_SESSION_VERSION,
+)
+from renku_data_services.session.constants import BUILD_RUN_KIND, BUILD_RUN_VERSION, TASK_RUN_KIND, TASK_RUN_VERSION
 
 
 async def main() -> None:
@@ -23,7 +29,12 @@ async def main() -> None:
     watcher = K8sWatcher(
         handler=k8s_object_handler(config.k8s_cache, config.metrics, rp_repo=config.rp_repo),
         clusters={c.id: c for c in clusters},
-        kinds=[AMALTHEA_SESSION_KIND, JUPYTER_SESSION_KIND],
+        kinds=[
+            f"{AMALTHEA_SESSION_KIND}.{AMALTHEA_SESSION_VERSION}",
+            f"{JUPYTER_SESSION_KIND}.{JUPYTER_SESSION_VERSION}",
+            f"{BUILD_RUN_KIND}.{BUILD_RUN_VERSION}",
+            f"{TASK_RUN_KIND}.{TASK_RUN_VERSION}",
+        ],
     )
     await watcher.start()
     logging.info("started watching resources")

components/renku_data_services/app_config/config.py

@@ -55,16 +55,25 @@
 )
 from renku_data_services.db_config import DBConfig
 from renku_data_services.git.gitlab import DummyGitlabAPI, GitlabAPI
-from renku_data_services.k8s.clients import DummyCoreClient, DummySchedulingClient, K8sCoreClient, K8sSchedulingClient
+from renku_data_services.k8s.clients import (
+    DummyCoreClient,
+    DummySchedulingClient,
+    K8sClusterClientsPool,
+    K8sCoreClient,
+    K8sSchedulingClient,
+)
+from renku_data_services.k8s.config import KubeConfigEnv
 from renku_data_services.k8s.quota import QuotaRepository
+from renku_data_services.k8s_watcher.db import K8sDbCache
 from renku_data_services.message_queue.config import RedisConfig
 from renku_data_services.message_queue.db import EventRepository, ReprovisioningRepository
 from renku_data_services.message_queue.interface import IMessageQueue
 from renku_data_services.message_queue.redis_queue import RedisQueue
 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 NotebooksConfig
+from renku_data_services.notebooks.config import NotebooksConfig, get_clusters
+from renku_data_services.notebooks.constants import AMALTHEA_SESSION_KIND, JUPYTER_SESSION_KIND
 from renku_data_services.platform.db import PlatformRepository
 from renku_data_services.project.db import (
     ProjectMemberRepository,
@@ -77,6 +86,7 @@
 from renku_data_services.search.reprovision import SearchReprovision
 from renku_data_services.secrets.db import LowLevelUserSecretsRepo, UserSecretsRepo
 from renku_data_services.session import crs as session_crs
+from renku_data_services.session.constants import BUILD_RUN_KIND, TASK_RUN_KIND
 from renku_data_services.session.db import SessionRepository
 from renku_data_services.session.k8s_client import ShipwrightClient
 from renku_data_services.solr.solr_client import SolrClientConfig
@@ -182,7 +192,7 @@ class BuildsConfig:
     tolerations: list[session_crs.Toleration] | None = None
 
     @classmethod
-    def from_env(cls, prefix: str = "", namespace: str = "") -> "BuildsConfig":
+    def from_env(cls, db: DBConfig, prefix: str = "", namespace: str = "") -> "BuildsConfig":
         """Create a config from environment variables."""
         enabled = os.environ.get(f"{prefix}IMAGE_BUILDERS_ENABLED", "false").lower() == "true"
         build_output_image_prefix = os.environ.get(f"{prefix}BUILD_OUTPUT_IMAGE_PREFIX")
@@ -216,11 +226,18 @@ def from_env(cls, prefix: str = "", namespace: str = "") -> "BuildsConfig":
         elif not enabled:
             shipwright_client = None
         else:
-            # TODO: is there a reason to use a different cache URL here?
-            cache_url = os.environ["NB_AMALTHEA_V2__CACHE_URL"]
+            # NOTE: we need to get an async client as a sync client can't be used in an async way
+            # But all the config code is not async, so we need to drop into the running loop, if there is one
+            kr8s_api = KubeConfigEnv().api()
+            k8s_db_cache = K8sDbCache(db.async_session_maker)
+            client = K8sClusterClientsPool(
+                clusters=get_clusters("/secrets/kube_configs", namespace=namespace, api=kr8s_api),
+                cache=k8s_db_cache,
+                kinds_to_cache=[AMALTHEA_SESSION_KIND, JUPYTER_SESSION_KIND, BUILD_RUN_KIND, TASK_RUN_KIND],
+            )
             shipwright_client = ShipwrightClient(
+                client=client,
                 namespace=namespace,
-                cache_url=cache_url,
             )
 
         node_selector: dict[str, str] | None = None
@@ -730,7 +747,7 @@ def from_env(cls, prefix: str = "") -> "Config":
         trusted_proxies = TrustedProxiesConfig.from_env(prefix)
         message_queue = RedisQueue(redis)
         nb_config = NotebooksConfig.from_env(db)
-        builds_config = BuildsConfig.from_env(prefix, k8s_namespace)
+        builds_config = BuildsConfig.from_env(db, prefix, k8s_namespace)
         posthog = PosthogConfig.from_env(prefix)
 
         return cls(

components/renku_data_services/k8s/clients.py

@@ -245,7 +245,7 @@ async def __list(self, _filter: K8sObjectFilter) -> AsyncIterable[APIObjectInClu
 
         try:
             res = await self.__cluster.api.async_get(
-                _filter.kind,
+                _filter.fully_qualified_kind,
                 *names,
                 label_selector=_filter.label_selector,
                 namespace=_filter.namespace,

components/renku_data_services/k8s/config.py

@@ -0,0 +1,111 @@
+"""Base config for k8s."""
+
+import glob
+import os
+
+import kr8s
+import yaml
+from sanic.log import logger
+
+from renku_data_services.k8s.constants import DEFAULT_K8S_CLUSTER
+from renku_data_services.k8s.models import Cluster, ClusterId
+
+
+class KubeConfig:
+    """Wrapper around kube config to get a kr8s api."""
+
+    def __init__(
+        self,
+        kubeconfig: str | None = None,
+        current_context_name: str | None = None,
+        ns: str | None = None,
+        sa: str | None = None,
+        url: str | None = None,
+    ) -> None:
+        self._kubeconfig = kubeconfig
+        self._ns = ns
+        self._current_context_name = current_context_name
+        self._sa = sa
+        self._url = url
+
+    def _sync_api(self) -> kr8s.Api | kr8s._AsyncApi:
+        return kr8s.api(
+            kubeconfig=self._kubeconfig,
+            namespace=self._ns,
+            context=self._current_context_name,
+        )
+
+    def _async_api(self) -> kr8s.asyncio.Api:
+        """Create an async api client from sync code.
+
+        Kr8s cannot return an AsyncAPI instance from sync code, and we can't easily make all our config code async,
+        so this method is a direct copy of the kr8s sync client code, just that it returns an async client.
+        """
+        ret = kr8s._async_utils.run_sync(kr8s.asyncio.api)(
+            url=self._url,
+            kubeconfig=self._kubeconfig,
+            serviceaccount=self._sa,
+            namespace=self._ns,
+            context=self._current_context_name,
+            _asyncio=True,  # This is the only line that is different from kr8s code
+        )
+        assert isinstance(ret, kr8s.asyncio.Api)
+        return ret
+
+    def api(self, _async: bool = True) -> kr8s.Api | kr8s._AsyncApi:
+        """Instantiate the Kr8s Api object based on the configuration."""
+        if _async:
+            return self._async_api()
+        else:
+            return self._sync_api()
+
+
+class KubeConfigEnv(KubeConfig):
+    """Get a kube config from the environment."""
+
+    def __init__(self) -> None:
+        super().__init__(ns=os.environ.get("K8S_NAMESPACE", "default"))
+
+
+class KubeConfigYaml(KubeConfig):
+    """Get a kube config from a yaml file."""
+
+    def __init__(self, kubeconfig: str) -> None:
+        super().__init__(kubeconfig=kubeconfig)
+
+        with open(kubeconfig) as stream:
+            _conf = yaml.safe_load(stream)
+
+        self._current_context_name = _conf.get("current-context", None)
+        if self._current_context_name is not None:
+            for context in _conf.get("contexts", []):
+                name = context.get("name", None)
+                inner = context.get("context", None)
+                if inner is not None and name is not None and name == self._current_context_name:
+                    self._ns = inner.get("namespace", None)
+                    break
+
+
+def get_clusters(kube_conf_root_dir: str, namespace: str, api: kr8s.asyncio.Api) -> list[Cluster]:
+    """Get all clusters accessible to the application."""
+    clusters = [Cluster(id=DEFAULT_K8S_CLUSTER, namespace=namespace, api=api)]
+
+    if os.path.exists(kube_conf_root_dir):
+        for filename in glob.glob(pathname="*.yaml", root_dir=kube_conf_root_dir):
+            try:
+                kube_config = KubeConfigYaml(filename)
+                cluster = Cluster(
+                    id=ClusterId(filename.removesuffix(".yaml")),
+                    namespace=kube_config.api().namespace,
+                    api=kube_config.api(),
+                )
+                clusters.append(cluster)
+                logger.info(f"Successfully loaded Kubernetes config: '{kube_conf_root_dir}/{filename}'")
+            except Exception as e:
+                logger.warning(
+                    f"Failed while loading '{kube_conf_root_dir}/{filename}', ignoring kube config. Error: {e}"
+                )
+    else:
+        logger.warning(f"Cannot open directory '{kube_conf_root_dir}', ignoring kube configs...")
+
+    return clusters

components/renku_data_services/k8s/constants.py

@@ -0,0 +1,7 @@
+"""Constant values for k8s."""
+
+from typing import Final
+
+from renku_data_services.k8s.models import ClusterId
+
+DEFAULT_K8S_CLUSTER: Final[ClusterId] = ClusterId("renkulab")

components/renku_data_services/k8s/models.py

@@ -140,9 +140,27 @@ class K8sObjectFilter:
     namespace: str | None = None
     cluster: ClusterId | None = None
     version: str | None = None
+    """Version should be of the form 'group/version', e.g. 'shipwright.io/v1' or 'core/v1'"""
+
     label_selector: dict[str, str] | None = None
     user_id: str | None = None
 
+    @property
+    def fully_qualified_kind(self) -> str:
+        """Returns the fully qualified kind string for this filter.
+
+        Note: This exists because kr8s has some methods where it only allows you to specify 'kind' and then has
+        weird logic to split that. This method is essentially the reverse of the kr8s logic so we can hand it a
+        string it will accept.
+        """
+        if not self.version:
+            return self.kind
+        if "/" in self.version:
+            # e.g. buildrun.shipwright.io/v1beta1
+            return f"{self.kind}.{self.version}"
+        # e.g. pod/v1
+        return f"{self.kind}/{self.version}"
+
 
 @dataclass(eq=True, frozen=True)
 class Cluster:

components/renku_data_services/k8s_watcher/core.py

@@ -22,6 +22,7 @@
 from renku_data_services.k8s.models import Cluster, ClusterId, K8sObject, K8sObjectMeta
 from renku_data_services.k8s_watcher.db import K8sDbCache
 from renku_data_services.notebooks.crs import State
+from renku_data_services.session.constants import DUMMY_TASK_RUN_USER_ID
 
 
 @dataclass
@@ -39,6 +40,11 @@ def user_id(self) -> str | None:
                 return cast(str, self.obj.metadata.labels["renku.io/userId"])
             case "amaltheasession":
                 return cast(str, self.obj.metadata.labels["renku.io/safe-username"])
+            case "buildrun":
+                return cast(str, self.obj.metadata.labels["renku.io/safe-username"])
+
+            case "taskrun":
+                return DUMMY_TASK_RUN_USER_ID
             case _:
                 return None
 

components/renku_data_services/notebooks/api/classes/k8s_client.py

@@ -16,6 +16,7 @@
 from renku_data_services.crc.db import ResourcePoolRepository
 from renku_data_services.errors import errors
 from renku_data_services.k8s.clients import K8sClusterClientsPool
+from renku_data_services.k8s.constants import DEFAULT_K8S_CLUSTER
 from renku_data_services.k8s.models import Cluster, ClusterId, K8sObject, K8sObjectFilter, K8sObjectMeta
 from renku_data_services.k8s_watcher.core import APIObjectInCluster
 from renku_data_services.notebooks.api.classes.auth import GitlabToken, RenkuTokens
@@ -25,7 +26,6 @@
 from renku_data_services.notebooks.util.kubernetes_ import find_env_var
 from renku_data_services.notebooks.util.retries import retry_with_exponential_backoff_async
 
-DEFAULT_K8S_CLUSTER: ClusterId = ClusterId("renkulab")
 sanitizer = kubernetes.client.ApiClient().sanitize_for_serialization