Airflow Dynamic Task Mapping identifiers in UI

Question

With the new Dynamic Task Mapping feature, it's possible to use the expand function and create a variable number of tasks based on the output of a previous task.

The problem is that, as shown in the documentation, the UI will only show the total number of tasks generated, this is particularly uncomfortable when the number of task instances is high and we need to debug a specific error.

I tried searching the docs and GitHub pages but all I could find is a suggestion. Is there anyone who found a solution or a workaround to fix this?

Answer 1

In Airflow 2.9 or above you can use map_index_template variable in your task mapping providing context of your task

@task(
    # optionally, you can set a custom index to display in the UI (Airflow 2.9+)
    map_index_template="{{ my_custom_map_index }}"
)
def add(x: int, y: int):

    # get the current context and define the custom map index variable
    from airflow.operators.python import get_current_context

    context = get_current_context()
    context["my_custom_map_index"] = "Input x=" + str(x)

    return x + y

added_values = add.partial(y=10).expand(x=[1, 2, 3])

How it appears in ui

Reference: https://www.astronomer.io/docs/learn/dynamic-tasks#dynamic-task-concepts