shahar1 commented on code in PR #75658:
URL: https://github.com/apache/airflow/pull/75658#discussion_r4484651136
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py:
##########
@@ -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(
Review Comment:
**major** — A failed Job now writes no pod logs in sync mode. This method
always delegates to `KubernetesJobOperator.execute_complete`, which raises on
`event["status"] == "error"` before its `if self.get_logs:` loop
(`operators/job.py:341-350`). Its `finally` block then deletes the pods under
the default `on_finish_action="delete_pod"`. Before this PR the sync path
fetched logs first and raised afterwards (`job.py:252-265`), so users lose logs
in the case where they need them most.
The sync path also used
`fetch_requested_container_logs(containers=self.container_logs)`. `_write_logs`
reads only the base container, so `container_logs` is now ignored.
Please write the logs here before delegating, honouring `container_logs`.
`test_execute_raises_when_job_failed` uses `set_pods([])`, so it cannot catch
this; a case with a failed pod would. Rule: "Flag any changed or added
behaviour without a corresponding test".
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py:
##########
@@ -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:
Review Comment:
**minor** — `pods` now includes every failed pod of retried attempts, and
the `_write_logs` loop in `super()` is unguarded (`KubernetesPodOperator`
catches `(HTTPError, ApiException)` around it at `pod.py:1244-1246`). A pod
with no readable logs, such as an evicted pod or one on a node that is gone,
would raise here. The task then fails after the Job is `Complete`, and a retry
reruns the whole Job. I have not reproduced this against a live API server.
Please guard the log read per pod.
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.py:
##########
@@ -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)
Review Comment:
**minor** — This repeats the setup of the parent `execute` but drops its
`parallelism < 1` check. `parallelism=0` used to raise "parallelism cannot be
less than 1 with `wait_until_job_complete=True`". Now the Job never creates
pods, and this loop polls forever. Please keep the check.
##########
providers/cncf/kubernetes/docs/operators.rst:
##########
@@ -935,6 +952,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.
Review Comment:
**nit** — The wait has no timeout of its own, so a wrong `queue_name` leaves
the Job suspended and the task waiting indefinitely. Before this PR the task
failed quickly with "No pods running". Please state this here and say how to
bound the wait.
##########
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:
Review Comment:
**minor** — The 404 branch only sees errors from `read_pod`.
`kubernetes.stream` wraps exec failures as `ApiException(status=0, reason=...)`
(`ws_client.py:596-597`), so a pod that disappears at exec time is re-raised
and fails the task while the Job keeps running. The `pod-deleted` test case
only models 404.
Separately, `extract_xcom_json` raises `ValueError` from `json.loads` when a
preempted or failed pod left a partial `return.json`, while the docs say the
XCom of those pods is dropped.
Suggested fix: on status 0, re-read the pod and skip it only if it is gone.
A blanket skip would hide a 403 on exec, and the sidecar would then keep the
Job alive forever. Treat invalid JSON as unread, since the final
`len(xcom_result) < succeeded` check still catches a succeeded pod.
##########
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(
Review Comment:
**minor** — `get_job_status` and `find_pod` use `@generic_api_retry`, but
this `list_namespaced_pod` call does not. One 503 during a wait that can last
days fails the task. The Job stays in Kueue holding quota, and the retry
creates a second Job. `operators/kueue.py:211` has the same gap, at a point
where the Job has already finished. A small `@generic_api_retry` helper could
serve both.
##########
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
<howto/connection:kubernetes>`
+ 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] = {}
Review Comment:
**minor** — `xcom_results` lives only in memory, while the deferring docs
say: "You should assume that a trigger instance can run *more than once*" and
"A trigger can suddenly be removed from one triggerer service and started on a
new one." After a triggerer restart, pods already read have a dead sidecar, so
the new instance skips them. The task then ends in `XComRetrievalError` after
the Job succeeded, and a retry reruns the Job.
This fails loudly and `KubernetesJobTrigger` has a similar limit, so it is
not a regression. With `completions=N` the exposure lasts most of the Job,
though, and a triggerer deployment is routine. I don't see a simple fix because
the sidecar must die for the pod to finish, so at minimum please document this
in the XCom section.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
