SameerMesiah97 commented on code in PR #70475: URL: https://github.com/apache/airflow/pull/70475#discussion_r3825381376
########## airflow-core/src/airflow/example_dags/example_deadline_callback.py: ########## @@ -0,0 +1,127 @@ +# +# 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. +""" +Example DAG demonstrating ExecutorCallback (DeadlineAlert) support in the Kubernetes Executor. + +The deadline is fixed to a date in the past so the callback fires immediately after each +DAG run is created, exercising the full callback-pod lifecycle on Kubernetes. +""" + +from __future__ import annotations + +import time +from datetime import datetime, timedelta, timezone + +import pendulum + +from airflow.sdk import DAG, task + +try: + from airflow.sdk.definitions.callback import SyncCallback + from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference + + _DEADLINE_AVAILABLE = True +except ImportError: + _DEADLINE_AVAILABLE = False + + +def deadline_callback(context, **kwargs): + """Simple deadline alert callback – logs context and exits cleanly.""" + dag_run = context.get("dag_run", {}) + dag_id = dag_run.get("dag_id", "unknown") + run_id = dag_run.get("dag_run_id", "unknown") + print(f"[deadline_callback] Deadline alert fired for dag_id={dag_id} run_id={run_id}") + + +def slow_deadline_callback(context, **kwargs): + """Slow deadline alert callback that sleeps 30s – used for scheduler-restart tests.""" + dag_run = context.get("dag_run", {}) + dag_id = dag_run.get("dag_id", "unknown") + run_id = dag_run.get("dag_run_id", "unknown") + print(f"[slow_deadline_callback] Starting for dag_id={dag_id} run_id={run_id}") + time.sleep(30) + print(f"[slow_deadline_callback] Done for dag_id={dag_id} run_id={run_id}") + + +def failing_deadline_callback(**_): + """Intentionally raises – used by the K8s executor failure integration test.""" + raise RuntimeError("Intentional callback failure for testing") + + +_PAST_DEADLINE = datetime(2020, 1, 1, tzinfo=timezone.utc) + +if _DEADLINE_AVAILABLE: + with DAG( + dag_id="example_deadline_callback", + schedule=None, + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + catchup=False, + tags=["example", "deadline", "callback"], + deadline=DeadlineAlert( + reference=DeadlineReference.FIXED_DATETIME(_PAST_DEADLINE), + interval=timedelta(hours=1), + callback=SyncCallback("airflow.example_dags.example_deadline_callback.deadline_callback"), + ), + ) as dag: + + @task + def dummy_task(): + """Placeholder task; the interesting work happens in the deadline callback.""" + print("dummy_task executed") + + dummy_task() + + with DAG( + dag_id="example_deadline_callback_slow", + schedule=None, + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + catchup=False, + tags=["example", "deadline", "callback"], + deadline=DeadlineAlert( + reference=DeadlineReference.FIXED_DATETIME(_PAST_DEADLINE), + interval=timedelta(hours=1), + callback=SyncCallback("airflow.example_dags.example_deadline_callback.slow_deadline_callback"), + ), + ) as dag_slow: + + @task + def dummy_task_slow(): + """Placeholder task for the slow-callback DAG.""" + print("dummy_task_slow executed") + + dummy_task_slow() + + with DAG( + dag_id="example_deadline_callback_failing", + schedule=None, + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + catchup=False, + tags=["example", "deadline", "callback"], + deadline=DeadlineAlert( + reference=DeadlineReference.FIXED_DATETIME(_PAST_DEADLINE), + interval=timedelta(hours=1), + callback=SyncCallback("airflow.example_dags.example_deadline_callback.failing_deadline_callback"), + ), + ) as dag_failing: + + @task + def dummy_task_failing(): + """Placeholder task for the failing-callback DAG.""" + print("dummy_task_failing executed") + + dummy_task_failing() Review Comment: Given that this DAG covers multiple scenarios rather than 'standard usage', I wonder whether this belongs in example dags at all. Could we keep the example DAG focused on demonstrating Kubernetes executor callback support and move the scheduler-restart/failure fixtures into the relevant integration test setup? ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py: ########## @@ -45,8 +69,8 @@ class FailureDetails(TypedDict, total=False): class KubernetesResults(NamedTuple): """Results from Kubernetes task execution.""" Review Comment: ```suggestion """Results from Kubernetes workload execution.""" ``` ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py: ########## @@ -1186,7 +1262,13 @@ def _adopt_completed_pods(self, kube_client: client.CoreV1Api) -> None: self.log.info("Failed to adopt pod %s. Reason: %s", pod.metadata.name, e) continue - ti_id = annotations_to_key(pod.metadata.annotations) + ti_id: TaskInstanceKey | CallbackKey + if AIRFLOW_V_3_3_PLUS and CALLBACK_POD_ANNOTATION_KEY in pod.metadata.annotations: + from airflow.models.callback import CallbackKey + + ti_id = CallbackKey(id=pod.metadata.annotations[CALLBACK_POD_ANNOTATION_KEY]) + else: + ti_id = annotations_to_key(pod.metadata.annotations) Review Comment: The above comment on adopting callbacks applies here too. ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py: ########## @@ -576,46 +594,76 @@ def _health_check_kube_watchers(self): ResourceVersion().resource_version[namespace] = "0" self.kube_watchers[namespace] = self._make_kube_watcher(namespace) + def _get_base_worker_pod(self, pod_template_file: str | None) -> k8s.V1Pod: + """Load the base worker pod from the template, raising if it is missing.""" + base_worker_pod = get_base_pod_from_template(pod_template_file, self.kube_config) + if not base_worker_pod: + raise AirflowException( + f"could not find a valid worker template yaml at {self.kube_config.pod_template_file}" + ) + return base_worker_pod + def run_next(self, next_job: KubernetesJob) -> None: """Receives the next job to run, builds the pod, and creates it.""" pod = self._build_pod_request(next_job) + if pod is None: + # Callback workloads build and submit their own pod inside _build_pod_request + # (via _run_next_callback) and return None; there is nothing left to create here. + return # the watcher will monitor pods, so we do not block. self.run_pod_async(pod, **self.kube_config.kube_client_request_args) self.log.debug("Kubernetes Job created!") - def _build_pod_request(self, next_job: KubernetesJob) -> k8s.V1Pod: + def _build_pod_request(self, next_job: KubernetesJob) -> k8s.V1Pod | None: """ Build the worker pod request object for a job. - Performs no API calls. May raise ``PodMutationHookException`` or - ``PodReconciliationError`` from the pod-mutation hook / reconciliation. + Performs no API calls (aside from callback workloads -- see below). May raise + ``PodMutationHookException`` or ``PodReconciliationError`` from the pod-mutation + hook / reconciliation. + + Returns ``None`` for callback workloads: those build *and submit* their own pod via + ``_run_next_callback`` because callbacks are not batched through ``run_next_batch``'s + async-creation path, so the caller must not attempt to create a pod for a ``None`` result. """ key = next_job.key command = next_job.command kube_executor_config = next_job.kube_executor_config pod_template_file = next_job.pod_template_file kube_image = next_job.kube_image or self.kube_config.kube_image + # Callback workloads follow a separate pod-construction path. + if AIRFLOW_V_3_3_PLUS and len(command) == 1: + from airflow.executors.workloads import ExecuteCallback + + workload = command[0] + if isinstance(workload, ExecuteCallback): + self._run_next_callback(workload.key, workload, pod_template_file) + return None Review Comment: Why does `_build_pod_request` now also submit the pod for callback workloads? This previously had a clear side-effect-free contract where it only constructed and returned the pod, with submission handled by the caller. Could `_run_next_callback` just construct and return the callback pod so we retain that separation and callbacks can follow the same pod creation path as tasks? ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py: ########## @@ -461,6 +461,81 @@ def construct_pod( return pod + @classmethod + def construct_callback_pod( + cls, + *, + namespace: str, + scheduler_job_id: str, + callback_id: str, + dag_id: str, + run_id: str, + kube_image: str, + args: list[str], + base_worker_pod: k8s.V1Pod, + with_mutation_hook: bool = False, + ) -> k8s.V1Pod: + """Create a Pod for executing an ExecuteCallback workload.""" + from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types import ( + CALLBACK_POD_ANNOTATION_KEY, + CALLBACK_WORKLOAD_TYPE_KEY, + ) + + # Derive a k8s-safe pod name from the callback UUID. + short_id = callback_id.replace("-", "")[:8] + pod_id = add_unique_suffix(name=f"callback-{short_id}", rand_len=8, max_len=POD_NAME_MAX_LENGTH) + + try: + image = base_worker_pod.spec.containers[0].image Review Comment: Is the precedence here intentional? `kube_image` is effectively only a fallback since we copy the image from `base_worker_pod` into the dynamic pod whenever it exists. This also seems different from the normal task pod construction where `kube_image` is passed into the dynamic pod and the normal reconciliation determines the final result. Could we follow the same approach here? ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py: ########## @@ -165,16 +171,25 @@ def _run( if event["type"] == "ERROR": return self.process_error(event) annotations = task.metadata.annotations - task_instance_related_annotations = { - "dag_id": annotations["dag_id"], - "task_id": annotations["task_id"], - logical_date_key: annotations.get(logical_date_key), - "run_id": annotations.get("run_id"), - "try_number": annotations["try_number"], - } - map_index = annotations.get("map_index") - if map_index is not None: - task_instance_related_annotations["map_index"] = map_index + if AIRFLOW_V_3_3_PLUS and CALLBACK_POD_ANNOTATION_KEY in annotations: + # Callback pod: forward only the annotations that process_watcher_task needs. + # Callback pods carry callback_id instead of task_id/try_number. + task_instance_related_annotations = { Review Comment: nit: could be `workload_related_annotations` to be more accurate. ########## kubernetes-tests/tests/kubernetes_tests/test_kubernetes_executor.py: ########## @@ -239,3 +244,296 @@ def test_pod_failure_logging_non_failed_state(self, mock_log): # Verify that kube_client methods were not called executor.kube_client.read_namespaced_pod.assert_not_called() + + [email protected](EXECUTOR != "KubernetesExecutor", reason="Only runs on KubernetesExecutor") +class TestKubernetesExecutorCallbackSupport(BaseK8STest): + """ + Integration tests for ExecutorCallback (DeadlineAlert / SyncCallback) support in the + Kubernetes executor. + + Prerequisites: + - The ``example_deadline_callback`` and ``example_deadline_callback_slow`` DAGs must + be loaded in the cluster (they live in airflow-core/src/airflow/example_dags/ and + are baked into the k8s image via ``breeze k8s build-k8s-image``). + - The executor must be KubernetesExecutor. + """ + + _FAST_DAG_ID = "example_deadline_callback" + _SLOW_DAG_ID = "example_deadline_callback_slow" + _FAILING_DAG_ID = "example_deadline_callback_failing" + + _CALLBACK_LABEL = "airflow-workload-type=callback" + _CALLBACK_ANNOTATION_KEY = "callback_id" + + @classmethod + def _get_callback_pods(cls, namespace: str = "airflow") -> list[dict]: + raw = check_output( + [ + "kubectl", + "get", + "pods", + "-n", + namespace, + "-l", + cls._CALLBACK_LABEL, + "-o", + "json", + ] + ) + return json.loads(raw)["items"] + + @classmethod + def _wait_for_callback_pod(cls, run_id: str, namespace: str = "airflow", timeout: int = 120) -> dict: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for pod in cls._get_callback_pods(namespace): + annotations = pod.get("metadata", {}).get("annotations", {}) + if annotations.get("run_id") == run_id: + return pod + time.sleep(1) + raise AssertionError(f"No callback pod for run_id={run_id!r} appeared within {timeout}s") + + @staticmethod + def _wait_for_pod_phase( + pod_name: str, + phases: list[str], + namespace: str = "airflow", + timeout: int = 120, + ) -> str: + """ + Block until *pod_name* reaches one of *phases*; return the reached phase. + + Returns ``"Deleted"`` if the pod is not found. Include ``"Deleted"`` in *phases* + when executor-driven deletion counts as success — pods are only removed after they + succeed (``delete_worker_pods=True`` default). + """ + deadline = time.monotonic() + timeout + phase = "" + while time.monotonic() < deadline: + result = subprocess.run( + [ + "kubectl", + "get", + "pod", + pod_name, + "-n", + namespace, + "-o", + "jsonpath={.status.phase}", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0 and ( + "NotFound" in result.stderr or "not found" in result.stderr.lower() + ): + phase = "Deleted" + else: + phase = result.stdout.strip() + if phase in phases: + return phase + time.sleep(2) + raise AssertionError( + f"Pod {pod_name!r} did not reach {phases} within {timeout}s (last seen phase: {phase!r})" + ) + + @staticmethod + def _wait_for_pod_gone(pod_name: str, namespace: str = "airflow", timeout: int = 60) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + result = subprocess.run( + ["kubectl", "get", "pod", pod_name, "-n", namespace], + capture_output=True, + text=True, + check=False, + ) + # kubectl exits non-zero and prints "NotFound" when the pod is gone + if result.returncode != 0 and "NotFound" in result.stderr: + return + time.sleep(5) + raise AssertionError(f"Pod {pod_name!r} was not deleted within {timeout}s") + + def _trigger_dag_run(self, dag_id: str) -> str: + result_json = self.start_dag(dag_id=dag_id, host=self.host) + dag_runs = result_json.get("dag_runs", []) + matching = [r for r in dag_runs if r["dag_id"] == dag_id] + assert matching, f"No dag runs returned for dag_id={dag_id!r}" + newest = max(matching, key=lambda r: r["queued_at"]) + return newest["dag_run_id"] + + @pytest.mark.execution_timeout(300) + def test_deadline_callback_executes_on_kubernetes(self): + """ + A DAG with a past deadline fires a SyncCallback that is executed as a Kubernetes pod. + The pod must reach the Succeeded phase. + """ + dag_id = self._FAST_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + print(f"[{dag_id}] dag_run_id={dag_run_id}") + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + pod_name = pod["metadata"]["name"] + print(f"[{dag_id}] callback pod appeared: {pod_name}") + + # The executor deletes pods immediately after they succeed (delete_worker_pods=True + # default), so "Deleted" is equally valid evidence of success. + phase = self._wait_for_pod_phase(pod_name, ["Succeeded", "Failed", "Deleted"], timeout=120) + assert phase in ("Succeeded", "Deleted"), ( + f"Callback pod {pod_name!r} reached phase {phase!r} instead of Succeeded/Deleted" + ) + + @pytest.mark.execution_timeout(300) + def test_callback_pod_annotations_and_labels(self): + """ + The callback pod must carry the expected Airflow annotations (callback_id, dag_id, run_id) + and labels (airflow-workload-type=callback, kubernetes_executor=True, airflow-worker=…). + Its container command must invoke execute_workload. + """ + dag_id = self._FAST_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + annotations = pod["metadata"]["annotations"] + labels = pod["metadata"]["labels"] + containers = pod["spec"]["containers"] + + assert self._CALLBACK_ANNOTATION_KEY in annotations, ( + f"Annotation {self._CALLBACK_ANNOTATION_KEY!r} missing from pod. Got: {annotations}" + ) + callback_id = annotations[self._CALLBACK_ANNOTATION_KEY] + assert re.fullmatch(r"[0-9a-f-]{36}", callback_id), ( + f"callback_id {callback_id!r} does not look like a UUID" + ) + assert annotations.get("dag_id") == dag_id, ( + f"Expected dag_id annotation {dag_id!r}, got {annotations.get('dag_id')!r}" + ) + assert "run_id" in annotations, f"run_id annotation missing. Got: {annotations}" + assert annotations.get("run_id"), "run_id annotation must be non-empty" + + # Task-specific annotations must NOT be present on callback pods. + for forbidden in ("task_id", "try_number", "map_index"): + assert forbidden not in annotations, f"Unexpected annotation {forbidden!r} found on callback pod" + + assert labels.get("airflow-workload-type") == "callback", ( + f"Expected label airflow-workload-type=callback, got: {labels}" + ) + assert labels.get("kubernetes_executor") == "True", ( + f"Expected label kubernetes_executor=True, got: {labels}" + ) + assert "airflow-worker" in labels, f"airflow-worker label missing. Got: {labels}" + + assert containers, "No containers found in callback pod spec" + args = containers[0].get("args", []) or [] + cmd = containers[0].get("command", []) or [] + full_cmd = cmd + args + assert any("execute_workload" in part for part in full_cmd), ( + f"Container command does not include 'execute_workload'. Full command: {full_cmd}" + ) + assert any("--json-string" in part for part in full_cmd), ( + f"Container command does not include '--json-string'. Full command: {full_cmd}" + ) + + @pytest.mark.execution_timeout(300) + def test_deadline_callback_pod_failure(self): + """ + When the callback pod exits with a non-zero code the watcher must emit a FAILED + state and the executor must not crash. + + Uses ``example_deadline_callback_failing``, whose callback always raises + ``RuntimeError``, causing ``sys.exit(1)`` in the pod and a ``Failed`` phase. + Failed pods are NOT auto-deleted (``delete_worker_pods_on_failure=False`` default), + so the pod stays and we can assert its phase before cleaning it up manually. + """ + dag_id = self._FAILING_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + print(f"[{dag_id}] dag_run_id={dag_run_id}") + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + pod_name = pod["metadata"]["name"] + print(f"[{dag_id}] callback pod appeared: {pod_name}") + + try: + # Failed pods are NOT auto-deleted (delete_worker_pods_on_failure=False default). + phase = self._wait_for_pod_phase(pod_name, ["Failed"], timeout=120) + assert phase == "Failed", f"Expected Failed phase, got {phase!r}" + + # Executor must not have crashed — scheduler pod must still be Running. + result = subprocess.run( + [ + "kubectl", + "get", + "pod", + "-n", + "airflow", + "-l", + "component=scheduler", + "-o", + "jsonpath={.items[0].status.phase}", + ], + capture_output=True, + text=True, + check=False, + ) + scheduler_phase = result.stdout.strip() + assert scheduler_phase == "Running", ( + f"Scheduler pod is in phase {scheduler_phase!r} after callback failure" + ) + finally: + # Clean up the failed pod manually (won't be auto-deleted by executor). + subprocess.run( + ["kubectl", "delete", "pod", pod_name, "-n", "airflow", "--ignore-not-found"], + capture_output=True, + check=False, + ) + + @pytest.mark.execution_timeout(300) + def test_callback_pod_is_cleaned_up_after_success(self): + """ + After the callback pod reaches Succeeded, the executor must delete it so no + orphaned callback pods linger in the namespace. + """ + dag_id = self._FAST_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + pod_name = pod["metadata"]["name"] + + # If the pod is already gone ("Deleted"), that itself is proof of executor-driven + # cleanup — skip the explicit deletion wait in that case. + phase = self._wait_for_pod_phase(pod_name, ["Succeeded", "Failed", "Deleted"], timeout=120) + if phase != "Deleted": + self._wait_for_pod_gone(pod_name, timeout=60) + + @pytest.mark.execution_timeout(400) + def test_callback_pod_survives_scheduler_restart(self): + """ + A callback pod running in Kubernetes must complete even when the scheduler + is restarted mid-execution. After restart, the new scheduler must re-adopt + the pod via the watcher label selector and record the terminal state. + """ + dag_id = self._SLOW_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + + # Wait until the callback pod is actually Running before killing the scheduler. + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + pod_name = pod["metadata"]["name"] + pre_restart_phase = self._wait_for_pod_phase( + pod_name, ["Running", "Succeeded", "Failed", "Deleted"], timeout=60 + ) + + self._delete_airflow_pod("scheduler") + self.ensure_resource_health("airflow-scheduler") + print(f"[{dag_id}] Scheduler restarted; waiting for callback pod to complete.") + + if pre_restart_phase in ("Succeeded", "Deleted"): + # Pod already completed before the restart — the test goal (pod completion + # despite scheduler lifecycle) is still met. + return + + # The slow callback sleeps for 30s; allow plenty of time for completion. + phase = self._wait_for_pod_phase(pod_name, ["Succeeded", "Failed", "Deleted"], timeout=180) + assert phase in ("Succeeded", "Deleted"), ( + f"Callback pod {pod_name!r} reached {phase!r} after scheduler restart (expected Succeeded/Deleted)" + ) Review Comment: Not sure if this tests actually covers what it is supposed to. Consdier a situation where the new scheduler never adopts the callback pod. The pod would continue running independently and eventually reach `Succeeded` anyway, so this test would still pass. Should we explicitly verify that the callback pod is adopted by the new scheduler and that its terminal state is processed? That's if we decide that callbacks should be adopted. ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py: ########## @@ -1034,6 +1098,19 @@ def adopt_launched_task( if TYPE_CHECKING: assert self.scheduler_job_id + if AIRFLOW_V_3_3_PLUS and CALLBACK_POD_ANNOTATION_KEY in pod.metadata.annotations: + new_worker_id_label = self._make_safe_label_value(self.scheduler_job_id) + from kubernetes.client.rest import ApiException + + try: + kube_client.patch_namespaced_pod( + name=pod.metadata.name, + namespace=pod.metadata.namespace, + body={"metadata": {"labels": {"airflow-worker": new_worker_id_label}}}, + ) + except ApiException as e: + self.log.info("Failed to adopt pod %s. Reason: %s", pod.metadata.name, e) + return Review Comment: This seems inconsitent with what has been done for the other executors (for e.g. `LambdaExecutor`) wrt adding callback support. . Do we actually want callback workloads to participate in scheduler adoption here, or should the lifecycle semantics be consistent across executors? ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py: ########## @@ -728,6 +776,42 @@ def _close_async_pod_client(self) -> None: self._async_loop = None self._async_pod_client = None + def _run_next_callback( + self, key: CallbackKey, workload: ExecuteCallback, pod_template_file: str | None + ) -> None: + """Build and submit a callback pod for an ExecuteCallback workload.""" + base_worker_pod = self._get_base_worker_pod(pod_template_file) + + # Extract dag_id and run_id from the standardised log_path: + # "executor_callbacks/<dag_id>/<run_id>/<callback_id>" + log_parts = (workload.log_path or "").split("/") + dag_id = log_parts[1] if len(log_parts) >= 4 else "" + run_id = log_parts[2] if len(log_parts) >= 4 else "" + Review Comment: It seems rather odd to parse critical metadata from logs. Should we be deriving `dag_id` and `run_id` by parsing `workload.log_path` here? what if the format changes? Is this approach not too brittle? If the callback pod requires this metadata, should it be exposed directly by the workload instead? ########## providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py: ########## @@ -461,6 +461,81 @@ def construct_pod( return pod + @classmethod + def construct_callback_pod( + cls, + *, + namespace: str, + scheduler_job_id: str, + callback_id: str, + dag_id: str, + run_id: str, + kube_image: str, + args: list[str], + base_worker_pod: k8s.V1Pod, + with_mutation_hook: bool = False, + ) -> k8s.V1Pod: + """Create a Pod for executing an ExecuteCallback workload.""" + from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types import ( + CALLBACK_POD_ANNOTATION_KEY, + CALLBACK_WORKLOAD_TYPE_KEY, + ) + + # Derive a k8s-safe pod name from the callback UUID. + short_id = callback_id.replace("-", "")[:8] + pod_id = add_unique_suffix(name=f"callback-{short_id}", rand_len=8, max_len=POD_NAME_MAX_LENGTH) + + try: + image = base_worker_pod.spec.containers[0].image + if not image: + image = kube_image + except Exception: Review Comment: Exception here is too broad. Not necessarily unjustified but I wish to know why a narrower exception was not used/ ########## kubernetes-tests/tests/kubernetes_tests/test_kubernetes_executor.py: ########## @@ -239,3 +244,296 @@ def test_pod_failure_logging_non_failed_state(self, mock_log): # Verify that kube_client methods were not called executor.kube_client.read_namespaced_pod.assert_not_called() + + [email protected](EXECUTOR != "KubernetesExecutor", reason="Only runs on KubernetesExecutor") +class TestKubernetesExecutorCallbackSupport(BaseK8STest): + """ + Integration tests for ExecutorCallback (DeadlineAlert / SyncCallback) support in the + Kubernetes executor. + + Prerequisites: + - The ``example_deadline_callback`` and ``example_deadline_callback_slow`` DAGs must + be loaded in the cluster (they live in airflow-core/src/airflow/example_dags/ and + are baked into the k8s image via ``breeze k8s build-k8s-image``). + - The executor must be KubernetesExecutor. + """ + + _FAST_DAG_ID = "example_deadline_callback" + _SLOW_DAG_ID = "example_deadline_callback_slow" + _FAILING_DAG_ID = "example_deadline_callback_failing" + + _CALLBACK_LABEL = "airflow-workload-type=callback" + _CALLBACK_ANNOTATION_KEY = "callback_id" + + @classmethod + def _get_callback_pods(cls, namespace: str = "airflow") -> list[dict]: + raw = check_output( + [ + "kubectl", + "get", + "pods", + "-n", + namespace, + "-l", + cls._CALLBACK_LABEL, + "-o", + "json", + ] + ) + return json.loads(raw)["items"] + + @classmethod + def _wait_for_callback_pod(cls, run_id: str, namespace: str = "airflow", timeout: int = 120) -> dict: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for pod in cls._get_callback_pods(namespace): + annotations = pod.get("metadata", {}).get("annotations", {}) + if annotations.get("run_id") == run_id: + return pod + time.sleep(1) + raise AssertionError(f"No callback pod for run_id={run_id!r} appeared within {timeout}s") + + @staticmethod + def _wait_for_pod_phase( + pod_name: str, + phases: list[str], + namespace: str = "airflow", + timeout: int = 120, + ) -> str: + """ + Block until *pod_name* reaches one of *phases*; return the reached phase. + + Returns ``"Deleted"`` if the pod is not found. Include ``"Deleted"`` in *phases* + when executor-driven deletion counts as success — pods are only removed after they + succeed (``delete_worker_pods=True`` default). + """ + deadline = time.monotonic() + timeout + phase = "" + while time.monotonic() < deadline: + result = subprocess.run( + [ + "kubectl", + "get", + "pod", + pod_name, + "-n", + namespace, + "-o", + "jsonpath={.status.phase}", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0 and ( + "NotFound" in result.stderr or "not found" in result.stderr.lower() + ): + phase = "Deleted" + else: + phase = result.stdout.strip() + if phase in phases: + return phase + time.sleep(2) + raise AssertionError( + f"Pod {pod_name!r} did not reach {phases} within {timeout}s (last seen phase: {phase!r})" + ) + + @staticmethod + def _wait_for_pod_gone(pod_name: str, namespace: str = "airflow", timeout: int = 60) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + result = subprocess.run( + ["kubectl", "get", "pod", pod_name, "-n", namespace], + capture_output=True, + text=True, + check=False, + ) + # kubectl exits non-zero and prints "NotFound" when the pod is gone + if result.returncode != 0 and "NotFound" in result.stderr: + return + time.sleep(5) + raise AssertionError(f"Pod {pod_name!r} was not deleted within {timeout}s") + + def _trigger_dag_run(self, dag_id: str) -> str: + result_json = self.start_dag(dag_id=dag_id, host=self.host) + dag_runs = result_json.get("dag_runs", []) + matching = [r for r in dag_runs if r["dag_id"] == dag_id] + assert matching, f"No dag runs returned for dag_id={dag_id!r}" + newest = max(matching, key=lambda r: r["queued_at"]) + return newest["dag_run_id"] + + @pytest.mark.execution_timeout(300) + def test_deadline_callback_executes_on_kubernetes(self): + """ + A DAG with a past deadline fires a SyncCallback that is executed as a Kubernetes pod. + The pod must reach the Succeeded phase. + """ + dag_id = self._FAST_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + print(f"[{dag_id}] dag_run_id={dag_run_id}") + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + pod_name = pod["metadata"]["name"] + print(f"[{dag_id}] callback pod appeared: {pod_name}") + + # The executor deletes pods immediately after they succeed (delete_worker_pods=True + # default), so "Deleted" is equally valid evidence of success. + phase = self._wait_for_pod_phase(pod_name, ["Succeeded", "Failed", "Deleted"], timeout=120) + assert phase in ("Succeeded", "Deleted"), ( + f"Callback pod {pod_name!r} reached phase {phase!r} instead of Succeeded/Deleted" + ) + + @pytest.mark.execution_timeout(300) + def test_callback_pod_annotations_and_labels(self): + """ + The callback pod must carry the expected Airflow annotations (callback_id, dag_id, run_id) + and labels (airflow-workload-type=callback, kubernetes_executor=True, airflow-worker=…). + Its container command must invoke execute_workload. + """ + dag_id = self._FAST_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + annotations = pod["metadata"]["annotations"] + labels = pod["metadata"]["labels"] + containers = pod["spec"]["containers"] + + assert self._CALLBACK_ANNOTATION_KEY in annotations, ( + f"Annotation {self._CALLBACK_ANNOTATION_KEY!r} missing from pod. Got: {annotations}" + ) + callback_id = annotations[self._CALLBACK_ANNOTATION_KEY] + assert re.fullmatch(r"[0-9a-f-]{36}", callback_id), ( + f"callback_id {callback_id!r} does not look like a UUID" + ) + assert annotations.get("dag_id") == dag_id, ( + f"Expected dag_id annotation {dag_id!r}, got {annotations.get('dag_id')!r}" + ) + assert "run_id" in annotations, f"run_id annotation missing. Got: {annotations}" + assert annotations.get("run_id"), "run_id annotation must be non-empty" + + # Task-specific annotations must NOT be present on callback pods. + for forbidden in ("task_id", "try_number", "map_index"): + assert forbidden not in annotations, f"Unexpected annotation {forbidden!r} found on callback pod" + + assert labels.get("airflow-workload-type") == "callback", ( + f"Expected label airflow-workload-type=callback, got: {labels}" + ) + assert labels.get("kubernetes_executor") == "True", ( + f"Expected label kubernetes_executor=True, got: {labels}" + ) + assert "airflow-worker" in labels, f"airflow-worker label missing. Got: {labels}" + + assert containers, "No containers found in callback pod spec" + args = containers[0].get("args", []) or [] + cmd = containers[0].get("command", []) or [] + full_cmd = cmd + args + assert any("execute_workload" in part for part in full_cmd), ( + f"Container command does not include 'execute_workload'. Full command: {full_cmd}" + ) + assert any("--json-string" in part for part in full_cmd), ( + f"Container command does not include '--json-string'. Full command: {full_cmd}" + ) + + @pytest.mark.execution_timeout(300) + def test_deadline_callback_pod_failure(self): + """ + When the callback pod exits with a non-zero code the watcher must emit a FAILED + state and the executor must not crash. + + Uses ``example_deadline_callback_failing``, whose callback always raises + ``RuntimeError``, causing ``sys.exit(1)`` in the pod and a ``Failed`` phase. + Failed pods are NOT auto-deleted (``delete_worker_pods_on_failure=False`` default), + so the pod stays and we can assert its phase before cleaning it up manually. + """ + dag_id = self._FAILING_DAG_ID + dag_run_id = self._trigger_dag_run(dag_id) + print(f"[{dag_id}] dag_run_id={dag_run_id}") + + pod = self._wait_for_callback_pod(dag_run_id, timeout=120) + pod_name = pod["metadata"]["name"] + print(f"[{dag_id}] callback pod appeared: {pod_name}") + + try: + # Failed pods are NOT auto-deleted (delete_worker_pods_on_failure=False default). + phase = self._wait_for_pod_phase(pod_name, ["Failed"], timeout=120) + assert phase == "Failed", f"Expected Failed phase, got {phase!r}" Review Comment: The test description says we're verifying that the watcher emits FAILED, but asserting the Kubernetes pod phase only verifies that the callback container failed. Is there a way to assert that the failed result was actually processed by the executor? -- 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]
