k8s.py

import logging
from collections.abc import Generator
from pathlib import Path
from typing import Literal

from kubernetes import client, config, dynamic

from jobq_server.exceptions import PodNotReadyError, WorkloadNotFound
from jobq_server.models import JobId
from jobq_server.utils.helpers import traverse
from jobq_server.utils.k8s import GroupVersionKind
from jobq_server.utils.kueue import KueueWorkload


class KubernetesService:
    def __init__(self):
        try:
            config.load_incluster_config()
            self._in_cluster = True
        except config.ConfigException:
            logging.warning(
                "Could not load in-cluster config, attempting to load Kubeconfig",
            )
            config.load_kube_config()
            self._in_cluster = False
        self._core_v1_api = client.CoreV1Api()

    @property
    def namespace(self) -> str:
        if not self._in_cluster:
            _, active_context = config.list_kube_config_contexts()
            current_namespace = traverse(active_context, "context.namespace")
        else:
            # When running in a cluster, determine the namespace from the mounted service account
            try:
                current_namespace = (
                    Path("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
                    .read_text()
                    .strip()
                )
            except FileNotFoundError as e:
                raise RuntimeError("Could not determine current namespace") from e
        return current_namespace

    def workload_for_managed_resource(
        self, uid: JobId, namespace: str | None = None
    ) -> KueueWorkload | None:
        try:
            return KueueWorkload.for_managed_resource(
                uid, namespace=namespace or self.namespace
            )
        except WorkloadNotFound:
            return None

    def _sanitize_log_kwargs(self, tail: int) -> dict[str, int]:
        return {"tail_lines": tail} if tail != -1 else {}

    def get_pod_logs(self, pod: client.V1Pod, tail: int = -1) -> str:
        try:
            return self._core_v1_api.read_namespaced_pod_log(
                pod.metadata.name,
                pod.metadata.namespace,
                **self._sanitize_log_kwargs(tail),
            )
        except client.ApiException as e:
            if e.status == 400:
                raise PodNotReadyError(
                    name=pod.metadata.name,
                    namespace=pod.metadata.namespace,
                ) from e
            raise

    def stream_pod_logs(
        self, pod: client.V1Pod, tail: int = -1
    ) -> Generator[str, None, None]:
        try:
            log_stream = self._core_v1_api.read_namespaced_pod_log(
                pod.metadata.name,
                pod.metadata.namespace,
                follow=True,
                _preload_content=False,
                **self._sanitize_log_kwargs(tail),
            )
            yield from log_stream
        except client.ApiException as e:
            if e.status == 400:
                raise PodNotReadyError(
                    name=pod.metadata.name,
                    namespace=pod.metadata.namespace,
                ) from e
            raise

    def delete_resource(
        self,
        gvk: GroupVersionKind,
        name: str,
        namespace: str,
        propagation_policy: Literal[
            "Foreground", "Background", "Orphan"
        ] = "Foreground",
    ) -> None:
        dyn = dynamic.DynamicClient(client.ApiClient())

        resource = dyn.resources.get(
            api_version=f"{gvk.group}/{gvk.version}" if gvk.group else gvk.version,
            kind=gvk.kind,
        )

        dyn.delete(
            resource,
            name=name,
            namespace=namespace,
            body=client.V1DeleteOptions(propagation_policy=propagation_policy),
        )

    def list_workloads(self, namespace: str | None = None) -> list[KueueWorkload]:
        api = client.CustomObjectsApi()
        workloads = api.list_namespaced_custom_object(
            group="kueue.x-k8s.io",
            version="v1beta1",
            namespace=namespace or self.namespace,
            plural="workloads",
        )
        self._core_v1_api.list_namespaced_pod(
            namespace=namespace or self.namespace,
        )
        return [
            KueueWorkload.model_validate(workload)
            for workload in workloads.get("items", [])
        ]

    def ensure_namespace(self, name: str) -> tuple[client.V1Namespace, bool]:
        """Create or look up a namespace by name

        Returns
        -------
        tuple[client.V1Namespace, bool]
            The namespace object and a boolean indicating whether it was created
        """

        try:
            return self._core_v1_api.read_namespace(name), False
        except client.ApiException as e:
            if e.status == 404:
                return self._core_v1_api.create_namespace(
                    client.V1Namespace(metadata=client.V1ObjectMeta(name=name))
                ), True
            raise

    def add_finalizer(
        self,
        resource: client.V1Namespace | client.V1CustomResourceDefinition,
        finalizer: str,
    ) -> None:
        """Add a finalizer to a Kubernetes resource"""
        if resource.metadata.finalizers is None:
            resource.metadata.finalizers = []
        if finalizer not in resource.metadata.finalizers:
            resource.metadata.finalizers.append(finalizer)

            if isinstance(resource, client.V1Namespace):
                self._core_v1_api.replace_namespace(resource.metadata.name, resource)
            else:
                api = client.CustomObjectsApi()
                api.replace_namespaced_custom_object(
                    group=resource.api_version.split("/")[0],
                    version=resource.api_version.split("/")[1],
                    namespace=resource.metadata.namespace,
                    plural=resource.kind.lower() + "s",
                    name=resource.metadata.name,
                    body=resource,
                )