Repository: apache/airflow
Pull request: AIP-72: Port over executor_config for K8sExecutor when using task sdk
<!--
 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.
 -->

<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.

Feel free to ping committers for the review!

In case of an existing issue, reference it using one of the following:

closes: #ISSUE
related: #ISSUE

How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->

closes: #46949 

This PR intends to complete the porting of KE to use task sdk as was done in #46860.

When we are sending the `ExecuteWorkload` payload to the KE pod, we do not need to send the `executor_config` in it. That is irrelevant to `supervise` function and is only using while creating the pod spec, like here: https://github.com/apache/airflow/blob/main/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py#L311-L316


<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).


Problem statement:
Port over `executor_config` for K8sExecutor using task sdk
### Body

This issue tracks the task for K8sExecutor now that its ported to use task sdk

### Committer

- [x] I acknowledge that I am a maintainer/committer of the Apache Airflow project.

Hints from issue discussion:

Commit message: Merge branch 'main' into AIP72-K8sExecutor-executor_config

Patch to review:
diff --git a/airflow/example_dags/example_kubernetes_executor.py b/airflow/example_dags/example_kubernetes_executor.py
index 395398ee93c8f..27ffb4cdad6d7 100644
--- a/airflow/example_dags/example_kubernetes_executor.py
+++ b/airflow/example_dags/example_kubernetes_executor.py
@@ -149,18 +149,6 @@ def non_root_task():
             print_stuff()
 
         third_task = non_root_task()
-
-        executor_config_other_ns = {
-            "pod_override": k8s.V1Pod(
-                metadata=k8s.V1ObjectMeta(namespace="test-namespace", labels={"release": "stable"})
-            )
-        }
-
-        @task(executor_config=executor_config_other_ns)
-        def other_namespace_task():
-            print_stuff()
-
-        other_ns_task = other_namespace_task()
         worker_container_repository = conf.get("kubernetes_executor", "worker_container_repository")
         worker_container_tag = conf.get("kubernetes_executor", "worker_container_tag")
 
@@ -231,9 +219,4 @@ def task_with_resource_limits():
 
         four_task = task_with_resource_limits()
 
-        (
-            start_task()
-            >> [volume_task, other_ns_task, sidecar_task]
-            >> third_task
-            >> [base_image_task, four_task]
-        )
+        (start_task() >> [volume_task, sidecar_task] >> third_task >> [base_image_task, four_task])
diff --git a/airflow/executors/workloads.py b/airflow/executors/workloads.py
index 5cb0102ce66df..3be2ad22d2399 100644
--- a/airflow/executors/workloads.py
+++ b/airflow/executors/workloads.py
@@ -61,6 +61,7 @@ class TaskInstance(BaseModel):
     pool_slots: int
     queue: str
     priority_weight: int
+    executor_config: dict | None = None
 
     # TODO: Task-SDK: Can we replace TastInstanceKey with just the uuid across the codebase?
     @property
diff --git a/kubernetes_tests/test_kubernetes_executor.py b/kubernetes_tests/test_kubernetes_executor.py
index 92e58d98118bd..63d1389b17103 100644
--- a/kubernetes_tests/test_kubernetes_executor.py
+++ b/kubernetes_tests/test_kubernetes_executor.py
@@ -26,9 +26,6 @@
 
 @pytest.mark.skipif(EXECUTOR != "KubernetesExecutor", reason="Only runs on KubernetesExecutor")
 class TestKubernetesExecutor(BaseK8STest):
-    @pytest.mark.skip(
-        reason="TODO: AIP-72 Porting over executor_config not yet done. Remove once #46892 is handled"
-    )
     @pytest.mark.execution_timeout(300)
     def test_integration_run_dag(self):
         dag_id = "example_kubernetes_executor"
@@ -54,9 +51,6 @@ def test_integration_run_dag(self):
         )
 
     @pytest.mark.execution_timeout(300)
-    @pytest.mark.skip(
-        reason="TODO: AIP-72 Porting over executor_config not yet done. Remove once #46892 is handled"
-    )
     def test_integration_run_dag_with_scheduler_failure(self):
         dag_id = "example_kubernetes_executor"
 
@@ -75,15 +69,6 @@ def test_integration_run_dag_with_scheduler_failure(self):
             timeout=300,
         )
 
-        self.monitor_task(
-            host=self.host,
-            dag_run_id=dag_run_id,
-            dag_id=dag_id,
-            task_id="other_namespace_task",
-            expected_final_state="success",
-            timeout=300,
-        )
-
         self.ensure_dag_expected_state(
             host=self.host,
             logical_date=logical_date,
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 3f6f49070ae0a..65a7f5225038a 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
@@ -291,8 +291,7 @@ def _process_workloads(self, workloads: list[workloads.All]) -> None:
             key = w.ti.key  # type: ignore[union-attr]
             queue = w.ti.queue  # type: ignore[union-attr]
 
-            # TODO: will be handled by https://github.com/apache/airflow/issues/46892
-            executor_config = {}  # type: ignore[var-annotated]
+            executor_config = w.ti.executor_config or {}  # type: ignore[union-attr]
 
             del self.queued_tasks[key]
             self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config)  # type: ignore[arg-type]
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 dda6b6e6b10be..0ef731783e62f 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
@@ -394,7 +394,7 @@ def run_next(self, next_job: KubernetesJobType) -> None:
 
             if isinstance(command[0], ExecuteTask):
                 workload = command[0]
-                ser_input = workload.model_dump_json()
+                ser_input = workload.model_dump_json(exclude={"ti": {"executor_config"}})
                 command = [
                     "python",
                     "-m",


Code review comments:
diff_hunk: @@ -291,8 +291,7 @@ def _process_workloads(self, workloads: list[workloads.All]) -> None:
             key = w.ti.key  # type: ignore[union-attr]
             queue = w.ti.queue  # type: ignore[union-attr]
 
-            # TODO: will be handled by https://github.com/apache/airflow/issues/46892
-            executor_config = {}  # type: ignore[var-annotated]
+            executor_config = w.ti.executor_config or {}  # type: ignore[union-attr]
path: providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
text: @user1:
Instead of the "ignore[union-attr]" we should do a type check and fail if `w` isn't ExecuteTask. Drive by refactor :) 

@author:
Thank you, TIL that mypy can be smart enough to gauge that :)