Fix KubernetesExecutor scheduler crash when pod_override is queued in-cluster

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py

@@ -36,6 +36,7 @@
 from typing import TYPE_CHECKING, Any
 
 from deprecated import deprecated
+from kubernetes.client import Configuration
 from kubernetes.dynamic import DynamicClient
 from sqlalchemy import select
 
@@ -73,6 +74,30 @@
     )
 
 
+def _reset_local_vars_configuration(obj: Any) -> None:
+    """
+    Reset ``local_vars_configuration`` to a fresh ``Configuration()`` on every kubernetes model object.
+
+    kubernetes-python-client v36.0.0 changed model constructors to call ``Configuration.get_default_copy()``
+    instead of ``Configuration()``, so objects created after incluster config was loaded capture the global
+    config whose ``refresh_api_key_hook`` is an unpicklable local closure function. We fall back to the
+    pre-v36 behaviour which is picklable. A fresh ``Configuration()`` keeps ``client_side_validation`` so
+    the worker-side ``reconcile_pods`` setters still work; ``None`` would break them.
+    """
+    if obj is None or isinstance(obj, (str, int, float, bool)):
+        return
+    if hasattr(obj, "openapi_types"):
+        obj.local_vars_configuration = Configuration()
+        for attr in obj.openapi_types:
+            _reset_local_vars_configuration(getattr(obj, attr, None))
+    elif isinstance(obj, dict):
+        for v in obj.values():
+            _reset_local_vars_configuration(v)
+    elif isinstance(obj, (list, tuple)):
+        for item in obj:
+            _reset_local_vars_configuration(item)
+
+
 class KubernetesExecutor(BaseExecutor):
     """Executor for Kubernetes."""
 
@@ -226,6 +251,9 @@ def execute_async(
             pod_template_file = executor_config.get("pod_template_file", None)
         else:
             pod_template_file = None
+
+        _reset_local_vars_configuration(kube_executor_config)
+
         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

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import pickle
 import random
 import re
 import string
@@ -25,7 +26,7 @@
 
 import pytest
 import yaml
-from kubernetes.client import models as k8s
+from kubernetes.client import Configuration, models as k8s
 from kubernetes.client.rest import ApiException
 from sqlalchemy import inspect
 from urllib3 import HTTPResponse
@@ -36,6 +37,7 @@
 from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import (
     KubernetesExecutor,
     PodReconciliationError,
+    _reset_local_vars_configuration,
 )
 from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types import (
     ADOPTED,
@@ -915,6 +917,123 @@ def test_pod_template_file_override_in_executor_config(
             finally:
                 executor.end()
 
+    @staticmethod
+    def _install_unpicklable_incluster_hook(pod):
+        """Simulate the in-cluster kubernetes v36 default config: an unpicklable local-closure hook
+        on every nested model object's ``Configuration`` (``InClusterConfigLoader._set_config.<locals>.
+        _refresh_api_key``)."""
+
+        def _make_hook():
+            def _refresh_api_key(config):
+                return None
+
+            return _refresh_api_key
+
+        for obj in (
+            pod,
+            pod.metadata,
+            pod.spec,
+            pod.spec.containers[0],
+            pod.spec.containers[0].resources,
+        ):
+            cfg = Configuration()
+            cfg.refresh_api_key_hook = _make_hook()
+            obj.local_vars_configuration = cfg
+
+    def test_reset_local_vars_configuration_strips_hook_recursively(self):
+        """``_reset_local_vars_configuration`` must replace every nested config with a hook-free one."""
+        # Build a pod with the unpicklable in-cluster hook on every nested config; confirm it is
+        # unpicklable first, then that the reset makes it picklable.
+        pod = k8s.V1Pod(
+            metadata=k8s.V1ObjectMeta(),
+            spec=k8s.V1PodSpec(
+                containers=[k8s.V1Container(name="base", resources=k8s.V1ResourceRequirements())]
+            ),
+        )
+        self._install_unpicklable_incluster_hook(pod)
+        with pytest.raises((pickle.PicklingError, AttributeError, TypeError)):
+            pickle.dumps(pod)
+
+        _reset_local_vars_configuration(pod)
+
+        for obj in (pod, pod.metadata, pod.spec, pod.spec.containers[0], pod.spec.containers[0].resources):
+            assert obj.local_vars_configuration.refresh_api_key_hook is None
+            assert obj.local_vars_configuration.client_side_validation is True
+        # now picklable
+        pickle.dumps(pod)
+
+    @pytest.mark.execution_timeout(30)
+    @pytest.mark.skipif(
+        AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+    )
+    @mock.patch(
+        "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.AirflowKubernetesScheduler.run_pod_async"
+    )
+    @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+    def test_execute_async_pod_override_picklable_and_reconcilable(
+        self, mock_get_kube_client, mock_run_pod_async, data_file
+    ):
+        """Regression: a ``pod_override`` carrying an in-cluster ``Configuration`` must be both picklable
+        and reconcilable.
+
+        kubernetes-python-client v36 model constructors capture the global in-cluster ``Configuration``,
+        whose ``refresh_api_key_hook`` is a local closure -> pickling the ``KubernetesJob`` onto the
+        multiprocessing queue crashes the scheduler. ``execute_async`` resets ``local_vars_configuration``
+        to a fresh ``Configuration()`` so the job pickles, while keeping ``client_side_validation`` so the
+        worker-side ``run_next``/``reconcile_pods`` still works (nulling it instead would raise
+        ``'NoneType' object has no attribute 'client_side_validation'``).
+        """
+        pod_override = k8s.V1Pod(
+            metadata=k8s.V1ObjectMeta(labels={"release": "stable"}),
+            spec=k8s.V1PodSpec(
+                containers=[
+                    k8s.V1Container(
+                        name="base",
+                        image="airflow:3.6",
+                        resources=k8s.V1ResourceRequirements(requests={"cpu": "100m", "memory": "384Mi"}),
+                    )
+                ]
+            ),
+        )
+        self._install_unpicklable_incluster_hook(pod_override)
+
+        # Sanity: the raw pod is unpicklable (reproduces the scheduler crash pre-fix).
+        with pytest.raises((pickle.PicklingError, AttributeError, TypeError)):
+            pickle.dumps(pod_override)
+
+        executor_template_file = data_file("executor/basic_template.yaml")
+        mock_get_kube_client.return_value = mock.patch("kubernetes.client.CoreV1Api", autospec=True)
+
+        with conf_vars({("kubernetes_executor", "pod_template_file"): None}):
+            executor = self.kubernetes_executor
+            executor.start()
+            try:
+                executor.execute_async(
+                    key=TaskInstanceKey("dag", "task", "run_id", 1),
+                    queue=None,
+                    command=["airflow", "tasks", "run", "true", "some_parameter"],
+                    executor_config={
+                        "pod_template_file": executor_template_file,
+                        "pod_override": pod_override,
+                    },
+                )
+
+                assert not executor.task_queue.empty()
+                job = executor.task_queue.get_nowait()
+                executor.task_queue.task_done()
+
+                # 1) The queued job pickles cleanly (this is the actual scheduler crash path).
+                pickle.dumps(job)
+                # 2) Contract preserved: kube_executor_config stays a V1Pod.
+                assert isinstance(job.kube_executor_config, k8s.V1Pod)
+                # 3) Worker-side reconcile/construct_pod succeeds (the null-config approach fails here).
+                self.kubernetes_executor.kube_scheduler.run_next(job)
+                mock_run_pod_async.assert_called_once()
+                built = mock_run_pod_async.call_args[0][0]
+                assert built.spec.containers[0].resources.requests == {"cpu": "100m", "memory": "384Mi"}
+            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")