feat: perf improvements for amalthea sessions (#411)

Author: olevski

components/renku_data_services/app_config/config.py

@@ -58,7 +58,7 @@
 from renku_data_services.message_queue.interface import IMessageQueue
 from renku_data_services.message_queue.redis_queue import RedisQueue
 from renku_data_services.namespace.db import GroupRepository
-from renku_data_services.notebooks.config import _NotebooksConfig
+from renku_data_services.notebooks.config import NotebooksConfig
 from renku_data_services.platform.db import PlatformRepository
 from renku_data_services.project.db import ProjectMemberRepository, ProjectRepository
 from renku_data_services.repositories.db import GitRepositoriesRepository
@@ -151,7 +151,7 @@ class Config:
     kc_api: IKeycloakAPI
     message_queue: IMessageQueue
     gitlab_url: str | None
-    nb_config: _NotebooksConfig
+    nb_config: NotebooksConfig
 
     secrets_service_public_key: rsa.RSAPublicKey
     """The public key of the secrets service, used to encrypt user secrets that only it can decrypt."""
@@ -544,7 +544,7 @@ def from_env(cls, prefix: str = "") -> "Config":
         sentry = SentryConfig.from_env(prefix)
         trusted_proxies = TrustedProxiesConfig.from_env(prefix)
         message_queue = RedisQueue(redis)
-        nb_config = _NotebooksConfig.from_env(db)
+        nb_config = NotebooksConfig.from_env(db)
 
         return cls(
             version=version,

components/renku_data_services/base_models/core.py

@@ -31,6 +31,11 @@ def is_authenticated(self) -> bool:
         """Indicates whether the user has successfully logged in."""
         return self.id is not None
 
+    @property
+    def is_anonymous(self) -> bool:
+        """Indicates whether the user is anonymous."""
+        return isinstance(self, AnonymousAPIUser)
+
     def get_full_name(self) -> str | None:
         """Generate the closest thing to a full name if the full name field is not set."""
         full_name = self.full_name or " ".join(filter(None, (self.first_name, self.last_name)))

components/renku_data_services/crc/db.py

@@ -292,6 +292,13 @@ async def get_classes(
             orms = res.scalars().all()
             return [orm.dump() for orm in orms]
 
+    async def get_resource_class(self, api_user: base_models.APIUser, id: int) -> models.ResourceClass:
+        """Get a specific resource class by its ID."""
+        classes = await self.get_classes(api_user, id)
+        if len(classes) == 0:
+            raise errors.MissingResourceError(message=f"The resource class with ID {id} cannot be found", quiet=True)
+        return classes[0]
+
     @_only_admins
     async def insert_resource_class(
         self,

components/renku_data_services/migrations/versions/584598f3b769_expand_and_separate_environments_from_.py

@@ -1,7 +1,7 @@
 """expand and separate environments from session launchers
 
 Revision ID: 584598f3b769
-Revises: 9058bf0a1a12
+Revises: 726d5d0e1f28
 Create Date: 2024-08-12 14:25:24.292285
 
 """
@@ -12,7 +12,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "584598f3b769"
-down_revision = "9058bf0a1a12"
+down_revision = "726d5d0e1f28"
 branch_labels = None
 depends_on = None
 

components/renku_data_services/notebooks/api/amalthea_patches/git_proxy.py

@@ -6,48 +6,53 @@
 
 from kubernetes import client
 
+from renku_data_services.base_models.core import AnonymousAPIUser, AuthenticatedAPIUser
 from renku_data_services.notebooks.api.amalthea_patches.utils import get_certificates_volume_mounts
+from renku_data_services.notebooks.api.classes.repository import GitProvider, Repository
+from renku_data_services.notebooks.config import NotebooksConfig
 
 if TYPE_CHECKING:
     # NOTE: If these are directly imported then you get circular imports.
     from renku_data_services.notebooks.api.classes.server import UserServer
 
 
-async def main_container(server: "UserServer") -> client.V1Container | None:
+async def main_container(
+    user: AnonymousAPIUser | AuthenticatedAPIUser,
+    config: NotebooksConfig,
+    repositories: list[Repository],
+    git_providers: list[GitProvider],
+) -> client.V1Container | None:
     """The patch that adds the git proxy container to a session statefulset."""
-    repositories = await server.repositories()
-    if not server.user.is_authenticated or not repositories:
+    if not user.is_authenticated or not repositories or user.access_token is None or user.refresh_token is None:
         return None
 
     etc_cert_volume_mount = get_certificates_volume_mounts(
-        server.config,
+        config,
         custom_certs=False,
         etc_certs=True,
         read_only_etc_certs=True,
     )
 
     prefix = "GIT_PROXY_"
-    git_providers = await server.git_providers()
-    repositories = await server.repositories()
     env = [
-        client.V1EnvVar(name=f"{prefix}PORT", value=str(server.config.sessions.git_proxy.port)),
-        client.V1EnvVar(name=f"{prefix}HEALTH_PORT", value=str(server.config.sessions.git_proxy.health_port)),
+        client.V1EnvVar(name=f"{prefix}PORT", value=str(config.sessions.git_proxy.port)),
+        client.V1EnvVar(name=f"{prefix}HEALTH_PORT", value=str(config.sessions.git_proxy.health_port)),
         client.V1EnvVar(
             name=f"{prefix}ANONYMOUS_SESSION",
-            value="false" if server.user.is_authenticated else "true",
+            value="false" if user.is_authenticated else "true",
         ),
-        client.V1EnvVar(name=f"{prefix}RENKU_ACCESS_TOKEN", value=str(server.user.access_token)),
-        client.V1EnvVar(name=f"{prefix}RENKU_REFRESH_TOKEN", value=str(server.user.refresh_token)),
-        client.V1EnvVar(name=f"{prefix}RENKU_REALM", value=server.config.keycloak_realm),
+        client.V1EnvVar(name=f"{prefix}RENKU_ACCESS_TOKEN", value=str(user.access_token)),
+        client.V1EnvVar(name=f"{prefix}RENKU_REFRESH_TOKEN", value=str(user.refresh_token)),
+        client.V1EnvVar(name=f"{prefix}RENKU_REALM", value=config.keycloak_realm),
         client.V1EnvVar(
             name=f"{prefix}RENKU_CLIENT_ID",
-            value=str(server.config.sessions.git_proxy.renku_client_id),
+            value=str(config.sessions.git_proxy.renku_client_id),
         ),
         client.V1EnvVar(
             name=f"{prefix}RENKU_CLIENT_SECRET",
-            value=str(server.config.sessions.git_proxy.renku_client_secret),
+            value=str(config.sessions.git_proxy.renku_client_secret),
         ),
-        client.V1EnvVar(name=f"{prefix}RENKU_URL", value="https://" + server.config.sessions.ingress.host),
+        client.V1EnvVar(name=f"{prefix}RENKU_URL", value="https://" + config.sessions.ingress.host),
         client.V1EnvVar(
             name=f"{prefix}REPOSITORIES",
             value=json.dumps([asdict(repo) for repo in repositories]),
@@ -60,7 +65,7 @@ async def main_container(server: "UserServer") -> client.V1Container | None:
         ),
     ]
     container = client.V1Container(
-        image=server.config.sessions.git_proxy.image,
+        image=config.sessions.git_proxy.image,
         security_context={
             "fsGroup": 100,
             "runAsGroup": 1000,
@@ -73,14 +78,14 @@ async def main_container(server: "UserServer") -> client.V1Container | None:
         liveness_probe={
             "httpGet": {
                 "path": "/health",
-                "port": server.config.sessions.git_proxy.health_port,
+                "port": config.sessions.git_proxy.health_port,
             },
             "initialDelaySeconds": 3,
         },
         readiness_probe={
             "httpGet": {
                 "path": "/health",
-                "port": server.config.sessions.git_proxy.health_port,
+                "port": config.sessions.git_proxy.health_port,
             },
             "initialDelaySeconds": 3,
         },
@@ -98,7 +103,8 @@ async def main(server: "UserServer") -> list[dict[str, Any]]:
     if not server.user.is_authenticated or not repositories:
         return []
 
-    container = await main_container(server)
+    git_providers = await server.git_providers()
+    container = await main_container(server.user, server.config, repositories, git_providers)
     if not container:
         return []
 

components/renku_data_services/notebooks/api/amalthea_patches/init_containers.py

@@ -3,72 +3,77 @@
 import json
 import os
 from dataclasses import asdict
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from typing import TYPE_CHECKING, Any
 
 from kubernetes import client
 
+from renku_data_services.base_models.core import AnonymousAPIUser, AuthenticatedAPIUser
 from renku_data_services.notebooks.api.amalthea_patches.utils import get_certificates_volume_mounts
-from renku_data_services.notebooks.config import _NotebooksConfig
+from renku_data_services.notebooks.api.classes.repository import GitProvider, Repository
+from renku_data_services.notebooks.config import NotebooksConfig
 
 if TYPE_CHECKING:
     # NOTE: If these are directly imported then you get circular imports.
     from renku_data_services.notebooks.api.classes.server import UserServer
 
 
-async def git_clone_container_v2(server: "UserServer") -> dict[str, Any] | None:
+async def git_clone_container_v2(
+    user: AuthenticatedAPIUser | AnonymousAPIUser,
+    config: NotebooksConfig,
+    repositories: list[Repository],
+    git_providers: list[GitProvider],
+    workspace_mount_path: PurePosixPath,
+    work_dir: PurePosixPath,
+    lfs_auto_fetch: bool = False,
+) -> dict[str, Any] | None:
     """Returns the specification for the container that clones the user's repositories for new operator."""
     amalthea_session_work_volume: str = "amalthea-volume"
-    repositories = await server.repositories()
     if not repositories:
         return None
 
     etc_cert_volume_mount = get_certificates_volume_mounts(
-        server.config,
+        config,
         custom_certs=False,
         etc_certs=True,
         read_only_etc_certs=True,
     )
 
-    user_is_anonymous = not server.user.is_authenticated
     prefix = "GIT_CLONE_"
     env = [
-        {
-            "name": f"{prefix}WORKSPACE_MOUNT_PATH",
-            "value": server.workspace_mount_path.as_posix(),
-        },
+        {"name": f"{prefix}WORKSPACE_MOUNT_PATH", "value": workspace_mount_path.as_posix()},
         {
             "name": f"{prefix}MOUNT_PATH",
-            "value": server.work_dir.as_posix(),
+            "value": work_dir.as_posix(),
         },
         {
             "name": f"{prefix}LFS_AUTO_FETCH",
-            "value": "1" if server.server_options.lfs_auto_fetch else "0",
+            "value": "1" if lfs_auto_fetch else "0",
         },
         {
             "name": f"{prefix}USER__USERNAME",
-            "value": server.user.email,
+            "value": user.email,
         },
         {
             "name": f"{prefix}USER__RENKU_TOKEN",
-            "value": str(server.user.access_token),
+            "value": str(user.access_token),
         },
-        {"name": f"{prefix}IS_GIT_PROXY_ENABLED", "value": "0" if user_is_anonymous else "1"},
+        {"name": f"{prefix}IS_GIT_PROXY_ENABLED", "value": "0" if user.is_anonymous else "1"},
         {
             "name": f"{prefix}SENTRY__ENABLED",
-            "value": str(server.config.sessions.git_clone.sentry.enabled).lower(),
+            "value": str(config.sessions.git_clone.sentry.enabled).lower(),
         },
         {
             "name": f"{prefix}SENTRY__DSN",
-            "value": server.config.sessions.git_clone.sentry.dsn,
+            "value": config.sessions.git_clone.sentry.dsn,
         },
         {
             "name": f"{prefix}SENTRY__ENVIRONMENT",
-            "value": server.config.sessions.git_clone.sentry.env,
+            "value": config.sessions.git_clone.sentry.env,
         },
         {
             "name": f"{prefix}SENTRY__SAMPLE_RATE",
-            "value": str(server.config.sessions.git_clone.sentry.sample_rate),
+            "value": str(config.sessions.git_clone.sentry.sample_rate),
         },
         {"name": "SENTRY_RELEASE", "value": os.environ.get("SENTRY_RELEASE")},
         {
@@ -80,12 +85,12 @@ async def git_clone_container_v2(server: "UserServer") -> dict[str, Any] | None:
             "value": str(Path(etc_cert_volume_mount[0]["mountPath"]) / "ca-certificates.crt"),
         },
     ]
-    if server.user.is_authenticated:
-        if server.user.email:
+    if user.is_authenticated:
+        if user.email:
             env.append(
-                {"name": f"{prefix}USER__EMAIL", "value": server.user.email},
+                {"name": f"{prefix}USER__EMAIL", "value": user.email},
             )
-        full_name = server.user.get_full_name()
+        full_name = user.get_full_name()
         if full_name:
             env.append(
                 {
@@ -105,7 +110,8 @@ async def git_clone_container_v2(server: "UserServer") -> dict[str, Any] | None:
         )
 
     # Set up git providers
-    required_git_providers = await server.required_git_providers()
+    required_provider_ids: set[str] = {r.provider for r in repositories if r.provider}
+    required_git_providers = [p for p in git_providers if p.id in required_provider_ids]
     for idx, provider in enumerate(required_git_providers):
         obj_env = f"{prefix}GIT_PROVIDERS_{idx}_"
         data = dict(id=provider.id, access_token_url=provider.access_token_url)
@@ -117,7 +123,7 @@ async def git_clone_container_v2(server: "UserServer") -> dict[str, Any] | None:
         )
 
     return {
-        "image": server.config.sessions.git_clone.image,
+        "image": config.sessions.git_clone.image,
         "name": "git-clone",
         "resources": {
             "requests": {
@@ -134,7 +140,7 @@ async def git_clone_container_v2(server: "UserServer") -> dict[str, Any] | None:
         },
         "volumeMounts": [
             {
-                "mountPath": server.workspace_mount_path.as_posix(),
+                "mountPath": workspace_mount_path.as_posix(),
                 "name": amalthea_session_work_volume,
             },
             *etc_cert_volume_mount,
@@ -156,7 +162,6 @@ async def git_clone_container(server: "UserServer") -> dict[str, Any] | None:
         read_only_etc_certs=True,
     )
 
-    user_is_anonymous = not server.user.is_authenticated
     prefix = "GIT_CLONE_"
     env = [
         {
@@ -179,7 +184,7 @@ async def git_clone_container(server: "UserServer") -> dict[str, Any] | None:
             "name": f"{prefix}USER__RENKU_TOKEN",
             "value": str(server.user.access_token),
         },
-        {"name": f"{prefix}IS_GIT_PROXY_ENABLED", "value": "0" if user_is_anonymous else "1"},
+        {"name": f"{prefix}IS_GIT_PROXY_ENABLED", "value": "0" if server.user.is_anonymous else "1"},
         {
             "name": f"{prefix}SENTRY__ENABLED",
             "value": str(server.config.sessions.git_clone.sentry.enabled).lower(),
@@ -288,7 +293,7 @@ async def git_clone(server: "UserServer") -> list[dict[str, Any]]:
     ]
 
 
-def certificates_container(config: _NotebooksConfig) -> tuple[client.V1Container, list[client.V1Volume]]:
+def certificates_container(config: NotebooksConfig) -> tuple[client.V1Container, list[client.V1Volume]]:
     """The specification for the container that setups self signed CAs."""
     init_container = client.V1Container(
         name="init-certificates",
@@ -321,7 +326,7 @@ def certificates_container(config: _NotebooksConfig) -> tuple[client.V1Container
     return (init_container, [volume_etc_certs, volume_custom_certs])
 
 
-def certificates(config: _NotebooksConfig) -> list[dict[str, Any]]:
+def certificates(config: NotebooksConfig) -> list[dict[str, Any]]:
     """Add a container that initializes custom certificate authorities for a session."""
     container, vols = certificates_container(config)
     api_client = client.ApiClient()

components/renku_data_services/notebooks/api/amalthea_patches/ssh.py

@@ -2,10 +2,10 @@
 
 from typing import Any
 
-from renku_data_services.notebooks.config import _NotebooksConfig
+from renku_data_services.notebooks.config import NotebooksConfig
 
 
-def main(config: _NotebooksConfig) -> list[dict[str, Any]]:
+def main(config: NotebooksConfig) -> list[dict[str, Any]]:
     """Adds the required configuration to the session statefulset for SSH access."""
     if not config.sessions.ssh.enabled:
         return []

components/renku_data_services/notebooks/api/amalthea_patches/utils.py

@@ -4,11 +4,11 @@
 
 from kubernetes import client
 
-from renku_data_services.notebooks.config import _NotebooksConfig
+from renku_data_services.notebooks.config import NotebooksConfig
 
 
 def get_certificates_volume_mounts(
-    config: _NotebooksConfig,
+    config: NotebooksConfig,
     etc_certs: bool = True,
     custom_certs: bool = True,
     read_only_etc_certs: bool = False,