Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,15 @@ def execute_async(

try:
kube_executor_config = PodGenerator.from_obj(executor_config)
if kube_executor_config is not None:
# Round-trip through a fresh Configuration rather than queuing the pod_override
# as constructed by user code. In-cluster, kubernetes-client 36.x's default
# Configuration carries a refresh_api_key_hook local closure that pickle cannot
# serialize; queuing this object as-is crashes the scheduler (not just the task)
# when task_queue.put() pickles it. See PodGenerator.deserialize_model_dict.
kube_executor_config = PodGenerator.deserialize_model_dict(
PodGenerator.serialize_pod(kube_executor_config)
)
except Exception:
self.log.error("Invalid executor_config for %s. Executor_config: %s", key, executor_config)
self.fail(key=key, info="Invalid executor_config passed")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2367,6 +2367,65 @@ def test_execute_async_retains_job_spec(self, mock_get_kube_client, mock_kuberne
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_pod_override_is_picklable_in_cluster(
self, mock_get_kube_client, mock_kubernetes_job_watcher
):
"""A pod_override built under an in-cluster Configuration must still be picklable.

kubernetes-client 36.x's default Configuration carries a ``refresh_api_key_hook``
local closure when running in-cluster (set by ``InClusterConfigLoader``). Any
``V1Pod``/nested model built without an explicit ``local_vars_configuration``
inherits that closure via ``Configuration.get_default_copy()`` and becomes
unpicklable -- crashing the scheduler (not just the task) when execute_async's
``task_queue.put()`` tries to pickle it.
"""
import pickle

from kubernetes.client.configuration import Configuration

def _unpicklable_refresh_hook(client_configuration):
pass

original_default = Configuration.get_default_copy()
bad_config = Configuration()
bad_config.refresh_api_key_hook = _unpicklable_refresh_hook
Configuration.set_default(bad_config)
try:
pod_override = k8s.V1Pod(
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
resources=k8s.V1ResourceRequirements(requests={"cpu": "100m", "memory": "384Mi"}),
)
]
),
)
# Sanity check the reproduction: an unprotected pod_override must not be
# picklable under this Configuration, otherwise this test proves nothing.
with pytest.raises(AttributeError, match="local object"):
pickle.dumps(pod_override)

executor = self.kubernetes_executor
executor.start()
try:
key = TaskInstanceKey("dag", "task", "run_id", 1, -1)
executor.execute_async(
key=key,
queue=None,
command=["airflow", "tasks", "run", "true", "some_parameter"],
executor_config={"pod_override": pod_override},
)
stored_config = executor.pod_launch_attempts[key].job.kube_executor_config
# Must not raise: this is what task_queue.put() does internally.
pickle.dumps(stored_config)
finally:
executor.end()
finally:
Configuration.set_default(original_default)

@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")
Expand Down