More flexible dynamic task mapping

[Tehada]

Description

Current implementation of dynamic task mapping as of airflow 2.3.0 is lacking some features to support different workflows. Particularly, it lacks "zip-like" functionality to dynamically spawn tasks. Consider this example (concept, not an actual code):

@task
def get_params():
    res = {
        "arg1": [1, 3],
        "arg2": [2, 4]
    }
    return res

@task
def run_task(arg1, arg2):
    print(arg1, arg2)

with DAG(dag_id="airflow_fun", start_date=datetime(2022, 5, 16)) as dag:
    p = get_params()
    run_task.expand_zip(arg1=p["arg1"], arg2=p["arg2"])

The desired behavior for this code is to run run_task twice: run_task(arg1=1, arg2=2) and run_task(arg1=3, arg2=4). It is not possible to implement such workflow in current implementation of dynamic tasks using expand API -- I want such API to appear in airflow.

It is possible, that creating something similar to current expand is not the best approach and may be there is another fundamentally different way, which will be simple to use and implement -- I am going to make some research of exising tools with this feature (prefect and dagster coming in mind first, luigi seems to be dead).

Also, @ashb noted in slack:

The first thing to work out and discuss is what is the syntax to the DAG author?

Which leads me to this plan (may be updated later):

1. Gather existing suggestions in this issue for further comparison.
2. Research goob and bad practices of this API in different tools.
3. Find out usecases and workflows, which will be not covered by suggested API (it is desirable to cover all usecases).

This plan and description are open for any feedback, suggestions, criticism.

Use case/motivation

One real-life usecase is to have a task get_params, which, when triggered, will list gcs bucket for new files for yesterday (number of files is variable day-to-day) and generate arguments for KubernetesPodOperator to run processing task in pod. But the processing is quite RAM-hungry, so I also predict RAM usage of particular processing task based on the file size in bucket. Then I want to run processing for each file in a separate pod with dynamically generated cmds, arguments and resources for each pod. With "zip-like" dynamic mapping this usecase will be trivially implementable.

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

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template!

[uranusjr]

If I understand correctly, this is totally possible with the current syntax to achieve what you want, except you need to zip and unpack manually.

@task
def zip_params(params):
    return zip(params["arg1"], params["arg2"])

@task
def run_task(param):
    arg1, arg2 = param
    print(arg1, arg2)

with DAG(dag_id="airflow_fun", start_date=datetime(2022, 5, 16)):
    p = zip_params(get_params())
    run_task.expand(param=p)

I do agree it would be useful to provide methods for common functional operations, however, and perhaps even support user-registered custom operators.

The syntax I have in mind is something like

import operator

with DAG(...):
    params = get_params().map(operator.itemgetter("arg1", "arg2")).zip()
    run_task.expand(arg1=params[0], arg2=params[1])

[Tehada]

@uranusjr, indeed, I slightly adjusted your example and it works fine. So the problem happens, when we want to use operators (BashOperator, KubernetesPodOperator, etc) -- they just don't seem to have a necessary API for full compatibility with taskflow API.

[uranusjr]

The main problem of using zipping against a classic operator class is not really about the zip operation itself, but to unpack the zipped iterator (i.e. the arg1, arg2 = param part). Even if you can produce a zip iterator (which you already can), it’s still difficult to use the iterator against a classic operator. I am yet to be able to come up with a reasonable API either; instantiating a class instance is (kind of) fundamentally not straightforward in a unctional programming context.

Note that you can also already kind of do this like below. There’re a lot of caveats on calling execute directly, but it’s kind of doable. (Not that I particularly recommend it though.)

@task
def run_task(param):
    image, arguments = param
    op = KubernetesPodOperator(task_id="my-task", image=image, arguments=arguments, ...)
    return op.execute()

[potiuk]

Also - another thing I usually recommend in this case - maybe this is not really a big problem that you cannot use it in classsic Airflow Operators. I think - personally - we should move away from classic operators and if you really want to use the "modern" way of writing DAGs, simply stop using them.

Pretty much all the "classic" operators are usually thin wrappers over the Hooks. This has always been the case, but with the taskflow api, this is much more "viable" than ever to use this pattern:

@task
def run_task(param):
    .... # params here
    hook = MyHook(connection=connection)
    hook.run_hook_methods_here()

Arguably - this is not as easy with KubernetesPodOperator - because it is pretty "thick" comparing to other operators and has a lot of logic for pod templates generation and such - so temporarily, the proposal by @uranusjr might make sense, but with some recent changes that are coming - i.e. #20578 and #21077 (they will come in 2.4 most likely) even KPO will move to the "taskflow" world.

Which might actaully mean that we do not have to do anything with classic operators here.

[ashb]

One downside to not using (classic) operators but just using PythonOperator (decoratored or not) i sht we loose much of the ability to see what the task is doing from a lineage PoV. (Both Lineage itself, and the task type in the UI is not useful anymore)

The OpenLineage folks are working on an idea to add the ability see what hooks are used in operators to the "Listener" interface, so that part might be solved soon.

[ashb]

@Tehada Please check out (and leave feedback) on the proposed interface in #24021

[uranusjr]

Now we have map and zip implemented (to be released in 2.4), do we have anything left to do in this issue?

[github-actions]

This issue has been automatically marked as stale because it has been open for 30 days with no response from the author. It will be closed in next 7 days if no further activity occurs from the issue author.

[github-actions]

This issue has been closed because it has not received response from the issue author.