diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index ee42f103fb3eb..ab8f636cb059b 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -588,7 +588,7 @@ class TaskDecoratorCollection: :param name: Name of the pod to run. This will be used (plus a random suffix if *random_name_suffix* is *True*) to generate a pod ID (DNS-1123 subdomain, containing only ``[a-z0-9.-]``). Defaults to - ``k8s_airflow_pod_{RANDOM_UUID}``. + ``k8s-airflow-pod-{python_callable.__name__}``. :param random_name_suffix: If *True*, will generate a random suffix. :param arguments: arguments of the entrypoint. (templated) The docker image's CMD is used if this is not provided. diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes.py index c35b492dcabe0..323101449dcaa 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes.py @@ -19,7 +19,6 @@ import base64 import os import pickle -import uuid from collections.abc import Sequence from shlex import quote from tempfile import TemporaryDirectory @@ -68,9 +67,15 @@ class _KubernetesDecoratedOperator(DecoratedOperator, KubernetesPodOperator): def __init__(self, namespace: str | None = None, use_dill: bool = False, **kwargs) -> None: self.use_dill = use_dill + + # If the name was not provided, we generate operator name from the python_callable + # we also instruct operator to add a random suffix to avoid collisions by default + op_name = kwargs.pop("name", f"k8s-airflow-pod-{kwargs['python_callable'].__name__}") + random_name_suffix = kwargs.pop("random_name_suffix", True) super().__init__( namespace=namespace, - name=kwargs.pop("name", f"k8s_airflow_pod_{uuid.uuid4().hex}"), + name=op_name, + random_name_suffix=random_name_suffix, cmds=["placeholder-command"], **kwargs, ) diff --git a/providers/cncf/kubernetes/tests/provider_tests/cncf/kubernetes/decorators/test_kubernetes.py b/providers/cncf/kubernetes/tests/provider_tests/cncf/kubernetes/decorators/test_kubernetes.py index 8f412c81b7f3c..528bce1780e2a 100644 --- a/providers/cncf/kubernetes/tests/provider_tests/cncf/kubernetes/decorators/test_kubernetes.py +++ b/providers/cncf/kubernetes/tests/provider_tests/cncf/kubernetes/decorators/test_kubernetes.py @@ -70,6 +70,16 @@ def mock_hook(): return mock.patch(HOOK_CLASS).start() +# Without this patch each time pod manager would try to extract logs from the pod +# and log an error about it's inability to get containers for the log +# {pod_manager.py:572} ERROR - Could not retrieve containers for the pod: ... +@pytest.fixture(autouse=True) +def mock_fetch_logs() -> mock.Mock: + f = mock.patch(f"{POD_MANAGER_CLASS}.fetch_requested_container_logs").start() + f.return_value = "logs" + return f + + def test_basic_kubernetes(dag_maker, session, mock_create_pod: mock.Mock, mock_hook: mock.Mock) -> None: with dag_maker(session=session) as dag: @@ -215,3 +225,104 @@ def f(): assert len(dag.task_group.children) == 1 teardown_task = dag.task_group.children["f"] assert teardown_task.is_teardown + + +@pytest.mark.parametrize( + "name", + ["no_name_in_args", None, "test_task_name"], + ids=["no_name_in_args", "name_set_to_None", "with_name"], +) +@pytest.mark.parametrize( + "random_name_suffix", + [True, False], + ids=["rand_suffix", "no_rand_suffix"], +) +def test_pod_naming( + dag_maker, + session, + mock_create_pod: mock.Mock, + name: str | None, + random_name_suffix: bool, +) -> None: + """ + Idea behind this test is to check naming conventions are respected in various + decorator arguments combinations scenarios. + + @task.kubernetes differs from KubernetesPodOperator in a way that it distinguishes + between no name argument was provided and name was set to None. + In the first case, the operator name is generated from the python_callable name, + in the second case default KubernetesPodOperator behavior is preserved. + """ + extra_kwargs = {"name": name} + if name == "no_name_in_args": + extra_kwargs.pop("name") + + with dag_maker(session=session) as dag: + + @task.kubernetes( + image="python:3.10-slim-buster", + in_cluster=False, + cluster_context="default", + config_file="/tmp/fake_file", + random_name_suffix=random_name_suffix, + namespace="default", + **extra_kwargs, # type: ignore + ) + def task_function_name(): + return 42 + + task_function_name() + + dr = dag_maker.create_dagrun() + (ti,) = dr.task_instances + session.add(ti) + session.commit() + + task_id = "task_function_name" + op = dag.get_task(task_id) + if name is not None: + assert isinstance(op.name, str) + + # If name was explicitly set to None, we expect the operator name to be None + if name is None: + assert op.name is None + # If name was not provided in decorator, it would be generated: + # f"k8s-airflow-pod-{python_callable.__name__}" + elif name == "no_name_in_args": + assert op.name == f"k8s-airflow-pod-{task_id}" + # Otherwise, we expect the name to be exactly the same as provided + else: + assert op.name == name + + op.execute(context=ti.get_template_context(session=session)) + pod_meta = mock_create_pod.call_args.kwargs["pod"].metadata + assert isinstance(pod_meta.name, str) + + # After execution pod names should not contain underscores + task_id_normalized = task_id.replace("_", "-") + + def check_op_name(name_arg: str | None) -> str: + if name_arg is None: + assert op.name is None + return task_id_normalized + + assert isinstance(op.name, str) + if name_arg == "no_name_in_args": + generated_name = f"k8s-airflow-pod-{task_id_normalized}" + assert op.name == generated_name + return generated_name + + normalized_name = name_arg.replace("_", "-") + assert op.name == normalized_name + + return normalized_name + + def check_pod_name(name_base: str): + if random_name_suffix: + assert pod_meta.name.startswith(f"{name_base}") + assert pod_meta.name != name_base + else: + assert pod_meta.name == name_base + + pod_name = check_op_name(name) + check_pod_name(pod_name)