This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch fix-k8s-executor-pod-override-pickle in repository https://gitbox.apache.org/repos/asf/airflow.git
commit fb24244909ce7baa4f2be55b7808c04b74af961d Author: Rahul Vats <[email protected]> AuthorDate: Mon Jun 22 09:55:26 2026 +0530 Fix KubernetesExecutor scheduler crash when pod_override is queued in-cluster When the scheduler runs in-cluster, the kubernetes client attaches an in-cluster Configuration to every V1Pod whose refresh_api_key_hook is a local closure (InClusterConfigLoader._set_config.<locals>._refresh_api_key). pickle cannot serialize a local closure, so putting a task's pod_override V1Pod on the executor's multiprocessing queue raises PicklingError and crashes the scheduler in a loop. This affects any in-cluster KubernetesExecutor deployment where a task sets a V1Pod pod_override, and is independent of the Airflow version; it surfaces with kubernetes client 36.x. Serialize the pod_override to a plain dict (dropping the Configuration) before it is queued, and rebuild the V1Pod worker-side in run_next. The worker already reconstructs its own kube client, so nothing is lost. This keeps pod_override working regardless of the kubernetes client version, instead of relying on a client version whose Configuration happens to be picklable. --- .../kubernetes/executors/kubernetes_executor.py | 15 +++++ .../executors/kubernetes_executor_utils.py | 5 ++ .../executors/test_kubernetes_executor.py | 67 +++++++++++++++++++++- 3 files changed, 85 insertions(+), 2 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 0dcc01537cb..85025b324e4 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 @@ -226,6 +226,21 @@ class KubernetesExecutor(BaseExecutor): pod_template_file = executor_config.get("pod_template_file", None) else: pod_template_file = None + + # Serialize any ``pod_override`` ``V1Pod`` to a plain dict before putting it on the + # multiprocessing queue. When the scheduler runs in-cluster, the kubernetes client attaches an + # in-cluster ``Configuration`` to every ``V1Pod`` whose ``refresh_api_key_hook`` is a local + # closure (``InClusterConfigLoader._set_config.<locals>._refresh_api_key``). ``pickle`` cannot + # serialize a local closure, so queuing a live ``V1Pod`` raises ``PicklingError`` and crashes the + # scheduler. ``run_next`` deserializes the dict back into a ``V1Pod`` worker-side. The same + # ``V1Pod`` object is also referenced by the workload's ``executor_config``, so sanitize that copy + # too (it is otherwise pickled as part of the workload, even though the worker rebuilds the pod + # override from ``kube_executor_config``). + if kube_executor_config is not None: + kube_executor_config = PodGenerator.serialize_pod(kube_executor_config) + if executor_config and executor_config.get("pod_override") is not None: + executor_config["pod_override"] = PodGenerator.serialize_pod(executor_config["pod_override"]) + self.event_buffer[key] = (TaskInstanceState.QUEUED, self.scheduler_job_id) self.task_queue.put(KubernetesJob(key, command, kube_executor_config, pod_template_file)) # We keep a temporary local record that we've handled this so we don't 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..698a1406884 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 @@ -559,6 +559,11 @@ class AirflowKubernetesScheduler(LoggingMixin): kube_executor_config = next_job.kube_executor_config pod_template_file = next_job.pod_template_file + # ``execute_async`` serializes the ``pod_override`` to a dict so it can be pickled onto the + # multiprocessing queue (a live in-cluster ``V1Pod`` is not picklable). Rebuild the ``V1Pod`` here. + if isinstance(kube_executor_config, dict): + kube_executor_config = PodGenerator.deserialize_model_dict(kube_executor_config) + dag_id, task_id, run_id, try_number, map_index = key if len(command) == 1: from airflow.executors.workloads import ExecuteTask 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..60f19b75904 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 @@ -865,8 +865,10 @@ class TestKubernetesExecutor: 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"} + # ``pod_override`` is serialized to a dict before being queued so it can be pickled onto + # the multiprocessing queue (a live in-cluster ``V1Pod`` is not picklable). + assert isinstance(expected_executor_config, dict) + assert expected_executor_config["metadata"]["labels"] == {"release": "stable"} assert expected_pod_template_file == executor_template_file self.kubernetes_executor.kube_scheduler.run_next(task) @@ -915,6 +917,67 @@ class TestKubernetesExecutor: finally: executor.end() + @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher") + @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_execute_async_queues_picklable_pod_override( + self, mock_get_kube_client, mock_kubernetes_job_watcher + ): + """Regression: a ``pod_override`` carrying an in-cluster ``Configuration`` must be picklable. + + When the scheduler runs in-cluster, the kubernetes client attaches a ``Configuration`` whose + ``refresh_api_key_hook`` is a local closure (``InClusterConfigLoader._set_config.<locals>. + _refresh_api_key``). ``pickle`` cannot serialize a local closure, so putting a live ``V1Pod`` on + the multiprocessing queue raised ``PicklingError`` and crashed the scheduler. ``execute_async`` + must serialize the pod to a dict so the queued ``KubernetesJob`` pickles cleanly. + """ + import pickle + + from kubernetes.client import Configuration + + pod_override = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(labels={"release": "stable"}), + spec=k8s.V1PodSpec(containers=[k8s.V1Container(name="base", image="airflow:3.6")]), + ) + + # Simulate the in-cluster config: an unpicklable local closure on the pod's Configuration. + def _make_unpicklable_hook(): + def _refresh_api_key(config): + return None + + return _refresh_api_key + + cfg = Configuration() + cfg.refresh_api_key_hook = _make_unpicklable_hook() + pod_override.metadata.local_vars_configuration = cfg + + # Sanity check: the raw pod is indeed not picklable (reproduces the crash pre-fix). + with pytest.raises((pickle.PicklingError, AttributeError, TypeError)): + pickle.dumps(pod_override) + + executor = self.kubernetes_executor + executor.start() + try: + executor.execute_async( + key=TaskInstanceKey("dag", "task", "run_id", 1, -1), + queue=None, + command=["airflow", "tasks", "run", "true", "some_parameter"], + executor_config={"pod_override": pod_override}, + ) + assert not executor.task_queue.empty() + job = executor.task_queue.get_nowait() + executor.task_queue.task_done() + + # The queued job (and its serialized pod_override) must pickle without error. + pickle.dumps(job) + assert isinstance(job.kube_executor_config, dict) + assert job.kube_executor_config["metadata"]["labels"] == {"release": "stable"} + + # And run_next must be able to rebuild a V1Pod from the serialized dict. + rebuilt = pod_generator.PodGenerator.deserialize_model_dict(job.kube_executor_config) + assert rebuilt.metadata.labels == {"release": "stable"} + 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")
