using `expand` inside @task_group will cause a `airflow.exceptions.DuplicateTaskIdFound`

[EdwardRadical]

Apache Airflow version

2.3.0 (latest released)

What happened

The unique ID of the task is seemingly not calculated correctly in base.py and it can cause in a duplicate key error, my guess is that this bit from /decorators/base.py:115:

if tg_task_id not in dag.task_ids:
    return task_id

doesn't make any sense and the actual return should be:

if tg_task_id not in dag.task_ids:
    return tg_task_id

But I didn't write it so there is maybe a reason down the line for return task_id. Someone that does know his way around TaskFlow's codebase should pitch in.

What you think should happen instead

task.expand should not lead to a airflow.exceptions.DuplicateTaskIdFound inside task groups.

How to reproduce

Pick example_task_group_decorator.py, change it so that task_1 returns a list, and call task_2 with .expand(value=task_1).

A modification as simple as this will already kill the DAG:

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""Example DAG demonstrating the usage of the @taskgroup decorator."""

import pendulum

from airflow.decorators import task, task_group
from airflow.models.dag import DAG


# [START howto_task_group_decorator]
# Creating Tasks
@task
def task_start():
    """Empty Task which is First Task of Dag"""
    return '[Task_start]'


@task
def task_1(value: int) -> list[str]:
    """Empty Task1"""
    return [f'[ Task1 {value} ]', f'[ Task1 {value+1} ]']


@task
def task_2(value: list[str]) -> str:
    """Empty Task2"""
    return f'[ Task2 {value} ]'


@task
def task_3(value: str) -> None:
    """Empty Task3"""
    print(f'[ Task3 {value} ]')


@task
def task_end() -> None:
    """Empty Task which is Last Task of Dag"""
    print('[ Task_End  ]')


# Creating TaskGroups
@task_group
def task_group_function(value: int) -> None:
    """TaskGroup for grouping related Tasks"""
    task_3(task_2.expand(value=task_1(value)))


# Executing Tasks and TaskGroups
with DAG(
    dag_id="example_task_group_decorator",
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
) as dag:
    start_task = task_start()
    end_task = task_end()
    for i in range(5):
        current_task_group = task_group_function(i)
        start_task >> current_task_group >> end_task

# [END howto_task_group_decorator]

Operating System

Manjaro Linux

Versions of Apache Airflow Providers

No response

Deployment

Virtualenv installation

Deployment details

No response

Anything else

No response

Are you willing to submit 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!

[tirkarthi]

39ea872 looks like a relevant commit . Also #12312 (comment)

cc: @turbaszek

[EdwardRadical]

I looked a bit into it.

Seems to me the result _get_unique_task_id was intended to be enriched by eventual task groups and whatnot, like here:

        if task_group:
            self.task_id = task_group.child_id(task_id)

But in the case of expand and MappedOperator it was not done. I think this was meant to be done as BaseOperator refrains to add the task to the group in case self.__from_mapped is set, but MappedOperator just doesn't care, adds everything to task group and dag without altering the task_id and the whole thing thankfully blows up at self.dag.add_task

        if self.task_group:
            self.task_group.add(self)
        if self.dag:
            self.dag.add_task(self)

It is notable that MappedOperator actually doesn't seem to care about logically separating the task mappings using the map_index, so as far as airflow knows they are perfect copies of the same task instance, hence, at the minimum attempt of nesting a mapped task somewhere, it goes haywire.

I would like to say that I have a decent grip over it but this screams subtle conceptual assumptions and I wouldn't be surprised if this wasn't the only gotcha in there, I wouldn't touch this with a five meters pole, this issue needs some expert knowledge.

On another note, checking for a collision in a dictionary for a key and then returning another key altogether as it was checked and deemed a valid unique ID, like done here:

        if tg_task_id not in dag.task_ids:
            return task_id

Is so much of a code smell that I almost thought it was a typo. Either the function returns an unique ID and it can be called _get_unique_task_id or it doesn't and you don't name it like that, but rather _get_task_id_stem.

[uranusjr]

It is notable that MappedOperator actually doesn't seem to care about logically separating the task mappings using the map_index, so as far as airflow knows they are perfect copies of the same task instance, hence, at the minimum attempt of nesting a mapped task somewhere, it goes haywire.

An instance of a task and a task instance are two different concepts in Airflow (it’s super confusing I know). A task is a thing in a DAG (think about the boxes you see in the Graph View), while a task instance is something that gets executed inside a worker. A mapped task spawns multiple task instances to run in the worker, but those task instances belong to the same logical task, which doesn’t have a notion of map index, but only know whether itself is being mapped (MappedOperator) or not.

In your DAG attached above, each loop iteration creates a new task from task_2 that is a mapped task and can be expanded into multiple task instances. Task instances expanded from the same task do share an identical task, but that’s not the issue here; the issue is the DAG needs to differentiate different tasks created from task_2, via some mechanism that is not currently implemented correctly.

[EdwardRadical]

An instance of a task and a task instance are two different concepts in Airflow (it’s super confusing I know). A task is a thing in a DAG (think about the boxes you see in the Graph View), while a task instance is something that gets executed inside a worker. A mapped task spawns multiple task instances to run in the worker, but those task instances belong to the same logical task, which doesn’t have a notion of map index, but only know whether itself is being mapped (MappedOperator) or not.

In your DAG attached above, each loop iteration creates a new task from task_2 that is a mapped task and can be expanded into multiple task instances. Task instances expanded from the same task do share an identical task, but that’s not the issue here; the issue is the DAG needs to differentiate different tasks created from task_2, via some mechanism that is not currently implemented correctly.

I might have explained myself incorrectly and mixed up terms: I know that tasks and task instances are separate concepts, what I meant to say is that if the mapped task instances are all coming from the same task definition, either the task is mapping itself to several instances, each uniquely identified and separable, or the task is mapping itself to a set of mapped tasks that are uniquely identified and separable, which will in turn generate one task instance.

What happens now is neither: at the point in which expand kicks in, where the task is already conceptually an operator, the operators generated by expand do not differ in the slightest from the original operator, in fact the whole implementation of MappedOperator looks like something that has been considerably hindered by behavior defined in the base classes that should have been addressed there to allow MappedOperator to work as it should, which is one more reason why expand does not work as one would think it would (if expand is used after expand it does not expand once for each mapped instance, rather it's a 1:1 mapping over the first set of mapped instances, you want each mapped instance to generate n mapped instances? too bad, not doable).

[uranusjr]

I think this is hitting the late binding Python gotcha? #21278

Not entirely sure, it could be related to how tasks are actually generated internally.

[EdwardRadical]

I don't see how it relates to be honest.

The issue here seems to be that _get_unique_task_id, return the base, unchanged, task_id each time expand is called as the task is part of a task group:

        if tg_task_id not in dag.task_ids:
            return task_id

This normally works because task groups do overwrite this task_id later by prefixing the task group identifier before adding it to the DAG:

        if task_group:
            self.task_id = task_group.child_id(task_id)

expand is also using this _get_unique_task_id as all operators do when instantiated, but since this happens at runtime self.task_id = task_group.child_id(task_id) is never executed and the task_id is never overwritten with an unique task_id by the task group.

This doesn't seem like a case of lazy initialization gone wrong, this is simply an actual bug.

[EdwardRadical]

At the end of the day it seems that

        if task_group:
            self.task_id = task_group.child_id(task_id)

Is only executed by inheritors of BaseOperator. MappedOperator does not inherit from it so it never does execute the line that would prefix correctly the task_id, hence the duplicate key error thrown later, when self.dag.add_task is called. Maybe a simple fix would be to add those two lines to MappedOperator too, although there are other places that fiddle with the task_id.