Skip to content

Commit

Permalink
feat: perf improvements for amalthea sessions (#411)
Browse files Browse the repository at this point in the history
  • Loading branch information
olevski committed Oct 4, 2024
1 parent ae87e85 commit fd2e56c
Show file tree
Hide file tree
Showing 19 changed files with 251 additions and 117 deletions.
6 changes: 3 additions & 3 deletions components/renku_data_services/app_config/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,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
Expand Down Expand Up @@ -145,7 +145,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."""
Expand Down Expand Up @@ -511,7 +511,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,
Expand Down
5 changes: 5 additions & 0 deletions components/renku_data_services/base_models/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)))
Expand Down
7 changes: 7 additions & 0 deletions components/renku_data_services/crc/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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]),
Expand All @@ -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,
Expand All @@ -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,
},
Expand All @@ -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 []

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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")},
{
Expand All @@ -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(
{
Expand All @@ -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)
Expand All @@ -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": {
Expand All @@ -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,
Expand All @@ -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 = [
{
Expand All @@ -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(),
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 []
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from renku_data_services.notebooks.api.classes.repository import GitProvider, Repository
from renku_data_services.notebooks.api.schemas.secrets import K8sUserSecrets
from renku_data_services.notebooks.api.schemas.server_options import ServerOptions
from renku_data_services.notebooks.config import _NotebooksConfig
from renku_data_services.notebooks.config import NotebooksConfig
from renku_data_services.notebooks.crs import JupyterServerV1Alpha1
from renku_data_services.notebooks.errors.programming import DuplicateEnvironmentVariableError
from renku_data_services.notebooks.errors.user import MissingResourceError
Expand All @@ -46,7 +46,7 @@ def __init__(
k8s_client: K8sClient,
workspace_mount_path: PurePosixPath,
work_dir: PurePosixPath,
config: _NotebooksConfig,
config: NotebooksConfig,
internal_gitlab_user: APIUser,
using_default_image: bool = False,
is_image_private: bool = False,
Expand Down Expand Up @@ -380,7 +380,7 @@ def __init__(
k8s_client: K8sClient,
workspace_mount_path: PurePosixPath,
work_dir: PurePosixPath,
config: _NotebooksConfig,
config: NotebooksConfig,
gitlab_project: Project | None,
internal_gitlab_user: APIUser,
using_default_image: bool = False,
Expand Down Expand Up @@ -506,7 +506,7 @@ def __init__(
workspace_mount_path: PurePosixPath,
work_dir: PurePosixPath,
repositories: list[Repository],
config: _NotebooksConfig,
config: NotebooksConfig,
internal_gitlab_user: APIUser,
using_default_image: bool = False,
is_image_private: bool = False,
Expand Down
Loading

0 comments on commit fd2e56c

Please sign in to comment.