This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 512a0445875 Decide pod_template and image based on Coordinator for
lang-SDK tasks on KubernetesExecutor (#68713)
512a0445875 is described below
commit 512a0445875c950ff50f9fbd890751da20e92a09
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Mon Jun 29 20:10:02 2026 +0900
Decide pod_template and image based on Coordinator for lang-SDK tasks on
KubernetesExecutor (#68713)
---
.../kubernetes/executors/kubernetes_executor.py | 82 ++++++++-
.../executors/kubernetes_executor_types.py | 1 +
.../executors/kubernetes_executor_utils.py | 3 +-
.../executors/test_kubernetes_executor.py | 198 ++++++++++++++++++++-
4 files changed, 278 insertions(+), 6 deletions(-)
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
index 4ddb189390b..154024add26 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
@@ -198,6 +198,60 @@ class KubernetesExecutor(BaseExecutor):
scheduler_job_id=self.scheduler_job_id,
)
+ def _coordinator_extra(self, queue: str | None) -> dict[str, Any] | None:
+ """
+ Return the ``extra`` mapping a coordinator declares for *queue*, if
any.
+
+ Read from the coordinator's declarative ``[sdk]`` config without
importing
+ or instantiating the coordinator. The coordinator manager only exists
on
+ Airflow 3.3+; on older Task SDKs the import fails and we fall back to
no
+ extra. A malformed ``[sdk] coordinators`` / ``queue_to_coordinator``
config
+ must not crash the scheduler on this first lookup either, so an invalid
+ config also falls back to no extra. The exception types are imported
from
+ ``airflow.sdk`` so they match whatever Task SDK actually raised them.
+ """
+ if not queue:
+ return None
+ try:
+ from airflow.sdk.exceptions import AirflowConfigException
+ from airflow.sdk.execution_time.coordinator import
get_coordinator_manager
+ except ImportError:
+ return None
+ try:
+ return get_coordinator_manager().extra_for_queue(queue)
+ except (AirflowConfigException, ValueError):
+ self.log.warning(
+ "Ignoring coordinator config for queue %s: invalid [sdk]
coordinator config",
+ queue,
+ exc_info=True,
+ )
+ return None
+
+ def _coordinator_pod_template_file(self, extra: dict[str, Any]) -> str |
None:
+ """
+ Return the pod template declared in a coordinator's *extra* mapping,
if any.
+
+ Lets a queue routed to a non-Python coordinator (via ``[sdk]
+ queue_to_coordinator``) launch its worker pod from a
coordinator-specific
+ template — for example an image carrying the JVM for a Java
coordinator.
+ """
+ return extra.get("pod_template_file")
+
+ def _coordinator_kube_image(self, extra: dict[str, Any]) -> str | None:
+ """
+ Return the worker base image declared in a coordinator's *extra*
mapping, if any.
+
+ The base container image is never taken from a pod template; it comes
+ from ``kube_image``
(``worker_container_repository:worker_container_tag``)
+ or a per-task ``pod_override``. A coordinator may declare its own
+ ``worker_container_repository`` and ``worker_container_tag`` in
``extra``
+ (e.g. a JRE-bearing image for a Java coordinator); both are required to
+ compose an override, otherwise the executor default applies.
+ """
+ if (repo := extra.get("worker_container_repository")) and (tag :=
extra.get("worker_container_tag")):
+ return f"{repo}:{tag}"
+ return None
+
def execute_async(
self,
key: TaskInstanceKey,
@@ -225,8 +279,34 @@ class KubernetesExecutor(BaseExecutor):
pod_template_file = executor_config.get("pod_template_file", None)
else:
pod_template_file = None
+
+ coordinator_kube_image: str | None = None
+ if (coordinator_extra := self._coordinator_extra(queue)) is not None:
+ # A coordinator-level pod_template wins (e.g. a JVM image for
JavaCoordinator)
+ coordinator_pod_template_file =
self._coordinator_pod_template_file(coordinator_extra)
+ if coordinator_pod_template_file is not None:
+ self.log.debug(
+ "Using coordinator-declared pod template %s for task %s in
queue %s",
+ coordinator_pod_template_file,
+ key,
+ queue,
+ )
+ pod_template_file = coordinator_pod_template_file
+
+ # The base image is not carried by a pod template, so a
coordinator routes
+ # its worker base image separately (e.g. a JRE image for a Java
queue).
+ if (coordinator_kube_image :=
self._coordinator_kube_image(coordinator_extra)) is not None:
+ self.log.debug(
+ "Using coordinator-declared base image %s for task %s in
queue %s",
+ coordinator_kube_image,
+ key,
+ queue,
+ )
+
self.event_buffer[key] = (TaskInstanceState.QUEUED,
self.scheduler_job_id)
- self.task_queue.put(KubernetesJob(key, command, kube_executor_config,
pod_template_file))
+ self.task_queue.put(
+ KubernetesJob(key, command, kube_executor_config,
pod_template_file, coordinator_kube_image)
+ )
def queue_workload(self, workload: workloads.All, session: Session | None)
-> None:
from airflow.executors import workloads
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
index f8e03f1f04c..45d2ea301f7 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
@@ -75,6 +75,7 @@ class KubernetesJob(NamedTuple):
command: Sequence[str]
kube_executor_config: Any
pod_template_file: str | None
+ kube_image: str | None = None
ALL_NAMESPACES = "ALL_NAMESPACES"
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
index af719ada9e4..787ea73237f 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
@@ -558,6 +558,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
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
dag_id, task_id, run_id, try_number, map_index = key
if len(command) == 1:
@@ -586,7 +587,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
pod_id=create_unique_id(dag_id, task_id),
dag_id=dag_id,
task_id=task_id,
- kube_image=self.kube_config.kube_image,
+ kube_image=kube_image,
try_number=try_number,
map_index=map_index,
date=None,
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
index bc1c2a97f55..65f5820482b 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
@@ -39,6 +39,7 @@ from
airflow.providers.cncf.kubernetes.executors.kubernetes_executor import (
)
from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types
import (
ADOPTED,
+ KubernetesJob,
KubernetesResults,
KubernetesWatch,
)
@@ -66,7 +67,11 @@ except ImportError:
from airflow.utils.state import State, TaskInstanceState
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
+from tests_common.test_utils.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+ AIRFLOW_V_3_3_PLUS,
+)
try:
# Check whether a module-level function from stats is importable.
@@ -863,11 +868,10 @@ class TestKubernetesExecutor:
assert not executor.task_queue.empty()
task = executor.task_queue.get_nowait()
- _, _, expected_executor_config, expected_pod_template_file =
task
executor.task_queue.task_done()
# Test that the correct values have been put to queue
- assert expected_executor_config.metadata.labels == {"release":
"stable"}
- assert expected_pod_template_file == executor_template_file
+ assert task.kube_executor_config.metadata.labels ==
{"release": "stable"}
+ assert task.pod_template_file == executor_template_file
self.kubernetes_executor.kube_scheduler.run_next(task)
mock_run_pod_async.assert_called_once_with(
@@ -915,6 +919,192 @@ class TestKubernetesExecutor:
finally:
executor.end()
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+ @pytest.mark.parametrize(
+ ("coordinator_extra", "executor_config", "expected_template"),
+ [
+ pytest.param(
+ {"pod_template_file": "/coord/java.yaml"},
+ None,
+ "/coord/java.yaml",
+ id="coordinator-template-used",
+ ),
+ pytest.param(
+ {"pod_template_file": "/coord/java.yaml"},
+ {"pod_template_file": "/from/executor_config.yaml"},
+ "/coord/java.yaml",
+ id="coordinator-template-wins",
+ ),
+ pytest.param(
+ None,
+ {"pod_template_file": "/from/executor_config.yaml"},
+ "/from/executor_config.yaml",
+ id="executor-config-used-without-coordinator",
+ ),
+ ],
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ @mock.patch.object(KubernetesExecutor, "_coordinator_extra")
+ def test_coordinator_pod_template_file_used_for_queue(
+ self,
+ mock_coordinator_extra,
+ mock_get_kube_client,
+ mock_kubernetes_job_watcher,
+ coordinator_extra,
+ executor_config,
+ expected_template,
+ ):
+ """A queue coordinator's template overrides executor_config; without a
coordinator, executor_config is used."""
+ mock_coordinator_extra.return_value = coordinator_extra
+ executor = self.kubernetes_executor
+ executor.start()
+ try:
+ executor.execute_async(
+ key=TaskInstanceKey("dag", "task", "run_id", 1, -1),
+ queue="java",
+ command=["airflow", "tasks", "run", "true", "some_parameter"],
+ executor_config=executor_config,
+ )
+ assert not executor.task_queue.empty()
+ queued_job = executor.task_queue.get_nowait()
+ executor.task_queue.task_done()
+ assert queued_job.pod_template_file == expected_template
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+ @pytest.mark.parametrize(
+ ("extra", "expected"),
+ [
+ pytest.param({"pod_template_file": "/coord/go.yaml"},
"/coord/go.yaml", id="template-in-extra"),
+ pytest.param({"other": "value"}, None,
id="extra-without-template"),
+ ],
+ )
+ def test_coordinator_pod_template_file_reads_extra(self, extra, expected):
+ """The template is read straight from the coordinator ``extra``
mapping passed in."""
+ assert self.kubernetes_executor._coordinator_pod_template_file(extra)
== expected
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ @mock.patch.object(
+ KubernetesExecutor,
+ "_coordinator_extra",
+ return_value={"worker_container_repository": "repo/java",
"worker_container_tag": "1"},
+ )
+ def test_coordinator_kube_image_carried_on_job(
+ self,
+ mock_coordinator_extra,
+ mock_get_kube_client,
+ mock_kubernetes_job_watcher,
+ ):
+ """A coordinator base image resolved by queue rides on the queued
job."""
+ executor = self.kubernetes_executor
+ executor.start()
+ try:
+ executor.execute_async(
+ key=TaskInstanceKey("dag", "task", "run_id", 1, -1),
+ queue="java",
+ command=["airflow", "tasks", "run", "true", "some_parameter"],
+ )
+ queued_job = executor.task_queue.get_nowait()
+ executor.task_queue.task_done()
+ assert queued_job.kube_image == "repo/java:1"
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+ @pytest.mark.parametrize(
+ ("extra", "expected"),
+ [
+ pytest.param(
+ {"worker_container_repository": "repo/java",
"worker_container_tag": "17"},
+ "repo/java:17",
+ id="repository-and-tag",
+ ),
+ pytest.param({"worker_container_repository": "repo/java"}, None,
id="repository-only"),
+ pytest.param({"worker_container_tag": "17"}, None, id="tag-only"),
+ pytest.param({"other": "value"}, None, id="extra-without-image"),
+ ],
+ )
+ def test_coordinator_kube_image_reads_extra(self, extra, expected):
+ """The base image is composed straight from the coordinator ``extra``
mapping passed in."""
+ assert self.kubernetes_executor._coordinator_kube_image(extra) ==
expected
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+ def test_coordinator_extra_skips_lookup_without_queue(self):
+ """No queue means no coordinator lookup (and no Task SDK import)."""
+ with
mock.patch("airflow.sdk.execution_time.coordinator.get_coordinator_manager") as
mock_get_manager:
+ assert self.kubernetes_executor._coordinator_extra(None) is None
+ mock_get_manager.assert_not_called()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+ def test_coordinator_extra_returns_none_on_old_task_sdk(self):
+ """Pre-3.3 Task SDKs lack get_coordinator_manager; the import error
falls back to None."""
+ with mock.patch.dict("sys.modules",
{"airflow.sdk.execution_time.coordinator": None}):
+ assert self.kubernetes_executor._coordinator_extra("java") is None
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="The coordinator
interface only support since 3.3+")
+ @pytest.mark.parametrize("exc_type", ["airflow_config_exception",
"value_error"])
+ def test_coordinator_extra_falls_back_on_invalid_config(self, exc_type):
+ """A malformed ``[sdk]`` coordinator config must degrade gracefully,
not crash the scheduler."""
+ if exc_type == "airflow_config_exception":
+ from airflow.sdk.exceptions import AirflowConfigException
+
+ exc = AirflowConfigException("invalid json")
+ else:
+ exc = ValueError("invalid coordinator key")
+ with
mock.patch("airflow.sdk.execution_time.coordinator.get_coordinator_manager") as
mock_get_manager:
+ mock_get_manager.return_value.extra_for_queue.side_effect = exc
+ assert self.kubernetes_executor._coordinator_extra("java") is None
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
+ )
+ @pytest.mark.parametrize(
+ ("job_image", "use_default"),
+ [
+ pytest.param("repo/java:17", False, id="job-image-wins"),
+ pytest.param(None, True, id="falls-back-to-kube_config"),
+ ],
+ )
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.AirflowKubernetesScheduler.run_pod_async"
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.PodGenerator")
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_base_pod_from_template"
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_run_next_applies_job_kube_image(
+ self,
+ mock_get_kube_client,
+ mock_get_base_pod,
+ mock_pod_generator,
+ mock_run_pod_async,
+ job_image,
+ use_default,
+ ):
+ """``run_next`` uses the job's coordinator image, falling back to the
kube_config default."""
+ executor = self.kubernetes_executor
+ executor.start()
+ try:
+ scheduler = executor.kube_scheduler
+ scheduler.run_next(
+ KubernetesJob(
+ key=TaskInstanceKey("dag", "task", "run_id", 1, -1),
+ command=["airflow", "tasks", "run", "true",
"some_parameter"],
+ kube_executor_config=None,
+ pod_template_file=None,
+ kube_image=job_image,
+ )
+ )
+ expected = scheduler.kube_config.kube_image if use_default else
job_image
+ assert
mock_pod_generator.construct_pod.call_args.kwargs["kube_image"] == expected
+ finally:
+ executor.end()
+
@pytest.mark.db_test
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
