Fix pre-mature evaluation of tasks in mapped task group

[avkirilishin]

closes: #34023

Redo of the #34337
Relates to #35541 (comment)

@uranusjr Can you please take a look at the approach? I haven't checked the performance yet, but if the approach is okay, I will check it if needed. Perhaps it can be optimized.

And I found that something goes wrong with the example from the get_relevant_upstream_map_indexes method. When I checked the dag:

importpendulumfromairflowimportDAGfromairflow.decoratorsimporttask, task_groupwithDAG(
'mul-in-tg',
schedule='@daily',
start_date=pendulum.DateTime(2023, 12, 26),
) asdag:
@taskdefupstream(inp):
returninp@taskdefthis_task(v): # This is self.task.returnv*2@task_groupdeftg1(inp):
val=upstream(inp) # This is the upstream task.this_task(val) # When inp is 1, val here should resolve to 2.returnvalval=tg1.expand(inp=[1, 2, 3]) # This val is the same object returned by tg1.@taskdefanother_task(inp, val):
print("(inp, val)", (inp, val))
@task_groupdeftg2(inp):
another_task(inp, val) # val here should resolve to [2, 4, 6].tg2.expand(inp=["a", "b"])

I observed:

Is this the expected behavior?

[potiuk]

cc: @uranusjr - following discussion in #35541 - do you think that one should fix the problem ?

[eladkal]

@avkirilishin can you rebase and resolve conflics?

[ephraimbuddy]

Can you add this test:

def test_mapped_tasks_in_mapped_task_group_waits_for_upstreams_to_complete(dag_maker, session):
"""Test that one failed trigger rule works well in mapped task group"""
with dag_maker() as dag:
@dag.task
def t1():
return [1, 2, 3]
@task_group("tg1")
def tg1(a):
@dag.task()
def t2(a):
return a
@dag.task(trigger_rule=TriggerRule.ONE_FAILED)
def t3(a):
return a
t2(a) >> t3(a)
t = t1()
tg1.expand(a=t)
dr = dag_maker.create_dagrun()
ti = dr.get_task_instance(task_id="t1")
ti.run()
dr.task_instance_scheduling_decisions()
ti3 = dr.get_task_instance(task_id="tg1.t3")
assert not ti3.state

[uranusjr]

Hmm, I’d prefer we don’t call this function for this case at all instead (i.e. try to ensure the task is expanded).

Also this check is not correct. If ti_count is exactly 1 if the task is mapped (against a list of length 1).

[uranusjr]

Can we do this outside of the blocks instead of repeating it everywhere? This is both cumbersome and error-prone.

[eladkal]

@avkirilishin are you still working on this PR?

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.