From f0a7bb3996925278973d0bdf48b4ac47815e518f Mon Sep 17 00:00:00 2001 From: raphaelauv Date: Thu, 29 Oct 2026 14:28:59 +0100 Subject: [PATCH] Wait for Kueue Jobs through admission and preemption KubernetesStartKueueJobOperator inherited the KubernetesJobOperator wait, which looks for the Job's pods right after creating it. Kueue keeps the Job suspended without pods until it admits it, and may preempt it (suspend it again and delete its pods) then admit it again, so the task failed to find pods before admission, and after a preemption waited on, or read logs and XCom from, pods that had been deleted. Least-privilege Roles that grant get on jobs but not on jobs/status also got a 403 when job operators and triggers waited for a Job: Kubernetes RBAC checks subresources separately, although the Job object carries the same status. closes: #74373 --- docs/spelling_wordlist.txt | 2 + providers/cncf/kubernetes/docs/changelog.rst | 6 + .../cncf/kubernetes/docs/kubernetes_rbac.rst | 17 +- providers/cncf/kubernetes/docs/operators.rst | 86 ++++++ providers/cncf/kubernetes/provider.yaml | 1 + .../cncf/kubernetes/get_provider_info.py | 1 + .../cncf/kubernetes/hooks/kubernetes.py | 11 +- .../cncf/kubernetes/operators/kueue.py | 101 ++++++++ .../cncf/kubernetes/triggers/kueue.py | 245 ++++++++++++++++++ .../cncf/kubernetes/hooks/test_kubernetes.py | 8 +- .../cncf/kubernetes/operators/test_kueue.py | 129 ++++++++- .../cncf/kubernetes/triggers/test_kueue.py | 203 +++++++++++++++ 12 files changed, 793 insertions(+), 17 deletions(-) create mode 100644 providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/kueue.py create mode 100644 providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_kueue.py diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index ec0b8690b3986..34ec67a4f034e 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1748,6 +1748,8 @@ subprocess subprocesses subqueries subquery +subresource +subresources SubscriberClient subscriptionId substring diff --git a/providers/cncf/kubernetes/docs/changelog.rst b/providers/cncf/kubernetes/docs/changelog.rst index 829cc47ba9e4d..a464fa489d716 100644 --- a/providers/cncf/kubernetes/docs/changelog.rst +++ b/providers/cncf/kubernetes/docs/changelog.rst @@ -33,6 +33,12 @@ Changelog team's ``[kubernetes_executor]`` config. Like all team config, these don't fall back to the un-prefixed section. If your teams connected using the un-prefixed values, set them for each team. +.. note:: + Job operators and triggers now read the Job object instead of the ``jobs/status`` subresource, so they + need ``get`` on ``jobs`` (``batch`` API group) rather than on ``jobs/status``. If the ``ServiceAccount`` + of your triggerer or workers was granted ``get`` on ``jobs/status`` but not on ``jobs``, grant it + ``get`` on ``jobs``. + 10.25.0 ....... diff --git a/providers/cncf/kubernetes/docs/kubernetes_rbac.rst b/providers/cncf/kubernetes/docs/kubernetes_rbac.rst index 971c7c0677016..4d81f86225f27 100644 --- a/providers/cncf/kubernetes/docs/kubernetes_rbac.rst +++ b/providers/cncf/kubernetes/docs/kubernetes_rbac.rst @@ -128,23 +128,24 @@ permissions: - apiGroups: ["batch"] resources: ["jobs"] verbs: ["create", "get", "list", "watch", "patch", "delete"] - - apiGroups: ["batch"] - resources: ["jobs/status"] - verbs: ["get", "watch"] + +The Job status is read from the Job object itself, so no permission on the +``jobs/status`` subresource is needed. When a job operator waits for completion, streams pod logs, reads XCom, or cleans up discovered pods, it also needs the relevant pod permissions from the pod launch section. For deferrable job operators, the worker creates the Job before deferring and -the triggerer polls Job status. If the triggerer uses in-cluster credentials, -bind the triggerer's ``ServiceAccount`` to the Job status permissions. When +the triggerer polls the Job. If the triggerer uses in-cluster credentials, +grant the triggerer's ``ServiceAccount`` permission to get jobs. When XCom is enabled, the triggerer also needs permission to get pods and exec into -the XCom sidecar container. The worker reads pod logs after the task resumes. +the XCom sidecar container, and for ``KubernetesStartKueueJobOperator`` to list +pods. The worker reads pod logs after the task resumes. ``KubernetesPatchJobOperator`` only needs permission to patch jobs. -``KubernetesDeleteJobOperator`` reads Job status before deleting a Job, so it -needs permission to get ``jobs/status`` and delete jobs. +``KubernetesDeleteJobOperator`` reads the Job before deleting it, so it +needs permission to get and delete jobs. Cleanup permissions ------------------- diff --git a/providers/cncf/kubernetes/docs/operators.rst b/providers/cncf/kubernetes/docs/operators.rst index 9c491a275ca54..aaa2e9c8a4e53 100644 --- a/providers/cncf/kubernetes/docs/operators.rst +++ b/providers/cncf/kubernetes/docs/operators.rst @@ -861,6 +861,23 @@ When the task is killed, ``on_kill`` deletes the Job (with foreground cascade). For discovered pods, deletion is controlled by ``on_kill_action``: ``delete_pod`` attempts direct pod deletion and ``keep_pod`` skips it. +XCom of a Job +^^^^^^^^^^^^^ + +With ``wait_until_job_complete=True`` and ``do_xcom_push=True``, each pod of the Job writes its XCom to +``/airflow/xcom/return.json``, as with ``KubernetesPodOperator``. The operator returns one value per pod, +so the ``return_value`` XCom is a list. When the list has a single value, it is unwrapped and +``return_value`` is that value, unless ``unwrap_single=False``. A pod that writes no XCom gives ``None``. + +The operator also pushes these XCom, whatever ``do_xcom_push`` is: + +* ``job_name`` and ``job_namespace``: the name and namespace of the Job, pushed when the Job is created. +* ``job``: the Job as a dictionary, including its ``status`` (the finished Job when + ``wait_until_job_complete=True``). + +``KubernetesJobOperator`` reads the XCom of the pods it found when the Job started. +For ``KubernetesStartKueueJobOperator``, see :ref:`howto/operator:KubernetesStartKueueJobOperator`. + .. _howto/operator:KubernetesDeleteJobOperator: @@ -930,6 +947,75 @@ you to start a Kueue job in a Kubernetes cluster :start-after: [START howto_operator_k8s_install_kueue] :end-before: [END howto_operator_k8s_install_kueue] +Waiting for the Job +^^^^^^^^^^^^^^^^^^^ + +With ``wait_until_job_complete=True``, the operator waits on the Job itself rather than on its pods: +the Job stays suspended without pods until Kueue admits it, and Kueue may preempt it (suspend it again +and delete its running pods) then admit it again, any number of times. New pods, with new names, are +created at each admission. Pods that already succeeded are not deleted by a preemption and still count +for the Job. Once the Job is finished, logs are read from the pods the Job has at that time. + +XCom +^^^^ + +With ``do_xcom_push=True``, the XCom of each pod is read at the first poll after its base container ends, +while the Job runs: until then, the XCom sidecar keeps the pod, and so the Job, from finishing. Once the Job is +finished, only the XCom of the pods that are ``Succeeded`` at that time are kept: the XCom of pods +deleted by a preemption and of failed pods are dropped. ``return_value`` is the list of those XCom, +one per succeeded pod, ordered by pod name (see :ref:`howto/operator:kubernetesjoboperator` for the +other XCom pushed and ``unwrap_single``). + +If the XCom of a succeeded pod could not be read, for example because the pod was deleted before the +task finished, the task fails rather than returning an incomplete list. A pod that writes no XCom does +not fail the task: its value is ``None``. + +Number of pods: ``completions`` and ``parallelism`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The number of values in ``return_value`` is the number of pods that succeeded, set by the ``completions`` +and ``parallelism`` parameters of the Job: + +.. list-table:: + :header-rows: 1 + + * - Settings + - Pods running at once + - The Job is finished when + - ``return_value`` + * - Defaults: ``parallelism=1``, ``completions`` not set + - 1 + - one pod succeeded + - the value of that pod + * - ``completions=N`` + - up to ``parallelism`` + - ``N`` pods succeeded + - a list of ``N`` values + * - ``parallelism=M`` (``M`` > 1), ``completions`` not set (work queue) + - up to ``M`` + - one pod succeeded and all the others ended + - a list of one value per pod that succeeded, from 1 to ``M`` + +Failed pods are retried up to ``backoff_limit`` and never add a value. ``completions`` can also come +from ``full_job_spec`` or ``job_template_file`` when the parameter is not set. + +Pod names end with a random suffix, so the order of ``return_value`` does not follow the order in which +pods ran. With ``completion_mode="Indexed"``, each pod gets its index, from ``0`` to ``completions - 1``, +in the ``JOB_COMPLETION_INDEX`` environment variable, and pod names start with ``--``. +As names are sorted as text (``10`` comes before ``2``), write the index into the XCom and sort the list +on it in the downstream task when the order matters: + +.. code-block:: bash + + echo "{\"index\": $JOB_COMPLETION_INDEX, \"rows\": 42}" > /airflow/xcom/return.json + +Kueue admits and preempts the Job as a whole: it reserves the quota for the pods the Job runs at once +(``parallelism``, at most ``completions``), so a high ``parallelism`` may make the Job wait longer for +admission. + +Reference +^^^^^^^^^ + For further information, look at: * `Kubernetes Documentation `__ diff --git a/providers/cncf/kubernetes/provider.yaml b/providers/cncf/kubernetes/provider.yaml index 3d8f4aa31917d..031473fa61a76 100644 --- a/providers/cncf/kubernetes/provider.yaml +++ b/providers/cncf/kubernetes/provider.yaml @@ -176,6 +176,7 @@ triggers: python-modules: - airflow.providers.cncf.kubernetes.triggers.pod - airflow.providers.cncf.kubernetes.triggers.job + - airflow.providers.cncf.kubernetes.triggers.kueue secrets-backends: diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py index 47e5e5cf5367b..7d92864d62c9a 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py @@ -73,6 +73,7 @@ def get_provider_info(): "python-modules": [ "airflow.providers.cncf.kubernetes.triggers.pod", "airflow.providers.cncf.kubernetes.triggers.job", + "airflow.providers.cncf.kubernetes.triggers.kueue", ], } ], diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py index b4809db6cf367..75f4019e889d1 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py @@ -633,13 +633,14 @@ def get_job_status(self, job_name: str, namespace: str) -> V1Job: """ Get job with status of specified name and namespace. + Reads the Job object rather than the ``jobs/status`` subresource: both return the same ``status``, + but RBAC checks subresources separately, so this only needs ``get`` on ``jobs``. + :param job_name: Name of Job to fetch. :param namespace: Namespace of the Job. :return: Job object """ - return self.batch_v1_client.read_namespaced_job_status( - name=job_name, namespace=namespace, pretty=True - ) + return self.batch_v1_client.read_namespaced_job(name=job_name, namespace=namespace, pretty=True) def wait_until_job_complete(self, job_name: str, namespace: str, job_poll_interval: float = 10) -> V1Job: """ @@ -1325,12 +1326,14 @@ async def get_job_status(self, name: str, namespace: str) -> V1Job: """ Get job's status object. + Reads the Job object, see :meth:`KubernetesHook.get_job_status`. + :param name: Name of the pod. :param namespace: Name of the pod's namespace. """ async with self.get_conn() as connection: v1_api = async_client.BatchV1Api(connection) - job: V1Job = await v1_api.read_namespaced_job_status( + job: V1Job = await v1_api.read_namespaced_job( name=name, namespace=namespace, ) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py index 4b12616375720..27baed3755f98 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py @@ -19,13 +19,26 @@ from __future__ import annotations import json +import time from collections.abc import Sequence from functools import cached_property +from typing import TYPE_CHECKING, Any from kubernetes.utils import FailToCreateError from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook from airflow.providers.cncf.kubernetes.operators.job import KubernetesJobOperator +from airflow.providers.cncf.kubernetes.triggers.kueue import ( + JOB_NAME_LABEL, + JOB_STATE_MESSAGES, + KubernetesKueueJobTrigger, + KueueJobState, + build_job_event, + collect_xcom_results, + get_job_state, + get_succeeded_pod_xcoms, +) +from airflow.providers.cncf.kubernetes.utils.pod_manager import XComRetrievalError from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_1_PLUS from airflow.providers.common.compat.sdk import AirflowException @@ -34,6 +47,9 @@ else: from airflow.models import BaseOperator +if TYPE_CHECKING: + from airflow.sdk import Context + class KubernetesInstallKueueOperator(BaseOperator): """ @@ -105,6 +121,12 @@ class KubernetesStartKueueJobOperator(KubernetesJobOperator): """ Executes a Kubernetes Job in Kueue. + With ``wait_until_job_complete=True``, the operator polls the Job itself until it is finished, instead + of waiting on the pods found when the Job started: Kueue keeps the Job suspended without pods until it + admits it, and may preempt it (suspending it again and deleting its pods) then admit it again any + number of times. The XCom of each pod is read while the Job runs, and kept only for the pods that are + ``Succeeded`` once the Job is finished. Logs are read from the pods the Job has once it is finished. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:KubernetesStartKueueJobOperator` @@ -133,3 +155,82 @@ def __init__(self, queue_name: str, *args, **kwargs) -> None: self.suspend = True self.labels.update({"kueue.x-k8s.io/queue-name": self.queue_name}) self.annotations.update({"kueue.x-k8s.io/queue-name": self.queue_name}) + + def execute(self, context: Context): + if not self.wait_until_job_complete: + return super().execute(context) + + self.name = self._set_name(self.name) + self.job_request_obj = self.build_job_request_obj(context) + self.job = self.create_job(job_request_obj=self.job_request_obj) # must set `self.job` for `on_kill` + + ti = context["ti"] + ti.xcom_push(key="job_name", value=self.job.metadata.name) + ti.xcom_push(key="job_namespace", value=self.job.metadata.namespace) + + if self.deferrable: + self.defer( + trigger=KubernetesKueueJobTrigger( + job_name=self.job.metadata.name, + job_namespace=self.job.metadata.namespace, + base_container_name=self.base_container_name, + kubernetes_conn_id=self.kubernetes_conn_id, + cluster_context=self.cluster_context, + config_file=self.config_file, + in_cluster=self.in_cluster, + poll_interval=self.job_poll_interval, + do_xcom_push=self.do_xcom_push, + ), + method_name="execute_complete", + ) + + return self.execute_complete( + context, self._wait_until_job_finished(self.job.metadata.name, self.job.metadata.namespace) + ) + + def _wait_until_job_finished(self, job_name: str, namespace: str) -> dict[str, Any]: + """Poll the Job until it is finished, and return the event ``KubernetesKueueJobTrigger`` yields.""" + state: KueueJobState | None = None + xcom_results: dict[str, str] = {} + while True: + job = self.hook.get_job_status(job_name=job_name, namespace=namespace) + if (new_state := get_job_state(job)) != state: + state = new_state + self.log.info("The job '%s' %s", job_name, JOB_STATE_MESSAGES[state]) + if state == KueueJobState.FINISHED: + return build_job_event(job, xcom_results) + if state == KueueJobState.RUNNING and self.do_xcom_push: + collect_xcom_results( + self.client, self.pod_manager, job_name, namespace, self.base_container_name, xcom_results + ) + time.sleep(self.job_poll_interval) + + def execute_complete(self, context: Context, event: dict, **kwargs): + # The pods found before a preemption by Kueue were deleted: read the ones the Job has now + pods = sorted( + self.client.list_namespaced_pod( + namespace=event["namespace"], label_selector=f"{JOB_NAME_LABEL}={event['name']}" + ).items, + key=lambda pod: pod.metadata.name, + ) + for pod in pods: + self.log_matching_pod(pod=pod, context=context) + xcom_result = get_succeeded_pod_xcoms(pods, event["xcom_result"]) + result = super().execute_complete( + context, + { + **event, + "pod_names": [pod.metadata.name for pod in pods], + "pod_namespace": event["namespace"], + "xcom_result": xcom_result, + }, + **kwargs, + ) + # Only reached for a successful Job: super() raises for a failed one + succeeded = (event["job"].get("status") or {}).get("succeeded") or 0 + if self.do_xcom_push and len(xcom_result) < succeeded: + raise XComRetrievalError( + f"Read the XCom of {len(xcom_result)} of the {succeeded} pods that succeeded in the job " + f"'{event['name']}': the other pods were deleted, or their XCom could not be read" + ) + return result diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/kueue.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/kueue.py new file mode 100644 index 0000000000000..2d891206ccf0a --- /dev/null +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/kueue.py @@ -0,0 +1,245 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Wait for a Kubernetes Job run through Kueue.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Iterable +from enum import Enum +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from kubernetes.client.rest import ApiException + +from airflow.providers.cncf.kubernetes.hooks.kubernetes import ( + JOB_FINAL_STATUS_CONDITION_TYPES, + AsyncKubernetesHook, + KubernetesHook, +) +from airflow.providers.cncf.kubernetes.utils.container import container_is_running, container_is_terminated +from airflow.providers.cncf.kubernetes.utils.pod_manager import PodManager, XComRetrievalError +from airflow.providers.cncf.kubernetes.utils.xcom_sidecar import PodDefaults +from airflow.triggers.base import BaseTrigger, TriggerEvent + +if TYPE_CHECKING: + from kubernetes.client import CoreV1Api, models as k8s + +# Set by the Job controller on the pods it creates for a Job +JOB_NAME_LABEL = "batch.kubernetes.io/job-name" + + +class KueueJobState(str, Enum): + """State of a Job run through Kueue.""" + + SUSPENDED = "suspended" + RUNNING = "running" + FINISHED = "finished" + + +JOB_STATE_MESSAGES = { + KueueJobState.SUSPENDED: "is suspended: waiting for admission by Kueue, or preempted by Kueue", + KueueJobState.RUNNING: "is admitted by Kueue: its pods are running", + KueueJobState.FINISHED: "is finished", +} + + +def get_job_state(job: k8s.V1Job) -> KueueJobState: + """Return the state of a Job run through Kueue; finished means Complete or Failed.""" + conditions = (job.status.conditions if job.status else None) or [] + if any(c.type in JOB_FINAL_STATUS_CONDITION_TYPES and c.status == "True" for c in conditions): + return KueueJobState.FINISHED + # The Job is created suspended; Kueue resumes it on admission and suspends it again on preemption + if job.spec and job.spec.suspend: + return KueueJobState.SUSPENDED + return KueueJobState.RUNNING + + +def build_job_event(job: k8s.V1Job, xcom_results: dict[str, str]) -> dict[str, Any]: + """ + Build the event of a finished Job, as passed to ``KubernetesStartKueueJobOperator.execute_complete``. + + :param job: the finished Job. + :param xcom_results: the raw XCom read from each pod, by pod name. + """ + error_message = KubernetesHook.is_job_failed(job=job) + return { + "name": job.metadata.name, + "namespace": job.metadata.namespace, + "status": "error" if error_message else "success", + "message": f"Job failed with error: {error_message}" + if error_message + else "Job completed successfully", + "job": job.to_dict(), + "xcom_result": xcom_results, + } + + +def get_succeeded_pod_xcoms(pods: Iterable[k8s.V1Pod], xcom_results: dict[str, str]) -> list[str]: + """Return the XCom read from the pods that succeeded, skipping failed pods and preempted runs.""" + return [ + xcom_results[pod.metadata.name] + for pod in pods + if pod.metadata.name in xcom_results and pod.status and pod.status.phase == "Succeeded" + ] + + +def collect_xcom_results( + client: CoreV1Api, + pod_manager: PodManager, + job_name: str, + namespace: str, + base_container_name: str, + xcom_results: dict[str, str], +) -> None: + """ + Read into ``xcom_results`` the XCom of each pod of the Job whose base container has ended. + + Reading the XCom also stops the XCom sidecar: until then the pod, and so the Job, cannot finish. + Pods already read are skipped, so this is called at each poll. + """ + pods = client.list_namespaced_pod( + namespace=namespace, label_selector=f"{JOB_NAME_LABEL}={job_name}" + ).items + for pod in pods: + if ( + pod.metadata.name in xcom_results + or not container_is_terminated(pod, base_container_name) + or not container_is_running(pod, PodDefaults.SIDECAR_CONTAINER_NAME) + ): + continue + try: + xcom_results[pod.metadata.name] = pod_manager.extract_xcom(pod, ignore_kill_failure=False) + except XComRetrievalError as e: + pod_manager.log.warning("Cannot read the XCom of the pod %s: %s", pod.metadata.name, e) + except ApiException as e: + if e.status != 404: + raise + pod_manager.log.warning("The pod %s was deleted before its XCom was read", pod.metadata.name) + + +class KubernetesKueueJobTrigger(BaseTrigger): + """ + Wait until a Job run through Kueue is finished, whatever Kueue does with it meanwhile. + + Unlike :class:`~airflow.providers.cncf.kubernetes.triggers.job.KubernetesJobTrigger`, it does not wait + on pods found when the Job started: Kueue may preempt the Job (its pods are deleted) and admit it again + (new pods start) any number of times. Only the Job itself is polled, and the XCom are read from the + pods the Job has at each poll. + + :param job_name: The name of the job. + :param job_namespace: The namespace of the job. + :param base_container_name: The name of the base container in the pod. + :param kubernetes_conn_id: The :ref:`kubernetes connection id ` + for the Kubernetes cluster. + :param cluster_context: Context that points to kubernetes cluster. + :param config_file: Path to kubeconfig file. + :param poll_interval: Polling period in seconds to check for the status. + :param in_cluster: run kubernetes client with in_cluster configuration. + :param do_xcom_push: If True, the content of the file + /airflow/xcom/return.json in the container will also be pushed to an + XCom when the container completes. + """ + + def __init__( + self, + job_name: str, + job_namespace: str, + base_container_name: str, + kubernetes_conn_id: str | None = None, + poll_interval: float = 10.0, + cluster_context: str | None = None, + config_file: str | None = None, + in_cluster: bool | None = None, + do_xcom_push: bool = False, + ): + super().__init__() + self.job_name = job_name + self.job_namespace = job_namespace + self.base_container_name = base_container_name + self.kubernetes_conn_id = kubernetes_conn_id + self.poll_interval = poll_interval + self.cluster_context = cluster_context + self.config_file = config_file + self.in_cluster = in_cluster + self.do_xcom_push = do_xcom_push + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize KubernetesKueueJobTrigger arguments and classpath.""" + return ( + "airflow.providers.cncf.kubernetes.triggers.kueue.KubernetesKueueJobTrigger", + { + "job_name": self.job_name, + "job_namespace": self.job_namespace, + "base_container_name": self.base_container_name, + "kubernetes_conn_id": self.kubernetes_conn_id, + "poll_interval": self.poll_interval, + "cluster_context": self.cluster_context, + "config_file": self.config_file, + "in_cluster": self.in_cluster, + "do_xcom_push": self.do_xcom_push, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Poll the Job until it is finished and yield a TriggerEvent.""" + state: KueueJobState | None = None + xcom_results: dict[str, str] = {} + while True: + job = await self.hook.get_job_status(name=self.job_name, namespace=self.job_namespace) + if (new_state := get_job_state(job)) != state: + state = new_state + self.log.info("The job '%s' %s", self.job_name, JOB_STATE_MESSAGES[state]) + if state == KueueJobState.FINISHED: + yield TriggerEvent(build_job_event(job, xcom_results)) + return + if state == KueueJobState.RUNNING and self.do_xcom_push: + await asyncio.to_thread(self._collect_xcom_results, xcom_results) + await asyncio.sleep(self.poll_interval) + + def _collect_xcom_results(self, xcom_results: dict[str, str]) -> None: + # Run in a thread: the sync client loads its connection and kube config on first use + collect_xcom_results( + self.sync_hook.core_v1_client, + self.pod_manager, + self.job_name, + self.job_namespace, + self.base_container_name, + xcom_results, + ) + + @cached_property + def hook(self) -> AsyncKubernetesHook: + return AsyncKubernetesHook( + conn_id=self.kubernetes_conn_id, + in_cluster=self.in_cluster, + config_file=self.config_file, + cluster_context=self.cluster_context, + ) + + @cached_property + def sync_hook(self) -> KubernetesHook: + return KubernetesHook( + conn_id=self.kubernetes_conn_id, + in_cluster=self.in_cluster, + config_file=self.config_file, + cluster_context=self.cluster_context, + ) + + @cached_property + def pod_manager(self) -> PodManager: + return PodManager(kube_client=self.sync_hook.core_v1_client) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py index 4fed8799e5728..41e5ca73c3a06 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py @@ -698,12 +698,12 @@ def test_delete_custom_object( @patch("kubernetes.config.kube_config.KubeConfigMerger") @patch(f"{HOOK_MODULE}.KubernetesHook.batch_v1_client") def test_get_job_status(self, mock_client, mock_kube_config_merger, mock_kube_config_loader): - job_expected = mock_client.read_namespaced_job_status.return_value + job_expected = mock_client.read_namespaced_job.return_value hook = KubernetesHook() job_actual = hook.get_job_status(job_name=JOB_NAME, namespace=NAMESPACE) - mock_client.read_namespaced_job_status.assert_called_once_with( + mock_client.read_namespaced_job.assert_called_once_with( name=JOB_NAME, namespace=NAMESPACE, pretty=True ) assert job_actual == job_expected @@ -2027,7 +2027,7 @@ async def test_read_logs_decodes_off_the_event_loop(self, lib_method, mock_to_th mock_to_thread.assert_awaited_once_with(_split_log_bytes, raw_bytes) @pytest.mark.asyncio - @mock.patch(KUBE_BATCH_API.format("read_namespaced_job_status")) + @mock.patch(KUBE_BATCH_API.format("read_namespaced_job")) async def test_get_job_status(self, lib_method, kube_config_loader): lib_method.return_value = self.mock_await_result(None) @@ -2042,7 +2042,7 @@ async def test_get_job_status(self, lib_method, kube_config_loader): namespace=NAMESPACE, ) - lib_method.assert_called_once() + lib_method.assert_called_once_with(name=JOB_NAME, namespace=NAMESPACE) @pytest.mark.asyncio @mock.patch(HOOK_MODULE + ".asyncio.sleep") diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_kueue.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_kueue.py index 907d0e44023a8..5506f68f60f3f 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_kueue.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_kueue.py @@ -21,6 +21,7 @@ from unittest import mock import pytest +from kubernetes.client import models as k8s from kubernetes.utils import FailToCreateError from airflow.models.dag import DAG @@ -29,7 +30,11 @@ KubernetesInstallKueueOperator, KubernetesStartKueueJobOperator, ) -from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.cncf.kubernetes.triggers.kueue import KubernetesKueueJobTrigger, build_job_event +from airflow.providers.cncf.kubernetes.utils.pod_manager import XComRetrievalError +from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred + +from unit.cncf.kubernetes.triggers.test_kueue import JOB_NAME, NAMESPACE, build_job, build_pod TEST_TASK_ID = "test_task" TEST_K8S_CONN_ID = "test_kubernetes_conn_id" @@ -193,3 +198,125 @@ def test_init_suspend_is_none(self, mock_log): assert operator.labels == {"kueue.x-k8s.io/queue-name": TEST_QUEUE_NAME} assert operator.annotations == {"kueue.x-k8s.io/queue-name": TEST_QUEUE_NAME} mock_log.info.assert_called_once_with(expected_info_message, "KubernetesStartKueueJobOperator") + + +class TestKubernetesStartKueueJobOperatorWait: + @pytest.fixture(autouse=True) + def setup_mocks(self): + operator_class = KubernetesStartKueueJobOperator + with ( + mock.patch.object(operator_class, "hook") as self.mock_hook, + mock.patch.object(operator_class, "client") as self.mock_client, + mock.patch.object(operator_class, "pod_manager") as self.mock_pod_manager, + mock.patch.object(operator_class, "_write_logs") as self.mock_write_logs, + mock.patch.object(operator_class, "post_complete_action"), + mock.patch.object(operator_class, "build_job_request_obj", return_value=build_job(True)), + mock.patch.object( + operator_class, "create_job", side_effect=lambda job_request_obj: job_request_obj + ), + ): + self.mock_hook.get_pod.side_effect = lambda name, namespace: self.pods[name] + yield + + def set_pods(self, *pod_lists: list[k8s.V1Pod]) -> None: + """Make each list_namespaced_pod call return the next list of pods.""" + self.pods = {pod.metadata.name: pod for pods in pod_lists for pod in pods} + self.mock_client.list_namespaced_pod.side_effect = [k8s.V1PodList(items=pods) for pods in pod_lists] + + def build_operator(self, **kwargs) -> KubernetesStartKueueJobOperator: + return KubernetesStartKueueJobOperator( + task_id=TEST_TASK_ID, + queue_name=TEST_QUEUE_NAME, + namespace=NAMESPACE, + wait_until_job_complete=True, + job_poll_interval=0, + **kwargs, + ) + + @mock.patch.object(KubernetesJobOperator, "execute", autospec=True) + def test_execute_without_wait_runs_job_operator(self, mock_execute): + operator = self.build_operator() + operator.wait_until_job_complete = False + + assert operator.execute({}) == mock_execute.return_value + self.mock_hook.get_job_status.assert_not_called() + + def test_execute_waits_through_preemption_and_reads_pods_of_last_run(self): + """Kueue admits the Job, preempts it (its first pod is deleted), admits it again, then it completes.""" + last_run = build_pod("last-run", base="terminated", sidecar="running", phase="Succeeded") + self.set_pods([build_pod("first-run", base="running", sidecar="running")], [last_run], [last_run]) + self.mock_hook.get_job_status.side_effect = [ + build_job(True), + build_job(False), + build_job(True), + build_job(False), + build_job(False, "Complete"), + ] + self.mock_pod_manager.extract_xcom.return_value = '{"a": 1}' + + result = self.build_operator(do_xcom_push=True, get_logs=True).execute({"ti": mock.MagicMock()}) + + assert result == {"a": 1} + self.mock_hook.get_job_status.assert_called_with(job_name=JOB_NAME, namespace=NAMESPACE) + self.mock_pod_manager.extract_xcom.assert_called_once_with(last_run, ignore_kill_failure=False) + self.mock_write_logs.assert_called_once_with(last_run) + + def test_execute_raises_when_job_failed(self): + self.set_pods([]) + self.mock_hook.get_job_status.side_effect = [build_job(True), build_job(False, "Failed")] + + with pytest.raises(AirflowException, match="Job failed with error: BackoffLimitExceeded"): + self.build_operator().execute({"ti": mock.MagicMock()}) + + def test_execute_deferrable_defers_to_kueue_job_trigger(self): + with pytest.raises(TaskDeferred) as deferred: + self.build_operator(deferrable=True).execute({"ti": mock.MagicMock()}) + + assert isinstance(deferred.value.trigger, KubernetesKueueJobTrigger) + assert deferred.value.method_name == "execute_complete" + self.mock_client.list_namespaced_pod.assert_not_called() + + def test_execute_complete_returns_xcom_of_succeeded_pods_of_the_job(self): + self.set_pods( + [ + build_pod("succeeded", base="terminated", phase="Succeeded"), + build_pod("failed", base="terminated", phase="Failed"), + ] + ) + event = build_job_event( + build_job(False, "Complete", succeeded=1), + {"succeeded": '{"a": 1}', "failed": '{"b": 2}', "preempted": "3"}, + ) + + result = self.build_operator(do_xcom_push=True, unwrap_single=False).execute_complete( + {"ti": mock.MagicMock()}, event + ) + + assert result == [{"a": 1}] + + def test_execute_complete_orders_xcom_by_pod_name(self): + self.set_pods( + [build_pod(name, base="terminated", phase="Succeeded") for name in ("job-1-b", "job-0-a")] + ) + event = build_job_event( + build_job(False, "Complete", succeeded=2), {"job-1-b": '{"index": 1}', "job-0-a": '{"index": 0}'} + ) + + result = self.build_operator(do_xcom_push=True).execute_complete({"ti": mock.MagicMock()}, event) + + assert result == [{"index": 0}, {"index": 1}] + + @pytest.mark.parametrize( + ("pod_names", "xcom_results"), + [ + pytest.param(["read"], {"read": '{"a": 1}', "deleted": '{"b": 2}'}, id="succeeded-pod-deleted"), + pytest.param(["read", "not-read"], {"read": '{"a": 1}'}, id="xcom-of-succeeded-pod-not-read"), + ], + ) + def test_execute_complete_raises_when_xcom_of_a_succeeded_pod_is_missing(self, pod_names, xcom_results): + self.set_pods([build_pod(name, base="terminated", phase="Succeeded") for name in pod_names]) + event = build_job_event(build_job(False, "Complete", succeeded=2), xcom_results) + + with pytest.raises(XComRetrievalError, match="Read the XCom of 1 of the 2 pods that succeeded"): + self.build_operator(do_xcom_push=True).execute_complete({"ti": mock.MagicMock()}, event) + self.mock_write_logs.assert_called() diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_kueue.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_kueue.py new file mode 100644 index 0000000000000..4aa757375e446 --- /dev/null +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_kueue.py @@ -0,0 +1,203 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from contextlib import nullcontext +from unittest import mock + +import pytest +from kubernetes.client import CoreV1Api, models as k8s +from kubernetes.client.rest import ApiException + +from airflow.providers.cncf.kubernetes.triggers.kueue import ( + KubernetesKueueJobTrigger, + KueueJobState, + build_job_event, + collect_xcom_results, + get_job_state, +) +from airflow.providers.cncf.kubernetes.utils.pod_manager import PodManager, XComRetrievalError +from airflow.triggers.base import TriggerEvent + +TRIGGER_PATH = "airflow.providers.cncf.kubernetes.triggers.kueue.{}" +JOB_NAME = "job-test" +NAMESPACE = "team-a" +TRIGGER_KWARGS = { + "job_name": JOB_NAME, + "job_namespace": NAMESPACE, + "base_container_name": "base", + "kubernetes_conn_id": "kubernetes_default", + "poll_interval": 0, + "cluster_context": "cluster", + "config_file": "/path/to/config", + "in_cluster": False, + "do_xcom_push": True, +} + + +def build_job(suspend: bool, condition: str | None = None, succeeded: int | None = None) -> k8s.V1Job: + conditions = ( + [k8s.V1JobCondition(type=condition, status="True", reason="BackoffLimitExceeded")] + if condition + else [] + ) + return k8s.V1Job( + metadata=k8s.V1ObjectMeta(name=JOB_NAME, namespace=NAMESPACE), + spec=k8s.V1JobSpec(suspend=suspend, template=k8s.V1PodTemplateSpec()), + status=k8s.V1JobStatus(conditions=conditions, succeeded=succeeded), + ) + + +def build_pod( + name: str, base: str | None = None, sidecar: str | None = None, phase: str = "Running" +) -> k8s.V1Pod: + """Build a pod of the Job, with its base and XCom sidecar containers "running", "terminated" or None.""" + statuses = [ + k8s.V1ContainerStatus( + name=container_name, + image="", + image_id="", + ready=state == "running", + restart_count=0, + state=k8s.V1ContainerState( + running=k8s.V1ContainerStateRunning() if state == "running" else None, + terminated=k8s.V1ContainerStateTerminated(exit_code=0) if state == "terminated" else None, + ), + ) + for container_name, state in (("base", base), ("airflow-xcom-sidecar", sidecar)) + if state + ] + return k8s.V1Pod( + metadata=k8s.V1ObjectMeta( + name=name, + namespace=NAMESPACE, + labels={"try_number": "1", "batch.kubernetes.io/job-name": JOB_NAME}, + ), + status=k8s.V1PodStatus(phase=phase, container_statuses=statuses), + ) + + +@pytest.mark.parametrize( + ("job", "expected_state"), + [ + (build_job(True), KueueJobState.SUSPENDED), + (build_job(False), KueueJobState.RUNNING), + (build_job(False, "Complete"), KueueJobState.FINISHED), + (build_job(False, "Failed"), KueueJobState.FINISHED), + ], + ids=["suspended", "running", "complete", "failed"], +) +def test_get_job_state(job, expected_state): + assert get_job_state(job) == expected_state + + +class TestCollectXcomResults: + def setup_method(self): + self.client = mock.create_autospec(CoreV1Api, instance=True) + self.pod_manager = mock.create_autospec(PodManager, instance=True) + + def test_reads_only_new_pods_with_terminated_base_and_running_sidecar(self): + ready = build_pod("ready", base="terminated", sidecar="running") + self.client.list_namespaced_pod.return_value = k8s.V1PodList( + items=[ + build_pod("pending"), + build_pod("running", base="running", sidecar="running"), + ready, + build_pod("already-read", base="terminated", sidecar="running"), + ] + ) + self.pod_manager.extract_xcom.return_value = '{"a": 1}' + xcom_results = {"already-read": "{}"} + + collect_xcom_results(self.client, self.pod_manager, JOB_NAME, NAMESPACE, "base", xcom_results) + + self.client.list_namespaced_pod.assert_called_once_with( + namespace=NAMESPACE, label_selector=f"batch.kubernetes.io/job-name={JOB_NAME}" + ) + self.pod_manager.extract_xcom.assert_called_once_with(ready, ignore_kill_failure=False) + assert xcom_results == {"already-read": "{}", "ready": '{"a": 1}'} + + @pytest.mark.parametrize( + ("error", "expectation"), + [ + (XComRetrievalError("sidecar stopped"), nullcontext()), + (ApiException(status=404), nullcontext()), + (ApiException(status=403), pytest.raises(ApiException)), + ], + ids=["sidecar-stopped", "pod-deleted", "other-api-error"], + ) + def test_skips_pods_stopped_or_deleted_by_preemption(self, error, expectation): + self.client.list_namespaced_pod.return_value = k8s.V1PodList( + items=[build_pod("preempted", base="terminated", sidecar="running")] + ) + self.pod_manager.extract_xcom.side_effect = error + xcom_results: dict[str, str] = {} + + with expectation: + collect_xcom_results(self.client, self.pod_manager, JOB_NAME, NAMESPACE, "base", xcom_results) + assert xcom_results == {} + + +class TestKubernetesKueueJobTrigger: + def test_serialize(self): + assert KubernetesKueueJobTrigger(**TRIGGER_KWARGS).serialize() == ( + "airflow.providers.cncf.kubernetes.triggers.kueue.KubernetesKueueJobTrigger", + TRIGGER_KWARGS, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("do_xcom_push", "expected_collect_calls", "expected_xcom_result"), + [ + pytest.param(True, 2, {"pod": "{}"}, id="xcom-read-while-running"), + pytest.param(False, 0, {}, id="without-xcom"), + ], + ) + @mock.patch(TRIGGER_PATH.format("collect_xcom_results"), autospec=True) + @mock.patch(TRIGGER_PATH.format("KubernetesKueueJobTrigger.sync_hook")) + @mock.patch(TRIGGER_PATH.format("KubernetesKueueJobTrigger.hook")) + async def test_run_waits_through_preemption( + self, + mock_hook, + mock_sync_hook, + mock_collect, + do_xcom_push, + expected_collect_calls, + expected_xcom_result, + ): + """Kueue admits the Job, preempts it, admits it again, then the Job completes.""" + jobs = [ + build_job(True), + build_job(False), + build_job(True), + build_job(False), + build_job(False, "Complete"), + ] + mock_hook.get_job_status = mock.AsyncMock(side_effect=jobs) + + def read_xcom(client, pod_manager, job_name, namespace, base_container_name, xcom_results): + xcom_results["pod"] = "{}" + + mock_collect.side_effect = read_xcom + trigger = KubernetesKueueJobTrigger(**TRIGGER_KWARGS) + trigger.do_xcom_push = do_xcom_push + + events = [event async for event in trigger.run()] + + assert events == [TriggerEvent(build_job_event(jobs[-1], expected_xcom_result))] + mock_hook.get_job_status.assert_awaited_with(name=JOB_NAME, namespace=NAMESPACE) + assert mock_collect.call_count == expected_collect_calls