20339056754

Committed 18 Dec 2025 01:46PM UTC coverage: 86.009%. First build

Build # 20339056754

Build Type

Pull #1128

github

Committed by

web-flow

Commit Message

Merge fe560b56c into c5f960729

Pull Request Pull Request #1128: feat: set lastInteraction time via session patch and return willHibernateAt

Run Details

84 of 96 new or added lines in 11 files covered. (87.5%)

24037 of 27947 relevant lines covered (86.01%)

1.51 hits per line

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

87.1

/components/renku_data_services/notebooks/crs.py

"""Custom resource definition with proper names from the autogenerated code."""

from __future__ import annotations

import re
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta
from typing import Any, cast, override
from urllib.parse import urlunparse

from kubernetes.utils import parse_duration, parse_quantity
from kubernetes.utils.duration import format_duration
from pydantic import BaseModel, Field, field_serializer, field_validator, model_serializer
from pydantic.types import HashableItemType
from ulid import ULID

from renku_data_services.base_models.core import RESET, 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
from renku_data_services.notebooks.cr_amalthea_session import Affinity as _Affinity
from renku_data_services.notebooks.cr_amalthea_session import (
    Authentication,
    CodeRepository,
    DataSource,
    EmptyDir,
    ExtraContainer,
    ExtraVolume,
    ExtraVolumeMount,
    ImagePullPolicy,
    ImagePullSecret,
    Ingress,
    InitContainer,
    MatchExpression,
    NodeAffinity,
    NodeSelectorTerm,
    PodAffinity,
    PodAntiAffinity,
    Preference,
    PreferredDuringSchedulingIgnoredDuringExecutionItem,
    ReconcileStrategy,
    RemoteSecretRef,
    RequiredDuringSchedulingIgnoredDuringExecution,
    RequiredDuringSchedulingIgnoredDuringExecutionItem,
    Size,
    State,
    Status,
    Storage,
    TlsSecret,
    Toleration,
)
from renku_data_services.notebooks.cr_amalthea_session import (
    Culling as _ASCulling,
)
from renku_data_services.notebooks.cr_amalthea_session import EnvItem2 as SessionEnvItem
from renku_data_services.notebooks.cr_amalthea_session import Item4 as SecretAsVolumeItem
from renku_data_services.notebooks.cr_amalthea_session import Limits6 as _Limits
from renku_data_services.notebooks.cr_amalthea_session import Limits7 as LimitsStr
from renku_data_services.notebooks.cr_amalthea_session import Location as SessionLocation
from renku_data_services.notebooks.cr_amalthea_session import Model as _ASModel
from renku_data_services.notebooks.cr_amalthea_session import (
    PreferredDuringSchedulingIgnoredDuringExecutionItem1 as PreferredPodAffinityItem,
)
from renku_data_services.notebooks.cr_amalthea_session import (
    PreferredDuringSchedulingIgnoredDuringExecutionItem2 as PreferredPodAntiAffinityItem,
)
from renku_data_services.notebooks.cr_amalthea_session import Requests6 as _Requests
from renku_data_services.notebooks.cr_amalthea_session import Requests7 as RequestsStr
from renku_data_services.notebooks.cr_amalthea_session import (
    RequiredDuringSchedulingIgnoredDuringExecutionItem1 as RequiredPodAntiAffinityItem,
)
from renku_data_services.notebooks.cr_amalthea_session import Resources3 as _Resources
from renku_data_services.notebooks.cr_amalthea_session import Secret1 as SecretAsVolume
from renku_data_services.notebooks.cr_amalthea_session import SecretRef as _SecretRef
from renku_data_services.notebooks.cr_amalthea_session import Session as _ASSession
from renku_data_services.notebooks.cr_amalthea_session import ShmSize1 as ShmSizeStr
from renku_data_services.notebooks.cr_amalthea_session import Size1 as SizeStr
from renku_data_services.notebooks.cr_amalthea_session import Spec as _ASSpec
from renku_data_services.notebooks.cr_amalthea_session import Type as AuthenticationType
from renku_data_services.notebooks.cr_amalthea_session import Type1 as CodeRepositoryType
from renku_data_services.notebooks.cr_base import BaseCRD
from renku_data_services.notebooks.cr_jupyter_server import Model as _JSModel
from renku_data_services.notebooks.cr_jupyter_server import Patch
from renku_data_services.notebooks.cr_jupyter_server import Spec as JupyterServerSpec
from renku_data_services.notebooks.cr_jupyter_server import Type as PatchType


class Metadata(BaseModel):
    """Basic k8s metadata spec."""

    class Config:
        """Do not exclude unknown properties."""

        extra = "allow"

    name: str
    namespace: str | None = None
    labels: dict[str, str] = Field(default_factory=dict)
    annotations: dict[str, str] = Field(default_factory=dict)
    uid: str | None = None
    creationTimestamp: datetime | None = None
    deletionTimestamp: datetime | None = None


class ComputeResources(BaseModel):
    """Resource requests from k8s values."""

    cpu: float | None = None
    memory: int | None = None
    storage: int | None = None
    gpu: int | None = None

    @field_validator("cpu", mode="before")
    @classmethod
    def _convert_k8s_cpu(cls, val: Any) -> Any:
        if val is None:
            return None
        if hasattr(val, "root"):
            val = val.root
        return float(parse_quantity(val))

    @field_validator("gpu", mode="before")
    @classmethod
    def _convert_k8s_gpu(cls, val: Any) -> Any:
        if val is None:
            return None
        if hasattr(val, "root"):
            val = val.root
        return round(parse_quantity(val), ndigits=None)

    @field_validator("memory", "storage", mode="before")
    @classmethod
    def _convert_k8s_bytes(cls, val: Any) -> Any:
        """Converts to gigabytes of base 10."""
        if val is None:
            return None
        if hasattr(val, "root"):
            val = val.root
        return round(parse_quantity(val) / 1_000_000_000, ndigits=None)


class JupyterServerV1Alpha1(_JSModel):
    """Jupyter server CRD."""

    kind: str = JUPYTER_SESSION_GVK.kind
    apiVersion: str = JUPYTER_SESSION_GVK.group_version
    metadata: Metadata

    def get_compute_resources(self) -> ComputeResources:
        """Convert the k8s resource requests and storage into usable values."""
        if self.spec is None:
            return ComputeResources()
        resource_requests: dict = self.spec.jupyterServer.resources.get("requests", {})
        resource_requests["storage"] = self.spec.storage.size
        return ComputeResources.model_validate(resource_requests)

    def resource_class_id(self) -> int:
        """Get the resource class from the annotations."""
        if "renku.io/resourceClassId" not in self.metadata.annotations:
            raise errors.ProgrammingError(
                message=f"The session with name {self.metadata.name} is missing its renku.io/resourceClassId annotation"
            )
        i = int(self.metadata.annotations["renku.io/resourceClassId"])
        return i


class CullingDurationParsingMixin(BaseCRD):
    """Amalthea session culling configuration."""

    @field_serializer("*", mode="wrap")
    def __serialize_duration_field(self, val: Any, nxt: Any, _info: Any) -> Any:
        if isinstance(val, timedelta):
            return format_duration(val)
        if val is RESET:
            return None
        return nxt(val)

    @field_validator(
        "maxAge",
        "maxFailedDuration",
        "maxHibernatedDuration",
        "maxIdleDuration",
        "maxStartingDuration",
        mode="wrap",
        check_fields=False,
    )
    @classmethod
    def __deserialize_duration(cls, val: Any, handler: Any) -> Any:
        if isinstance(val, str):
            try:
                return safe_parse_duration(val)
            except Exception:
                return handler(val)
        return handler(val)


class Culling(_ASCulling, CullingDurationParsingMixin):
    """Amalthea session culling configuration.

    A value of zero for any of the values indicates that the automatic action will never happen.
    I.e. a value of zero for `maxIdleDuration` indicates that the session will never be hibernated
    no matter how long it is idle.
    """

    pass


class Requests(_Requests):
    """Resource requests of type integer."""

    root: int


class Limits(_Limits):
    """Resource limits of type integer."""

    root: int


class Resources(_Resources):
    """Resource requests and limits spec.

    Overriding these is necessary because of
    https://docs.pydantic.dev/2.11/errors/validation_errors/#string_type.
    An integer model cannot have a regex pattern for validation in pydantic.
    But the code generation applies the pattern constraint to both the int and string variations
    of the fields. But the int variation runs and blows up at runtime only when an int is passed
    for validation.
    """

    limits: Mapping[str, LimitsStr | Limits] | None = None
    requests: Mapping[str, RequestsStr | Requests] | None = None


class SecretRef(_SecretRef, RemoteSecretRef):
    """Reference to a secret."""

    pass


class Session(_ASSession):
    """Amalthea spec.session schema."""

    resources: Resources | None = None
    remoteSecretRef: SecretRef | None = None


class AmaltheaSessionSpec(_ASSpec):
    """Amalthea session specification."""

    session: Session
    culling: Culling | None = None


class AmaltheaSessionV1Alpha1(_ASModel):
    """Amalthea session CRD."""

    kind: str = AMALTHEA_SESSION_GVK.kind
    apiVersion: str = AMALTHEA_SESSION_GVK.group_version
    # Here we overwrite the default from ASModel because it is too weakly typed
    metadata: Metadata  # type: ignore[assignment]
    spec: AmaltheaSessionSpec

    def get_compute_resources(self) -> ComputeResources:
        """Convert the k8s resource requests and storage into usable values."""
        resource_requests: dict = {}
        if self.spec.session.resources is not None:
            reqs = self.spec.session.resources.requests or {}
            reqs = {k: parse_quantity(v.root if hasattr(v, "root") else v) for k, v in reqs.items()}
            resource_requests = {
                **reqs,
                "storage": parse_quantity(self.spec.session.storage.size.root),
            }
        return ComputeResources.model_validate(resource_requests)

    @property
    def project_id(self) -> ULID:
        """Get the project ID from the annotations."""
        if "renku.io/project_id" not in self.metadata.annotations:
            raise errors.ProgrammingError(
                message=f"The session with name {self.metadata.name} is missing its project_id annotation"
            )
        return cast(ULID, ULID.from_str(self.metadata.annotations["renku.io/project_id"]))

    @property
    def launcher_id(self) -> ULID:
        """Get the launcher ID from the annotations."""
        if "renku.io/launcher_id" not in self.metadata.annotations:
            raise errors.ProgrammingError(
                message=f"The session with name {self.metadata.name} is missing its launcher_id annotation"
            )
        return cast(ULID, ULID.from_str(self.metadata.annotations["renku.io/launcher_id"]))

    def resource_class_id(self) -> int:
        """Get the resource class from the annotations."""
        if "renku.io/resource_class_id" not in self.metadata.annotations:
            raise errors.ProgrammingError(
                message=f"The session with name {self.metadata.name} is missing its resource_class_id annotation"
            )
        return int(self.metadata.annotations["renku.io/resource_class_id"])

    def as_apispec(self) -> apispec.SessionResponse:
        """Convert the manifest into a form ready to be serialized and sent in a HTTP response."""
        if self.status is None:
            raise errors.ProgrammingError(
                message=f"The manifest for a session with name {self.metadata.name} cannot be serialized "
                f"because it is missing a status"
            )
        if self.spec is None:
            raise errors.ProgrammingError(
                message=f"The manifest for a session with name {self.metadata.name} cannot be serialized "
                "because it is missing the spec field"
            )
        if self.spec.session.resources is None:
            raise errors.ProgrammingError(
                message=f"The manifest for a session with name {self.metadata.name} cannot be serialized "
                "because it is missing the spec.session.resources field"
            )
        url = "None"
        if self.base_url is not None:
            url = self.base_url
        ready_containers = 0
        total_containers = 0
        if self.status.initContainerCounts is not None:
            ready_containers += self.status.initContainerCounts.ready or 0
            total_containers += self.status.initContainerCounts.total or 0
        if self.status.containerCounts is not None:
            ready_containers += self.status.containerCounts.ready or 0
            total_containers += self.status.containerCounts.total or 0

        if self.status.state in [State.Running, State.Hibernated, State.Failed]:
            # Amalthea is sometimes slow when (un)hibernating and still shows the old status, so we patch it here
            # so the client sees the correct state
            if not self.spec.hibernated and self.status.state == State.Hibernated:
                state = apispec.State3.starting
            elif self.spec.hibernated and self.status.state == State.Running:
                state = apispec.State3.hibernated
            else:
                state = apispec.State3(self.status.state.value.lower())
        elif self.status.state == State.RunningDegraded:
            state = apispec.State3.running
        elif self.status.state == State.NotReady and self.metadata.deletionTimestamp is not None:
            state = apispec.State3.stopping
        else:
            state = apispec.State3.starting

        will_delete_at: datetime | None = None
        match self.status, self.spec.culling:
            case (
                Status(state=State.Hibernated, hibernatedSince=hibernated_since),
                Culling(maxHibernatedDuration=max_hibernated),
            ) if hibernated_since and max_hibernated:
                will_delete_at = hibernated_since + max_hibernated

        return apispec.SessionResponse(
            image=self.spec.session.image,
            name=self.metadata.name,
            resources=apispec.SessionResources(
                requests=apispec.SessionResourcesRequests.model_validate(
                    self.get_compute_resources(), from_attributes=True
                )
                if self.spec.session.resources.requests is not None
                else None,
            ),
            started=self.metadata.creationTimestamp,
            lastInteraction=self.spec.culling.lastInteraction if self.spec.culling else None,
            status=apispec.SessionStatus(
                state=state,
                ready_containers=ready_containers,
                total_containers=total_containers,
                will_hibernate_at=self.status.willHibernateAt,
                will_delete_at=will_delete_at,
                message=self.status.error,
            ),
            url=url,
            project_id=str(self.project_id),
            launcher_id=str(self.launcher_id),
            resource_class_id=self.resource_class_id(),
        )

    @property
    def base_url(self) -> str | None:
        """Get the URL of the session, excluding the default URL from the session launcher."""
        if self.status.url and len(self.status.url) > 0:
            return self.status.url
        if self.spec is None or self.spec.ingress is None:
            return None
        scheme = "https" if self.spec and self.spec.ingress and self.spec.ingress.tlsSecret else "http"
        host = self.spec.ingress.host
        path = self.spec.session.urlPath if self.spec.session.urlPath else "/"
        if not path.endswith("/"):
            path += "/"
        params = None
        query = None
        fragment = None
        url = urlunparse((scheme, host, path, params, query, fragment))
        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: ResourcesPatch | ResetType | None = None
    shmSize: int | str | ResetType | None = None
    storage: Storage | ResetType | None = None
    imagePullPolicy: ImagePullPolicy | None = None
    extraVolumeMounts: list[ExtraVolumeMount] | ResetType | None = None


class AmaltheaSessionV1Alpha1MetadataPatch(BaseCRD):
    """Patch for the metadata of an amalthea session."""

    annotations: dict[str, str | ResetType] | ResetType | None = None


class NodeAffinityPatch(BaseCRD):
    """Patch for the node affinity of a session."""

    preferredDuringSchedulingIgnoredDuringExecution: (
        Sequence[PreferredDuringSchedulingIgnoredDuringExecutionItem] | None | ResetType
    ) = None
    requiredDuringSchedulingIgnoredDuringExecution: (
        RequiredDuringSchedulingIgnoredDuringExecution | None | ResetType
    ) = None


class PodAffinityPatch(BaseCRD):
    """Patch for the pod affinity of a session."""

    preferredDuringSchedulingIgnoredDuringExecution: Sequence[PreferredPodAffinityItem] | None | ResetType = None
    requiredDuringSchedulingIgnoredDuringExecution: (
        Sequence[RequiredDuringSchedulingIgnoredDuringExecutionItem] | None | ResetType
    ) = None


class PodAntiAffinityPatch(BaseCRD):
    """Patch for the pod anti affinity of a session."""

    preferredDuringSchedulingIgnoredDuringExecution: Sequence[PreferredPodAntiAffinityItem] | None | ResetType = None
    requiredDuringSchedulingIgnoredDuringExecution: Sequence[RequiredPodAntiAffinityItem] | None | ResetType = None


class AffinityPatch(BaseCRD):
    """Patch for the affinity of a session."""

    nodeAffinity: NodeAffinityPatch | ResetType | None = None
    podAffinity: PodAffinityPatch | ResetType | None = None
    podAntiAffinity: PodAntiAffinityPatch | ResetType | None = None


class CullingPatch(CullingDurationParsingMixin):
    """Patch for the culling durations of a session."""

    maxAge: timedelta | ResetType | None = None
    maxFailedDuration: timedelta | ResetType | None = None
    maxHibernatedDuration: timedelta | ResetType | None = None
    maxIdleDuration: timedelta | ResetType | None = None
    maxStartingDuration: timedelta | ResetType | None = None
    lastInteraction: datetime | ResetType | None = None


class AmaltheaSessionV1Alpha1SpecPatch(BaseCRD):
    """Patch for the spec of an amalthea session."""

    extraContainers: list[ExtraContainer] | ResetType | None = None
    extraVolumes: list[ExtraVolume] | ResetType | None = None
    hibernated: bool | 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: CullingPatch | ResetType | None = None
    service_account_name: str | ResetType | None = None


class AmaltheaSessionV1Alpha1Patch(BaseCRD):
    """Patch for an amalthea session."""

    metadata: AmaltheaSessionV1Alpha1MetadataPatch | None = None
    spec: AmaltheaSessionV1Alpha1SpecPatch

    def to_rfc7386(self) -> dict[str, Any]:
        """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:
    """Required because parse_duration from k8s can only deal with values with up to 5 digits.

    Values with 1 unit only (like seconds) and high values will cause things to fail.
    For example `parse_duration("1210500s")` will raise ValueError whereas `parse_duration("100s")` will be fine.
    This does not make the whole thing 100% foolproof but it eliminates errors like the above which
    we have seen in production.
    """
    if isinstance(val, timedelta):
        return val
    m = re.match(r"^([0-9]+)(h|m|s|ms)$", str(val))
    if m is not None:
        num = m.group(1)
        unit = m.group(2)
        match unit:
            case "h":
                return timedelta(hours=float(num))
            case "m":
                return timedelta(minutes=float(num))
            case "s":
                return timedelta(seconds=float(num))
            case "ms":
                return timedelta(microseconds=float(num))
    return cast(timedelta, parse_duration(val))


class Affinity(_Affinity):
    """Model for the affinity of an Amalthea session."""

    @property
    def is_empty(self) -> bool:
        """Indicates whether there is anything set at all in any of the affinity fields, or they are just all empty."""
        node_affinity_is_empty = not (
            self.nodeAffinity
            and (
                self.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution
                or self.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution
            )
        )
        pod_affinity_is_empty = not (
            self.podAffinity
            and (
                self.podAffinity.preferredDuringSchedulingIgnoredDuringExecution
                or self.podAffinity.requiredDuringSchedulingIgnoredDuringExecution
            )
        )
        pod_antiaffinity_is_empty = not (
            self.podAntiAffinity
            and (
                self.podAntiAffinity.preferredDuringSchedulingIgnoredDuringExecution
                or self.podAntiAffinity.requiredDuringSchedulingIgnoredDuringExecution
            )
        )
        return node_affinity_is_empty and pod_affinity_is_empty and pod_antiaffinity_is_empty

SwissDataScienceCenter / renku-data-services / 20339056754

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