fix: patching of session custom resources

Author: olevski

components/renku_data_services/base_models/core.py

@@ -7,8 +7,9 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum, StrEnum
-from typing import ClassVar, Never, NewType, Optional, Protocol, Self, TypeVar, overload
+from typing import Annotated, ClassVar, Never, NewType, Optional, Protocol, Self, TypeVar, overload
 
+from pydantic import PlainSerializer
 from sanic import Request
 
 from renku_data_services.errors import errors
@@ -400,9 +401,11 @@ async def authenticate(self, access_token: str, request: Request) -> AnyAPIUser:
         ...
 
 
-ResetType = NewType("ResetType", object)
+__ResetType = NewType("__ResetType", object)
+ResetType = Annotated[__ResetType, PlainSerializer(lambda _: None, return_type=None)]
 """This type represents that a value that may be None should be reset back to None or null.
 This type should have only one instance, defined in the same file as this type.
+The value will be serialized by pydantic as None.
 """
 
 RESET: ResetType = ResetType(object())

components/renku_data_services/notebooks/core_sessions.py

@@ -5,7 +5,7 @@
 import os
 import random
 import string
-from collections.abc import AsyncIterator, Sequence
+from collections.abc import AsyncIterator, Mapping, Sequence
 from datetime import timedelta
 from pathlib import PurePosixPath
 from typing import Protocol, TypeVar, cast
@@ -20,7 +20,7 @@
 
 import renku_data_services.notebooks.image_check as ic
 from renku_data_services.app_config import logging
-from renku_data_services.base_models import AnonymousAPIUser, APIUser, AuthenticatedAPIUser
+from renku_data_services.base_models import RESET, AnonymousAPIUser, APIUser, AuthenticatedAPIUser, ResetType
 from renku_data_services.base_models.metrics import MetricsService
 from renku_data_services.connected_services.db import ConnectedServicesRepository
 from renku_data_services.crc.db import ClusterRepository, ResourcePoolRepository
@@ -54,6 +54,7 @@
     Authentication,
     AuthenticationType,
     Culling,
+    CullingPatch,
     DataSource,
     ExtraContainer,
     ExtraVolume,
@@ -69,6 +70,7 @@
     Requests,
     RequestsStr,
     Resources,
+    ResourcesPatch,
     SecretAsVolume,
     SecretAsVolumeItem,
     Session,
@@ -91,6 +93,7 @@
 )
 from renku_data_services.notebooks.utils import (
     node_affinity_from_resource_class,
+    node_affinity_patch_from_resource_class,
     tolerations_from_resource_class,
 )
 from renku_data_services.project.db import ProjectRepository, ProjectSessionSecretRepository
@@ -462,6 +465,21 @@ async def request_session_secret_creation(
             )
 
 
+def resources_patch_from_resource_class(resource_class: ResourceClass) -> ResourcesPatch:
+    """Convert the resource class to a k8s resources spec."""
+    gpu_name = GpuKind.NVIDIA.value + "/gpu"
+    resources = resources_from_resource_class(resource_class)
+    requests: Mapping[str, Requests | RequestsStr | ResetType] | ResetType
+    limits: Mapping[str, Limits | LimitsStr | ResetType] | ResetType
+    defaul_requests = {"memory": RESET, "cpu": RESET, gpu_name: RESET}
+    default_limits = {"memory": RESET, "cpu": RESET, gpu_name: RESET}
+    if resources.requests:
+        requests = RESET if len(resources.requests.keys()) == 0 else {**defaul_requests, **resources.requests}
+    if resources.limits:
+        limits = RESET if len(resources.limits.keys()) == 0 else {**default_limits, **resources.limits}
+    return ResourcesPatch(requests=requests, limits=limits)
+
+
 def resources_from_resource_class(resource_class: ResourceClass) -> Resources:
     """Convert the resource class to a k8s resources spec."""
     requests: dict[str, Requests | RequestsStr] = {
@@ -528,6 +546,31 @@ def get_culling(
     )
 
 
+def get_culling_patch(
+    user: AuthenticatedAPIUser | AnonymousAPIUser, resource_pool: ResourcePool, nb_config: NotebooksConfig
+) -> CullingPatch:
+    """Get the patch for the culling durations of a session."""
+    culling = get_culling(user, resource_pool, nb_config)
+    patch = CullingPatch(
+        maxAge=RESET,
+        maxFailedDuration=RESET,
+        maxHibernatedDuration=RESET,
+        maxIdleDuration=RESET,
+        maxStartingDuration=RESET,
+    )
+    if culling.maxAge:
+        patch.maxAge = culling.maxAge
+    if culling.maxFailedDuration:
+        patch.maxFailedDuration = culling.maxFailedDuration
+    if culling.maxHibernatedDuration:
+        patch.maxHibernatedDuration = culling.maxHibernatedDuration
+    if culling.maxIdleDuration:
+        patch.maxIdleDuration = culling.maxIdleDuration
+    if culling.maxStartingDuration:
+        patch.maxStartingDuration = culling.maxStartingDuration
+    return patch
+
+
 async def __requires_image_pull_secret(nb_config: NotebooksConfig, image: str, internal_gitlab_user: APIUser) -> bool:
     """Determines if an image requires a pull secret based on its visibility and their GitLab access token."""
 
@@ -1030,29 +1073,39 @@ async def patch_session(
                 )
             )
         rp = await rp_repo.get_resource_pool_from_class(user, body.resource_class_id)
+        try:
+            old_rp = await rp_repo.get_resource_pool_from_class(user, session.resource_class_id())
+        except (errors.MissingResourceError, errors.UnauthorizedError, errors.ForbiddenError):
+            old_rp = None
         rc = rp.get_resource_class(body.resource_class_id)
         if not rc:
             raise errors.MissingResourceError(
                 message=f"The resource class you requested with ID {body.resource_class_id} does not exist"
             )
-        # TODO: reject session classes which change the cluster
+        if old_rp is not None and rp.cluster != old_rp.cluster:
+            raise errors.ValidationError(message="Changing resource pools with different clusters is not allowed.")
         if not patch.metadata:
             patch.metadata = AmaltheaSessionV1Alpha1MetadataPatch()
-        # Patch the resource class ID in the annotations
+        # Patch the resource pool and class ID in the annotations
+        patch.metadata.annotations = {"renku.io/resource_pool_id": str(rp.id)}
         patch.metadata.annotations = {"renku.io/resource_class_id": str(body.resource_class_id)}
         if not patch.spec.session:
             patch.spec.session = AmaltheaSessionV1Alpha1SpecSessionPatch()
-        patch.spec.session.resources = resources_from_resource_class(rc)
+        patch.spec.session.resources = resources_patch_from_resource_class(rc)
         # Tolerations
         tolerations = tolerations_from_resource_class(rc, nb_config.sessions.tolerations_model)
         patch.spec.tolerations = tolerations
         # Affinities
-        patch.spec.affinity = node_affinity_from_resource_class(rc, nb_config.sessions.affinity_model)
+        patch.spec.affinity = node_affinity_patch_from_resource_class(rc, nb_config.sessions.affinity_model)
         # Priority class (if a quota is being used)
-        patch.spec.priorityClassName = rc.quota
-        patch.spec.culling = get_culling(user, rp, nb_config)
+        if rc.quota is None:
+            patch.spec.priorityClassName = RESET
+        patch.spec.culling = get_culling_patch(user, rp, nb_config)
+        # Service account name
         if rp.cluster is not None:
-            patch.spec.service_account_name = rp.cluster.service_account_name
+            patch.spec.service_account_name = (
+                rp.cluster.service_account_name if rp.cluster.service_account_name is not None else RESET
+            )
 
     # If the session is being hibernated we do not need to patch anything else that is
     # not specifically called for in the request body, we can refresh things when the user resumes.
@@ -1126,6 +1179,8 @@ async def patch_session(
     if image_pull_secret:
         session_extras.concat(SessionExtraResources(secrets=[image_pull_secret]))
         patch.spec.imagePullSecrets = [ImagePullSecret(name=image_pull_secret.name, adopt=image_pull_secret.adopt)]
+    else:
+        patch.spec.imagePullSecrets = RESET
 
     # Construct session patch
     patch.spec.extraContainers = _make_patch_spec_list(

components/renku_data_services/notebooks/crs.py

@@ -14,6 +14,7 @@
 from pydantic.types import HashableItemType
 from ulid import ULID
 
+from renku_data_services.base_models.core import ResetType
 from renku_data_services.errors import errors
 from renku_data_services.notebooks import apispec
 from renku_data_services.notebooks.constants import AMALTHEA_SESSION_GVK, JUPYTER_SESSION_GVK
@@ -33,6 +34,8 @@
     MatchExpression,
     NodeAffinity,
     NodeSelectorTerm,
+    PodAffinity,
+    PodAntiAffinity,
     Preference,
     PreferredDuringSchedulingIgnoredDuringExecutionItem,
     ReconcileStrategy,
@@ -375,36 +378,61 @@ def base_url(self) -> str | None:
         return url
 
 
+class ResourcesPatch(BaseCRD):
+    """Resource requests and limits patch."""
+
+    limits: Mapping[str, LimitsStr | Limits | ResetType] | ResetType | None = None
+    requests: Mapping[str, RequestsStr | Requests | ResetType] | ResetType | None = None
+
+
 class AmaltheaSessionV1Alpha1SpecSessionPatch(BaseCRD):
     """Patch for the main session config."""
 
-    resources: Resources | None = None
-    shmSize: int | str | None = None
-    storage: Storage | None = None
+    resources: ResourcesPatch | ResetType | None = None
+    shmSize: int | str | ResetType | None = None
+    storage: Storage | ResetType | None = None
     imagePullPolicy: ImagePullPolicy | None = None
-    extraVolumeMounts: list[ExtraVolumeMount] | None = None
+    extraVolumeMounts: list[ExtraVolumeMount] | ResetType | None = None
 
 
 class AmaltheaSessionV1Alpha1MetadataPatch(BaseCRD):
     """Patch for the metadata of an amalthea session."""
 
-    annotations: dict[str, str] | None = None
+    annotations: dict[str, str | ResetType] | ResetType | None = None
+
+
+class AffinityPatch(BaseCRD):
+    """Patch for the affinity of a session."""
+
+    nodeAffinity: NodeAffinity | ResetType | None = None
+    podAffinity: PodAffinity | ResetType | None = None
+    podAntiAffinity: PodAntiAffinity | ResetType | None = None
+
+
+class CullingPatch(Culling):
+    """Patch for the culling durations of a session."""
+
+    maxAge: timedelta | ResetType | None = None  # type:ignore[assignment]
+    maxFailedDuration: timedelta | ResetType | None = None  # type:ignore[assignment]
+    maxHibernatedDuration: timedelta | ResetType | None = None  # type:ignore[assignment]
+    maxIdleDuration: timedelta | ResetType | None = None  # type:ignore[assignment]
+    maxStartingDuration: timedelta | ResetType | None = None  # type:ignore[assignment]
 
 
 class AmaltheaSessionV1Alpha1SpecPatch(BaseCRD):
     """Patch for the spec of an amalthea session."""
 
-    extraContainers: list[ExtraContainer] | None = None
-    extraVolumes: list[ExtraVolume] | None = None
+    extraContainers: list[ExtraContainer] | ResetType | None = None
+    extraVolumes: list[ExtraVolume] | ResetType | None = None
     hibernated: bool | None = None
-    initContainers: list[InitContainer] | None = None
-    imagePullSecrets: list[ImagePullSecret] | None = None
-    priorityClassName: str | None = None
-    tolerations: list[Toleration] | None = None
-    affinity: Affinity | None = None
+    initContainers: list[InitContainer] | ResetType | None = None
+    imagePullSecrets: list[ImagePullSecret] | ResetType | None = None
+    priorityClassName: str | ResetType | None = None
+    tolerations: list[Toleration] | ResetType | None = None
+    affinity: AffinityPatch | ResetType | None = None
     session: AmaltheaSessionV1Alpha1SpecSessionPatch | None = None
-    culling: Culling | None = None
-    service_account_name: str | None = None
+    culling: CullingPatch | ResetType | None = None
+    service_account_name: str | ResetType | None = None
 
 
 class AmaltheaSessionV1Alpha1Patch(BaseCRD):
@@ -414,8 +442,12 @@ class AmaltheaSessionV1Alpha1Patch(BaseCRD):
     spec: AmaltheaSessionV1Alpha1SpecPatch
 
     def to_rfc7386(self) -> dict[str, Any]:
-        """Generate the patch to be applied to the session."""
-        return self.model_dump(exclude_none=True)
+        """Generate the patch to be applied to the session.
+
+        Note that when the value for a key in the patch is anything other than a
+        dictionary then the rfc7386 patch will replace, not merge.
+        """
+        return self.model_dump(exclude_none=True, mode="json")
 
 
 def safe_parse_duration(val: Any) -> timedelta:

components/renku_data_services/notebooks/utils.py

@@ -1,8 +1,10 @@
 """Utilities for notebooks."""
 
 import renku_data_services.crc.models as crc_models
+from renku_data_services.base_models.core import RESET
 from renku_data_services.notebooks.crs import (
     Affinity,
+    AffinityPatch,
     MatchExpression,
     NodeAffinity,
     NodeSelectorTerm,
@@ -128,6 +130,21 @@ def node_affinity_from_resource_class(
     return affinity
 
 
+def node_affinity_patch_from_resource_class(
+    resource_class: crc_models.ResourceClass, default_affinity: Affinity
+) -> AffinityPatch:
+    """Create a patch for the session affinity."""
+    affinity = node_affinity_from_resource_class(resource_class, default_affinity)
+    patch = AffinityPatch(nodeAffinity=RESET, podAffinity=RESET, podAntiAffinity=RESET)
+    if affinity.nodeAffinity:
+        patch.nodeAffinity = affinity.nodeAffinity
+    if affinity.podAffinity:
+        patch.podAffinity = affinity.podAffinity
+    if affinity.podAntiAffinity:
+        patch.podAntiAffinity = affinity.podAntiAffinity
+    return patch
+
+
 def tolerations_from_resource_class(
     resource_class: crc_models.ResourceClass, default_tolerations: list[Toleration]
 ) -> list[Toleration]: