PythonKubernetesOperator and kubernetes taskflow decorator

[c-thiel]

Description

After the implementation of the Docker Taskflow Operator in 2.2, having a simple function on Kubernetes should be quite simple to accomplish.

One major difference I see is how to get the Code into the Pod.
My prefered way here would be to add an init container which takes care of this - much similar on how we use a Sidecar to extract XCom.

I would also prefer to add more defaults than the KubernetesPodOperator does by default to keep the call as simple and straightforward as possible. We could for example default the namespace to the namespace Airflow is running in if it is a Kubernetes Deployment. Also, we could generate a reasonable pod name from DAG-ID and Task ID.

Use case/motivation

Beeing able to run a task as simple as:

@task.kubernetes(image='my_image:1.1.0')
def sum(a, b):
    from my_package_in_version_1 import complex_sum
    return complex_sum(a,b)

would be awesome!
It would cleanly separate environments and resources between tasks just as the Docker or Virtualenv Operators do - and unlinke the @task.python decorator. We could specify the resources and environment of each task freely by specifying a resources Argument.

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[uranusjr]

cc @dimberman

[subkanthi]

Can I try this please if its not being worked on.

[potiuk]

Feel fre3 @subkanthi !

[Asciotti]

I just found out about TaskFlow recently while looking through some docs, it looks fantastic! It reminds me of what Prefect is doing and providing an easier interface to convert code -> running in dags.

We heavily use kubernetes-based operators, and would love to see this feature add.

It there a PR I can track on this @subkanthi ?

[subkanthi]

Not yet, will have it in a week hopefully.

[amin-rashidi-sp]

Thanks for doing this @subkanthi, any update on the progress?

[c-thiel]

Hello @subkanthi,
I just looked through your pull request and have a couple of remarks. Also from my side thanks for picking this up!!!

1. Regarding the default namespace - should we at least for in-cluster deployments default to the namespace of the current pod? This is very helpful if Airflow is run on kubernetes as it will rarely run in the default namespace. We can determine the namspace from "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
2. Script Path
   Currently the script is always been written to /tmp/script.py. I think it would be nice if this was a parameter.
3. Pod Name
   Should we maybe try to get the dag_id and task_id in there?
   I am currently using:

def generate_pod_name(task_id, dag_id=None):
    """Given a task_id and a dag_id, get a kubernetes pod name"""
    dag_id = dag_id if dag_id is not None else "xxx-unspecified"
    name = f"airflow-task-{_to_k8s_name(dag_id)}-{_to_k8s_name(task_id)}"
    name = re.sub(r"[^a-z0-9.-]+", "-", name.lower())
    name = name[:62]
    if name.endswith("-"):
        name = name[:-1]
    return name

Which is also not perfect I guess. The KubernetesPodOperator itself then adds another suffix.
Edit: The problem with the DAG ID is that it might be bound to the task a little too late. xxx-unspecified is of course a bad thing to do ;)

4. do_xcom_push
   Similar to the VirtualEnvOperator (I think?), we could determine a sensible default using:

def contains_explicit_return(f):
    return any(isinstance(node, ast.Return) for node in ast.walk(ast.parse(f)))

Let me know what you think!

[hterik]

I would advice you to look into the Airlflow Kubernetes Scheduler. If you already use Kubernetes as your main airflow scheduler you can simply use taskflow and have python functions running in kubernetes as default. If you don't use the scheduler but want to have some tasks running in kubernetes anyway, there is still a lot of useful functions to reuse or take inspiration from in the scheduler code. For example washing dag ids to be compatible to kubernetes and adding them as labels and annotations on the Pod.

See construct_pod function in airflow/kubernetes/pod_generator.py

[c-thiel]

@hterik I also think the Scheduler bit can be good inspiration.
I just want to address your first point briefly: I really do think that there are many use-cases for this: Having an operator allows much more freedoms than using the KubernetesExecutor only - for example in specifying the exact dependencies you need for individual tasks. Additionaly, spawning pods for every little thing really might stress the cluster. In our deployment we are runnig ~20 thousand tasks each morning in an hour and then a few more over the day. We tried the KubernetesExecutor but having to schedule a new pod for every sensor and mundane task was a bit too much for our control plane. Thus we are using Celery with lots of KubernetesOperator with self implemented taskflow API.

[hterik]

Did you see the CeleryKubernetes executor ? That allows mixing K8 with celery

Can also recommend ["executor_config"]["pod_override"] on taskflow api if you have need to specify exact dependencies per task. We have custom wrappers on top of this so you just have to write:

@node_type=HighPerformance
@task
def foo():
   ....

Actually it is a lot more granular than this but you get the idea.

For small lightweight sensors there is possibility to batch them into https://airflow.apache.org/docs/apache-airflow/stable/concepts/smart-sensors.html or https://airflow.apache.org/docs/apache-airflow/stable/concepts/deferring.html

Not against your idea of a custom operator, just handing out tips :)