zenml-nightly 0.84.1.dev20250805__py3-none-any.whl → 0.84.2.dev20250807__py3-none-any.whl

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -43,12 +43,16 @@ from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
 from kubernetes.client.rest import ApiException
 
+from zenml.integrations.kubernetes.constants import (
+    STEP_NAME_ANNOTATION_KEY,
+)
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_namespace_manifest,
     build_role_binding_manifest_for_service_account,
     build_secret_manifest,
     build_service_account_manifest,
 )
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 from zenml.utils.time_utils import utc_now
 
@@ -57,6 +61,35 @@ logger = get_logger(__name__)
 R = TypeVar("R")
 
 
+# This is to fix a bug in the kubernetes client which has some wrong
+# client-side validations that means the `on_exit_codes` field is
+# unusable. See https://github.com/kubernetes-client/python/issues/2056
+class PatchedFailurePolicyRule(k8s_client.V1PodFailurePolicyRule):  # type: ignore[misc]
+    """Patched failure policy rule."""
+
+    @property
+    def on_pod_conditions(self):  # type: ignore[no-untyped-def]
+        """On pod conditions.
+
+        Returns:
+            On pod conditions.
+        """
+        return self._on_pod_conditions
+
+    @on_pod_conditions.setter
+    def on_pod_conditions(self, on_pod_conditions):  # type: ignore[no-untyped-def]
+        """On pod conditions.
+
+        Args:
+            on_pod_conditions: On pod conditions.
+        """
+        self._on_pod_conditions = on_pod_conditions
+
+
+k8s_client.V1PodFailurePolicyRule = PatchedFailurePolicyRule
+k8s_client.models.V1PodFailurePolicyRule = PatchedFailurePolicyRule
+
+
 class PodPhase(enum.Enum):
     """Phase of the Kubernetes pod.
 
@@ -71,6 +104,14 @@ class PodPhase(enum.Enum):
     UNKNOWN = "Unknown"
 
 
+class JobStatus(enum.Enum):
+    """Status of a Kubernetes job."""
+
+    RUNNING = "Running"
+    SUCCEEDED = "Succeeded"
+    FAILED = "Failed"
+
+
 def is_inside_kubernetes() -> bool:
     """Check whether we are inside a Kubernetes cluster or on a remote host.
 
@@ -211,7 +252,9 @@ def get_pod(
         The found pod object. None if it's not found.
     """
     try:
-        return core_api.read_namespaced_pod(name=pod_name, namespace=namespace)
+        return retry_on_api_exception(core_api.read_namespaced_pod)(
+            name=pod_name, namespace=namespace
+        )
     except k8s_client.rest.ApiException as e:
         if e.status == 404:
             return None
@@ -603,6 +646,7 @@ def retry_on_api_exception(
     max_retries: int = 3,
     delay: float = 1,
     backoff: float = 1,
+    fail_on_status_codes: Tuple[int, ...] = (404,),
 ) -> Callable[..., R]:
     """Retry a function on API exceptions.
 
@@ -611,6 +655,7 @@ def retry_on_api_exception(
         max_retries: The maximum number of retries.
         delay: The delay between retries.
         backoff: The backoff factor.
+        fail_on_status_codes: The status codes to fail on immediately.
 
     Returns:
         The wrapped function with retry logic.
@@ -624,6 +669,9 @@ def retry_on_api_exception(
             try:
                 return func(*args, **kwargs)
             except ApiException as e:
+                if e.status in fail_on_status_codes:
+                    raise
+
                 retries += 1
                 if retries <= max_retries:
                     logger.warning("Error calling %s: %s.", func.__name__, e)
@@ -657,6 +705,86 @@ def create_job(
     )
 
 
+def get_job(
+    batch_api: k8s_client.BatchV1Api,
+    namespace: str,
+    job_name: str,
+) -> k8s_client.V1Job:
+    """Get a job by name.
+
+    Args:
+        batch_api: Kubernetes batch api.
+        namespace: Kubernetes namespace.
+        job_name: The name of the job to get.
+
+    Returns:
+        The job.
+    """
+    return retry_on_api_exception(batch_api.read_namespaced_job)(
+        name=job_name, namespace=namespace
+    )
+
+
+def list_jobs(
+    batch_api: k8s_client.BatchV1Api,
+    namespace: str,
+    label_selector: Optional[str] = None,
+) -> k8s_client.V1JobList:
+    """List jobs in a namespace.
+
+    Args:
+        batch_api: Kubernetes batch api.
+        namespace: Kubernetes namespace.
+        label_selector: The label selector to use.
+
+    Returns:
+        The job list.
+    """
+    return retry_on_api_exception(batch_api.list_namespaced_job)(
+        namespace=namespace,
+        label_selector=label_selector,
+    )
+
+
+def update_job(
+    batch_api: k8s_client.BatchV1Api,
+    namespace: str,
+    job_name: str,
+    annotations: Dict[str, str],
+) -> k8s_client.V1Job:
+    """Update a job.
+
+    Args:
+        batch_api: Kubernetes batch api.
+        namespace: Kubernetes namespace.
+        job_name: The name of the job to update.
+        annotations: The annotations to update.
+
+    Returns:
+        The updated job.
+    """
+    return retry_on_api_exception(batch_api.patch_namespaced_job)(
+        name=job_name,
+        namespace=namespace,
+        body={"metadata": {"annotations": annotations}},
+    )
+
+
+def is_step_job(job: k8s_client.V1Job) -> bool:
+    """Check if a job is a step job.
+
+    Args:
+        job: The job to check.
+
+    Returns:
+        Whether the job is a step job.
+    """
+    if not job.metadata or not job.metadata.annotations:
+        return False
+
+    return STEP_NAME_ANNOTATION_KEY in job.metadata.annotations
+
+
 def get_container_status(
     pod: k8s_client.V1Pod, container_name: str
 ) -> Optional[k8s_client.V1ContainerState]:
@@ -841,3 +969,242 @@ def wait_for_job_to_finish(
         time.sleep(backoff_interval)
         if exponential_backoff and backoff_interval < maximum_backoff:
             backoff_interval *= 2
+
+
+def check_job_status(
+    batch_api: k8s_client.BatchV1Api,
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    job_name: str,
+    fail_on_container_waiting_reasons: Optional[List[str]] = None,
+    container_name: Optional[str] = None,
+) -> Tuple[JobStatus, Optional[str]]:
+    """Check the status of a job.
+
+    Args:
+        batch_api: Kubernetes BatchV1Api client.
+        core_api: Kubernetes CoreV1Api client.
+        namespace: Kubernetes namespace.
+        job_name: Name of the job for which to wait.
+        fail_on_container_waiting_reasons: List of container waiting reasons
+            that will cause the job to fail.
+        container_name: Name of the container to check for failure.
+
+    Returns:
+        The status of the job and an error message if the job failed.
+    """
+    job: k8s_client.V1Job = retry_on_api_exception(
+        batch_api.read_namespaced_job
+    )(name=job_name, namespace=namespace)
+
+    if job.status.conditions:
+        for condition in job.status.conditions:
+            if condition.type == "Complete" and condition.status == "True":
+                return JobStatus.SUCCEEDED, None
+            if condition.type == "Failed" and condition.status == "True":
+                error_message = condition.message or "Unknown"
+                container_failure_reason = None
+                try:
+                    pods = core_api.list_namespaced_pod(
+                        label_selector=f"job-name={job_name}",
+                        namespace=namespace,
+                    ).items
+                    # Sort pods by creation timestamp, oldest first
+                    pods.sort(
+                        key=lambda pod: pod.metadata.creation_timestamp,
+                    )
+                    if pods:
+                        if (
+                            termination_reason
+                            := get_container_termination_reason(
+                                pods[-1], container_name or "main"
+                            )
+                        ):
+                            exit_code, reason = termination_reason
+                            if exit_code != 0:
+                                container_failure_reason = (
+                                    f"{reason}, exit_code={exit_code}"
+                                )
+                except Exception:
+                    pass
+
+                if container_failure_reason:
+                    error_message += f" (container failure reason: {container_failure_reason})"
+
+                return JobStatus.FAILED, error_message
+
+    if fail_on_container_waiting_reasons:
+        pod_list: k8s_client.V1PodList = retry_on_api_exception(
+            core_api.list_namespaced_pod
+        )(
+            namespace=namespace,
+            label_selector=f"job-name={job_name}",
+            field_selector="status.phase=Pending",
+        )
+        for pod in pod_list.items:
+            container_state = get_container_status(
+                pod, container_name or "main"
+            )
+
+            if (
+                container_state
+                and (waiting_state := container_state.waiting)
+                and waiting_state.reason in fail_on_container_waiting_reasons
+            ):
+                retry_on_api_exception(batch_api.delete_namespaced_job)(
+                    name=job_name,
+                    namespace=namespace,
+                    propagation_policy="Foreground",
+                )
+                error_message = (
+                    f"Detected container in state `{waiting_state.reason}`"
+                )
+                return JobStatus.FAILED, error_message
+
+    return JobStatus.RUNNING, None
+
+
+def create_config_map(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    name: str,
+    data: Dict[str, str],
+) -> None:
+    """Create a Kubernetes config map.
+
+    Args:
+        core_api: Kubernetes CoreV1Api client.
+        namespace: Kubernetes namespace.
+        name: Name of the config map to create.
+        data: Data to store in the config map.
+    """
+    retry_on_api_exception(core_api.create_namespaced_config_map)(
+        namespace=namespace,
+        body=k8s_client.V1ConfigMap(metadata={"name": name}, data=data),
+    )
+
+
+def update_config_map(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    name: str,
+    data: Dict[str, str],
+) -> None:
+    """Update a Kubernetes config map.
+
+    Args:
+        core_api: Kubernetes CoreV1Api client.
+        namespace: Kubernetes namespace.
+        name: Name of the config map to update.
+        data: Data to store in the config map.
+    """
+    retry_on_api_exception(core_api.patch_namespaced_config_map)(
+        namespace=namespace,
+        name=name,
+        body=k8s_client.V1ConfigMap(data=data),
+    )
+
+
+def get_config_map(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    name: str,
+) -> k8s_client.V1ConfigMap:
+    """Get a Kubernetes config map.
+
+    Args:
+        core_api: Kubernetes CoreV1Api client.
+        namespace: Kubernetes namespace.
+        name: Name of the config map to get.
+
+    Returns:
+        The config map.
+    """
+    return retry_on_api_exception(core_api.read_namespaced_config_map)(
+        namespace=namespace,
+        name=name,
+    )
+
+
+def delete_config_map(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    name: str,
+) -> None:
+    """Delete a Kubernetes config map.
+
+    Args:
+        core_api: Kubernetes CoreV1Api client.
+        namespace: Kubernetes namespace.
+        name: Name of the config map to delete.
+    """
+    retry_on_api_exception(core_api.delete_namespaced_config_map)(
+        namespace=namespace,
+        name=name,
+    )
+
+
+def get_parent_job_name(
+    core_api: k8s_client.CoreV1Api,
+    pod_name: str,
+    namespace: str,
+) -> Optional[str]:
+    """Get the name of the job that created a pod.
+
+    Args:
+        core_api: Kubernetes CoreV1Api client.
+        pod_name: Name of the pod.
+        namespace: Kubernetes namespace.
+
+    Returns:
+        The name of the job that created the pod, or None if the pod is not
+        associated with a job.
+    """
+    pod = get_pod(core_api, pod_name=pod_name, namespace=namespace)
+    if (
+        pod
+        and pod.metadata
+        and pod.metadata.labels
+        and (job_name := pod.metadata.labels.get("job-name", None))
+    ):
+        return cast(str, job_name)
+
+    return None
+
+
+def apply_default_resource_requests(
+    memory: str,
+    cpu: Optional[str] = None,
+    pod_settings: Optional[KubernetesPodSettings] = None,
+) -> KubernetesPodSettings:
+    """Applies default resource requests to a pod settings object.
+
+    Args:
+        memory: The memory resource request.
+        cpu: The CPU resource request.
+        pod_settings: The pod settings to update. A new one will be created
+            if not provided.
+
+    Returns:
+        The new or updated pod settings.
+    """
+    resources = {
+        "requests": {"memory": memory},
+    }
+    if cpu:
+        resources["requests"]["cpu"] = cpu
+    if not pod_settings:
+        pod_settings = KubernetesPodSettings(resources=resources)
+    elif not pod_settings.resources:
+        # We can't update the pod settings in place (because it's a frozen
+        # pydantic model), so we have to create a new one.
+        pod_settings = KubernetesPodSettings(
+            **pod_settings.model_dump(exclude_unset=True),
+            resources=resources,
+        )
+    else:
+        set_requests = pod_settings.resources.get("requests", {})
+        resources["requests"].update(set_requests)
+        pod_settings.resources["requests"] = resources["requests"]
+
+    return pod_settings