Bugfix: Move rendering of map_index_template so it renders for failed tasks as long as it was defined before the point of failure

[TJaniF]

In Airflow 2.9 map_index_template does not render when the task fails.

I moved the rendering into the finally of _execute_callable so it always happens.

TaskFlow:

Traditional operator:

Also attempted 2 unit tests :)

Apologies if there is already a PR addressing this. I did not see one.

cc: @RNHTTR

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[RNHTTR]

Nice! Thanks for cleaning up my mess :)

[uranusjr]

Need to fix tests

[uranusjr]

This should fix things…

[TJaniF]

@uranusjr thank you!

[rightx2]

I wonder what I've experienced is same type of the problem above:

1. clear dynnamic mapped tasks in a dag run
2. all dynamic mapped tasks go into state of running
3. Force to set dynmaic task box(or dag run) state success (mark state as success)

In this case, all map_index values are integer, instead of given names.

[TJaniF]

@rightx2 That is interesting. I just tried that and for me the map index values from the first run "stick", but maybe I misunderstood what you are doing :)

This is what I tried:

This is the DAG:

from airflow.decorators import dag, task

@dag(
    start_date=None,
    schedule=None,
    catchup=False,
)
def mapping_test():
    @task(
        map_index_template="{{ custom_map_index }}"
    )
    def add_one(num):
        import time
        
        if num > 10:
            time.sleep(10)
        else:
            time.sleep(2)
        
        from airflow.operators.python import get_current_context

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

        return num + 1
    
    add_one.expand(num=[1, 2, 3, 4, 5, 10, 11, 12, 13])


mapping_test()

I'm using Airflow 2.9.1 (Astro Runtime 11.3.0)

[rightx2]

@TJaniF

I think there was a mistake in my experiment. I believe I can clarify it now.

1. For whole newly added dynamic tasks:

   • When I start running dyanmic tasks for the first time and stop them immediately (marking them as success or failed, before they finish), all map_index values are integers.

2. For newly added individual tasks within pre-existing dynamic tasks:

   • Similarly, if I clear the tasks this time and stop them, only the map_index of the newly added individual task is an integer.

[TJaniF]

@rightx2 Thanks for the added explanation. I understand what you mean now and can reproduce it, marking as success (or as failed) will lead to an integer index even if the code to define the custom map index has already run. Thanks for flagging this!

@RNHTTR I think it would be nice if it worked for marking tasks the same way as for failed tasks with attempting to render the custom map index template even if interrupted. :)
If you agree I can open an issue and/or take a stab at it, though I can't promise that will happen soonish, am working on some other things rn 😅