14382014257

Committed 10 Apr 2025 01:42PM UTC coverage: 86.576% (+0.2%) from 86.351%

Build # 14382014257

Build Type

Pull #759

github

Committed by

web-flow

Commit Message

Merge 470ff1568 into 74eb7d965

Pull Request Pull Request #759: feat: add new service cache and migrations

Run Details

412 of 486 new or added lines in 15 files covered. (84.77%)

18 existing lines in 6 files now uncovered.

20232 of 23369 relevant lines covered (86.58%)

1.53 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.71

/components/renku_data_services/notebooks/config/__init__.py

"""Base motebooks svc configuration."""

import os
from dataclasses import dataclass, field
from typing import Any, Optional, Protocol, Self

import kr8s

from renku_data_services.base_models import APIUser
from renku_data_services.crc.db import ResourcePoolRepository
from renku_data_services.crc.models import ResourceClass
from renku_data_services.db_config.config import DBConfig
from renku_data_services.k8s.clients import K8sCoreClient, K8sSchedulingClient
from renku_data_services.k8s.quota import QuotaRepository
from renku_data_services.k8s_watcher.db import CachedK8sClient, Cluster, K8sDbCache
from renku_data_services.k8s_watcher.models import ClusterId
from renku_data_services.notebooks.api.classes.data_service import (
    CRCValidator,
    DummyCRCValidator,
    DummyGitProviderHelper,
    GitProviderHelper,
)
from renku_data_services.notebooks.api.classes.k8s_client import (
    K8sClient,
)
from renku_data_services.notebooks.api.classes.repository import GitProvider
from renku_data_services.notebooks.api.schemas.server_options import ServerOptions
from renku_data_services.notebooks.config.dynamic import (
    _CloudStorage,
    _GitConfig,
    _K8sConfig,
    _parse_str_as_bool,
    _SentryConfig,
    _ServerOptionsConfig,
    _SessionConfig,
    _UserSecrets,
)
from renku_data_services.notebooks.config.static import _ServersGetEndpointAnnotations
from renku_data_services.notebooks.constants import (
    AMALTHEA_SESSION_KIND,
    AMALTHEA_SESSION_VERSION,
    JUPYTER_SESSION_KIND,
    JUPYTER_SESSION_VERSION,
)
from renku_data_services.notebooks.crs import AmaltheaSessionV1Alpha1, JupyterServerV1Alpha1


def kr8s_async_api(
    url: Optional[str] = None,
    kubeconfig: Optional[str] = None,
    serviceaccount: Optional[str] = None,
    namespace: Optional[str] = None,
    context: Optional[str] = None,
) -> 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=url,
        kubeconfig=kubeconfig,
        serviceaccount=serviceaccount,
        namespace=namespace,
        context=context,
        _asyncio=True,  # This is the only line that is different from kr8s code
    )
    assert isinstance(ret, kr8s.asyncio.Api)
    return ret


class CRCValidatorProto(Protocol):
    """Compute resource control validator."""

    async def validate_class_storage(
        self,
        user: APIUser,
        class_id: int,
        storage: Optional[int] = None,
    ) -> ServerOptions:
        """Validate the resource class storage for the session."""
        ...

    async def get_default_class(self) -> ResourceClass:
        """Get the default resource class."""
        ...

    async def find_acceptable_class(
        self, user: APIUser, requested_server_options: ServerOptions
    ) -> Optional[ServerOptions]:
        """Find a suitable resource class based on resource requirements."""
        ...


class GitProviderHelperProto(Protocol):
    """Git provider protocol."""

    async def get_providers(self, user: APIUser) -> list[GitProvider]:
        """Get a list of git providers."""
        ...


class Kr8sApiStack:
    """Class maintaining a stack of current api clients.

    Used for testing.
    """

    stack: list[kr8s.Api] = list()

    def push(self, api: kr8s.Api) -> None:
        """Push a new api client onto the stack."""
        self.stack.append(api)

    def pop(self) -> kr8s.Api:
        """Pop the current kr8s api client from the stack."""
        return self.stack.pop()

    @property
    def current(self) -> kr8s.Api:
        """Get the currently active api client."""
        return self.stack[-1]

    def __getattribute__(self, name: str) -> Any:
        """Pass on requests to current api client."""
        if name in ["push", "pop", "current", "stack"]:
            return object.__getattribute__(self, name)
        return object.__getattribute__(self.current, name)


@dataclass
class NotebooksConfig:
    """The notebooks configuration."""

    server_options: _ServerOptionsConfig
    sessions: _SessionConfig
    sentry: _SentryConfig
    git: _GitConfig
    k8s: _K8sConfig
    k8s_cached_client: CachedK8sClient
    _kr8s_api: kr8s.asyncio.Api
    cloud_storage: _CloudStorage
    user_secrets: _UserSecrets
    crc_validator: CRCValidatorProto
    git_provider_helper: GitProviderHelperProto
    k8s_client: K8sClient[JupyterServerV1Alpha1]
    k8s_v2_client: K8sClient[AmaltheaSessionV1Alpha1]
    current_resource_schema_version: int = 1
    anonymous_sessions_enabled: bool = False
    ssh_enabled: bool = False
    service_prefix: str = "/notebooks"
    version: str = "0.0.0"
    keycloak_realm: str = "Renku"
    data_service_url: str = "http://renku-data-service"
    dummy_stores: bool = False
    session_get_endpoint_annotations: _ServersGetEndpointAnnotations = field(
        default_factory=_ServersGetEndpointAnnotations
    )
    session_id_cookie_name: str = "_renku_session"  # NOTE: This cookie name is set and controlled by the gateway

    @classmethod
    def from_env(cls, db_config: DBConfig) -> Self:
        """Create a configuration object from environment variables."""
        dummy_stores = _parse_str_as_bool(os.environ.get("DUMMY_STORES", False))
        sessions_config: _SessionConfig
        git_config: _GitConfig
        kr8s_api: kr8s.asyncio.Api
        data_service_url = os.environ.get("NB_DATA_SERVICE_URL", "http://127.0.0.1:8000")
        server_options = _ServerOptionsConfig.from_env()
        crc_validator: CRCValidatorProto
        git_provider_helper: GitProviderHelperProto
        k8s_namespace = os.environ.get("K8S_NAMESPACE", "default")
        quota_repo: QuotaRepository
        if dummy_stores:
            crc_validator = DummyCRCValidator()
            sessions_config = _SessionConfig._for_testing()
            git_provider_helper = DummyGitProviderHelper()
            git_config = _GitConfig("http://not.specified", "registry.not.specified")
            kr8s_api = Kr8sApiStack()  # type: ignore[assignment]
        else:
            quota_repo = QuotaRepository(K8sCoreClient(), K8sSchedulingClient(), namespace=k8s_namespace)
            rp_repo = ResourcePoolRepository(db_config.async_session_maker, quota_repo)
            crc_validator = CRCValidator(rp_repo)
            sessions_config = _SessionConfig.from_env()
            git_config = _GitConfig.from_env()
            git_provider_helper = GitProviderHelper(
                data_service_url, f"http://{sessions_config.ingress.host}", git_config.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 = kr8s_async_api()

        k8s_config = _K8sConfig.from_env()
        cluster_id = ClusterId("renkulab")
        clusters = {cluster_id: Cluster(id=cluster_id, namespace=k8s_config.renku_namespace, api=kr8s_api)}
        k8s_cached_client = CachedK8sClient(
            clusters=clusters,
            cache=K8sDbCache(db_config.async_session_maker),
            kinds_to_cache=[AMALTHEA_SESSION_KIND, JUPYTER_SESSION_KIND],
        )
        k8s_client = K8sClient(
            cached_client=k8s_cached_client,
            username_label="renku.io/userId",
            namespace=k8s_config.renku_namespace,
            cluster=cluster_id,
            server_type=JupyterServerV1Alpha1,
            server_kind=JUPYTER_SESSION_KIND,
            server_api_version=JUPYTER_SESSION_VERSION,
        )
        k8s_v2_client = K8sClient(
            cached_client=k8s_cached_client,
            # NOTE: v2 sessions have no userId label, the safe-username label is the keycloak user ID
            username_label="renku.io/safe-username",
            namespace=k8s_config.renku_namespace,
            cluster=cluster_id,
            server_type=AmaltheaSessionV1Alpha1,
            server_kind=AMALTHEA_SESSION_KIND,
            server_api_version=AMALTHEA_SESSION_VERSION,
        )
        return cls(
            server_options=server_options,
            sessions=sessions_config,
            sentry=_SentryConfig.from_env(),
            git=git_config,
            k8s=k8s_config,
            cloud_storage=_CloudStorage.from_env(),
            user_secrets=_UserSecrets.from_env(),
            current_resource_schema_version=1,
            anonymous_sessions_enabled=_parse_str_as_bool(os.environ.get("NB_ANONYMOUS_SESSIONS_ENABLED", False)),
            ssh_enabled=_parse_str_as_bool(os.environ.get("NB_SSH_ENABLED", False)),
            version=os.environ.get("NB_VERSION", "0.0.0"),
            keycloak_realm=os.environ.get("NB_KEYCLOAK_REALM", "Renku"),
            data_service_url=data_service_url,
            dummy_stores=dummy_stores,
            crc_validator=crc_validator,
            git_provider_helper=git_provider_helper,
            k8s_client=k8s_client,
            k8s_v2_client=k8s_v2_client,
            k8s_cached_client=k8s_cached_client,
            _kr8s_api=kr8s_api,
        )

1	"""Base motebooks svc configuration."""
2
3	import os	2✔
4	from dataclasses import dataclass, field	2✔
5	from typing import Any, Optional, Protocol, Self	2✔
6
7	import kr8s	2✔
8
9	from renku_data_services.base_models import APIUser	2✔
10	from renku_data_services.crc.db import ResourcePoolRepository	2✔
11	from renku_data_services.crc.models import ResourceClass	2✔
12	from renku_data_services.db_config.config import DBConfig	2✔
13	from renku_data_services.k8s.clients import K8sCoreClient, K8sSchedulingClient	2✔
14	from renku_data_services.k8s.quota import QuotaRepository	2✔
15	from renku_data_services.k8s_watcher.db import CachedK8sClient, Cluster, K8sDbCache	2✔
16	from renku_data_services.k8s_watcher.models import ClusterId	2✔
17	from renku_data_services.notebooks.api.classes.data_service import (	2✔
18	CRCValidator,
19	DummyCRCValidator,
20	DummyGitProviderHelper,
21	GitProviderHelper,
22	)
23	from renku_data_services.notebooks.api.classes.k8s_client import (	2✔
24	K8sClient,
25	)
26	from renku_data_services.notebooks.api.classes.repository import GitProvider	2✔
27	from renku_data_services.notebooks.api.schemas.server_options import ServerOptions	2✔
28	from renku_data_services.notebooks.config.dynamic import (	2✔
29	_CloudStorage,
30	_GitConfig,
31	_K8sConfig,
32	_parse_str_as_bool,
33	_SentryConfig,
34	_ServerOptionsConfig,
35	_SessionConfig,
36	_UserSecrets,
37	)
38	from renku_data_services.notebooks.config.static import _ServersGetEndpointAnnotations	2✔
39	from renku_data_services.notebooks.constants import (	2✔
40	AMALTHEA_SESSION_KIND,
41	AMALTHEA_SESSION_VERSION,
42	JUPYTER_SESSION_KIND,
43	JUPYTER_SESSION_VERSION,
44	)
45	from renku_data_services.notebooks.crs import AmaltheaSessionV1Alpha1, JupyterServerV1Alpha1	2✔
46
47
48	def kr8s_async_api(	2✔
49	url: Optional[str] = None,
50	kubeconfig: Optional[str] = None,
51	serviceaccount: Optional[str] = None,
52	namespace: Optional[str] = None,
53	context: Optional[str] = None,
54	) -> kr8s.asyncio.Api:
55	"""Create an async api client from sync code.
56
57	Kr8s cannot return an AsyncAPI instance from sync code, and we can't easily make all our config code async,
58	so this method is a direct copy of the kr8s sync client code, just that it returns an async client.
59	"""
NEW 60	ret = kr8s._async_utils.run_sync(kr8s.asyncio.api)(	×
61	url=url,
62	kubeconfig=kubeconfig,
63	serviceaccount=serviceaccount,
64	namespace=namespace,
65	context=context,
66	_asyncio=True, # This is the only line that is different from kr8s code
67	)
NEW 68	assert isinstance(ret, kr8s.asyncio.Api)	×
NEW 69	return ret	×
70
71
72	class CRCValidatorProto(Protocol):	2✔
73	"""Compute resource control validator."""
74
75	async def validate_class_storage(	2✔
76	self,
77	user: APIUser,
78	class_id: int,
79	storage: Optional[int] = None,
80	) -> ServerOptions:
81	"""Validate the resource class storage for the session."""
82	...	×
83
84	async def get_default_class(self) -> ResourceClass:	2✔
85	"""Get the default resource class."""
86	...	×
87
88	async def find_acceptable_class(	2✔
89	self, user: APIUser, requested_server_options: ServerOptions
90	) -> Optional[ServerOptions]:
91	"""Find a suitable resource class based on resource requirements."""
92	...	×
93
94
95	class GitProviderHelperProto(Protocol):	2✔
96	"""Git provider protocol."""
97
98	async def get_providers(self, user: APIUser) -> list[GitProvider]:	2✔
99	"""Get a list of git providers."""
100	...	×
101
102
103	class Kr8sApiStack:	2✔
104	"""Class maintaining a stack of current api clients.
105
106	Used for testing.
107	"""
108
109	stack: list[kr8s.Api] = list()	2✔
110
111	def push(self, api: kr8s.Api) -> None:	2✔
112	"""Push a new api client onto the stack."""
113	self.stack.append(api)	1✔
114
115	def pop(self) -> kr8s.Api:	2✔
116	"""Pop the current kr8s api client from the stack."""
117	return self.stack.pop()	1✔
118
119	@property	2✔
120	def current(self) -> kr8s.Api:	2✔
121	"""Get the currently active api client."""
122	return self.stack[-1]	1✔
123
124	def __getattribute__(self, name: str) -> Any:	2✔
125	"""Pass on requests to current api client."""
126	if name in ["push", "pop", "current", "stack"]:	1✔
127	return object.__getattribute__(self, name)	1✔
128	return object.__getattribute__(self.current, name)	1✔
129
130
131	@dataclass	2✔
132	class NotebooksConfig:	2✔
133	"""The notebooks configuration."""
134
135	server_options: _ServerOptionsConfig	2✔
136	sessions: _SessionConfig	2✔
137	sentry: _SentryConfig	2✔
138	git: _GitConfig	2✔
139	k8s: _K8sConfig	2✔
140	k8s_cached_client: CachedK8sClient	2✔
141	_kr8s_api: kr8s.asyncio.Api	2✔
142	cloud_storage: _CloudStorage	2✔
143	user_secrets: _UserSecrets	2✔
144	crc_validator: CRCValidatorProto	2✔
145	git_provider_helper: GitProviderHelperProto	2✔
146	k8s_client: K8sClient[JupyterServerV1Alpha1]	2✔
147	k8s_v2_client: K8sClient[AmaltheaSessionV1Alpha1]	2✔
148	current_resource_schema_version: int = 1	2✔
149	anonymous_sessions_enabled: bool = False	2✔
150	ssh_enabled: bool = False	2✔
151	service_prefix: str = "/notebooks"	2✔
152	version: str = "0.0.0"	2✔
153	keycloak_realm: str = "Renku"	2✔
154	data_service_url: str = "http://renku-data-service"	2✔
155	dummy_stores: bool = False	2✔
156	session_get_endpoint_annotations: _ServersGetEndpointAnnotations = field(	2✔
157	default_factory=_ServersGetEndpointAnnotations
158	)
159	session_id_cookie_name: str = "_renku_session" # NOTE: This cookie name is set and controlled by the gateway	2✔
160
161	@classmethod	2✔
162	def from_env(cls, db_config: DBConfig) -> Self:	2✔
163	"""Create a configuration object from environment variables."""
164	dummy_stores = _parse_str_as_bool(os.environ.get("DUMMY_STORES", False))	2✔
165	sessions_config: _SessionConfig
166	git_config: _GitConfig
167	kr8s_api: kr8s.asyncio.Api
168	data_service_url = os.environ.get("NB_DATA_SERVICE_URL", "http://127.0.0.1:8000")	2✔
169	server_options = _ServerOptionsConfig.from_env()	2✔
170	crc_validator: CRCValidatorProto
171	git_provider_helper: GitProviderHelperProto
172	k8s_namespace = os.environ.get("K8S_NAMESPACE", "default")	2✔
173	quota_repo: QuotaRepository
174	if dummy_stores:	2✔
175	crc_validator = DummyCRCValidator()	2✔
176	sessions_config = _SessionConfig._for_testing()	2✔
177	git_provider_helper = DummyGitProviderHelper()	2✔
178	git_config = _GitConfig("http://not.specified", "registry.not.specified")	2✔
179	kr8s_api = Kr8sApiStack() # type: ignore[assignment]	2✔
180	else:
181	quota_repo = QuotaRepository(K8sCoreClient(), K8sSchedulingClient(), namespace=k8s_namespace)	×
182	rp_repo = ResourcePoolRepository(db_config.async_session_maker, quota_repo)	×
183	crc_validator = CRCValidator(rp_repo)	×
184	sessions_config = _SessionConfig.from_env()	×
185	git_config = _GitConfig.from_env()	×
186	git_provider_helper = GitProviderHelper(	×
187	data_service_url, f"http://{sessions_config.ingress.host}", git_config.url
188	)
189	# NOTE: we need to get an async client as a sync client can't be used in an async way
190	# But all the config code is not async, so we need to drop into the running loop, if there is one
NEW 191	kr8s_api = kr8s_async_api()	×
192
193	k8s_config = _K8sConfig.from_env()	2✔
194	cluster_id = ClusterId("renkulab")	2✔
195	clusters = {cluster_id: Cluster(id=cluster_id, namespace=k8s_config.renku_namespace, api=kr8s_api)}	2✔
196	k8s_cached_client = CachedK8sClient(	2✔
197	clusters=clusters,
198	cache=K8sDbCache(db_config.async_session_maker),
199	kinds_to_cache=[AMALTHEA_SESSION_KIND, JUPYTER_SESSION_KIND],
200	)
201	k8s_client = K8sClient(	2✔
202	cached_client=k8s_cached_client,
203	username_label="renku.io/userId",
204	namespace=k8s_config.renku_namespace,
205	cluster=cluster_id,
206	server_type=JupyterServerV1Alpha1,
207	server_kind=JUPYTER_SESSION_KIND,
208	server_api_version=JUPYTER_SESSION_VERSION,
209	)
210	k8s_v2_client = K8sClient(	2✔
211	cached_client=k8s_cached_client,
212	# NOTE: v2 sessions have no userId label, the safe-username label is the keycloak user ID
213	username_label="renku.io/safe-username",
214	namespace=k8s_config.renku_namespace,
215	cluster=cluster_id,
216	server_type=AmaltheaSessionV1Alpha1,
217	server_kind=AMALTHEA_SESSION_KIND,
218	server_api_version=AMALTHEA_SESSION_VERSION,
219	)
220	return cls(	2✔
221	server_options=server_options,
222	sessions=sessions_config,
223	sentry=_SentryConfig.from_env(),
224	git=git_config,
225	k8s=k8s_config,
226	cloud_storage=_CloudStorage.from_env(),
227	user_secrets=_UserSecrets.from_env(),
228	current_resource_schema_version=1,
229	anonymous_sessions_enabled=_parse_str_as_bool(os.environ.get("NB_ANONYMOUS_SESSIONS_ENABLED", False)),
230	ssh_enabled=_parse_str_as_bool(os.environ.get("NB_SSH_ENABLED", False)),
231	version=os.environ.get("NB_VERSION", "0.0.0"),
232	keycloak_realm=os.environ.get("NB_KEYCLOAK_REALM", "Renku"),
233	data_service_url=data_service_url,
234	dummy_stores=dummy_stores,
235	crc_validator=crc_validator,
236	git_provider_helper=git_provider_helper,
237	k8s_client=k8s_client,
238	k8s_v2_client=k8s_v2_client,
239	k8s_cached_client=k8s_cached_client,
240	_kr8s_api=kr8s_api,
241	)

SwissDataScienceCenter / renku-data-services / 14382014257

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous