Added taskgroup decorator

[VBhojawala]

Added taskgroup decorator which can be used to create taskgroup from python callable. Inside python callable tasks can be grouped by calling tasks (python callable ) created with task decorator.

• taskgroup decorator takes both optional and keyword argument which are passed to Constructor of TaskGroup class.

• TaskGroup can be created with one of following syntax

@taskgroup
@taskgroup()
@taskgroup(group_id='group_name')

• TaskGroup class constructor takes one mandatory argument group_id, if not given in decorator it sets group_id to python callable name.

Following is a simple example demonstrating use of taskgroup decorator grouping multiple tasks.

@task
def task_1(value):
  return f'[ Task1 {value} ]'


@task
def task_2(value):
  print(f'[ Task2 {value} ]')


@taskgroup
def section_1(value):
  return task_2(task_1(value))

taskgroup decorator utilizes existing TaskGroup context manager currently used for creating TaskGroup, which means we can create nested taskgroup by created nested callable. Following is an example demonstrating use of nested taskgroup.

@task
def task_start():
    return '[Task_start]'

@task
def task_end():
    print(f'[ Task_End ]')

@task
def task_1(value):
    return f'[ Task1 {value} ]'

@task
def task_2(value):
    print(f'[ Task2 {value} ]')

@task
def task_3(value):
    return f'[ Task3 {value} ]'

@task
def task_4(value):
    print(f'[ Task4 {value} ]')

@taskgroup
def section_1(value):

    @taskgroup
    def section_2(value2):
        return task_4(task_3(value2))

    op1 = task_2(task_1(value))
    return section_2(op1) 

Dedicated test cases for taskgroup decorator is created in file
/tests/utils/test_task_group_decorator.py

Recent changes

• Added logic to append suffix to duplicate group_id.

closes: #11870

[VBhojawala]

Hi @casassg @yuqian90 @ashb,

I have raised new pull request with suggested changes. Kindly review it.

[casassg]

this looks great! Thanks for taking your time to contribute this!

a small nit: Could you add an extra test that calls a taskgroup decorated function 2 times.

@taskgroup
def training(paths):
  validate(train(paths))

@dag
def pipeline():
  training('some/path')
  training('some/other/path')

[github-actions]

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

[github-actions]

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

[casassg]

ty! This looks good to me. Will need a committer for a final review

[kaxil]

Can you rebase the PR on latest master please

[VBhojawala]

Hi @kaxil, I have rebased the PR.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest master at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[github-actions]

The Workflow run is cancelling this PR. Building images for the PR has failed. Follow the the workflow link to check the reason.

[github-actions]

The Workflow run is cancelling this PR. Building images for the PR has failed. Follow the the workflow link to check the reason.

[ashb]

I am not sure this is the right behavour!

This is separate task groups with the same id. That feels like a logic bug on the DAG author part and should throw an error.

What do others think of this?

[casassg]

I agree. But we can't retrieve the function for previously called decorated functions, which means that we can't know wether the function name is the same or not (aka wether its the same function being called 2 times or 2 functions w the same group_id). So while I agree this is an unintended side effect, I would leave it as is as it allows us to instatiate the same function several times

[ashb]

We can't check what task group decorator created the previous function, but we could check if the current TG was created by the current decorated function, and if not then it would fail.

We could keep a set of TG objects in the wrapper fn, and if the TG already exists on the DAG, and the TG is not in that set, then it was created by some other decorator -> raise an error.

[casassg]

Makes sense. I don't think it's a big deal either way. If users do this, even though it's weird and can be confusing, they can do it.

Mostly want to make sure we are valuing the right trade offs between complexity added and value of having this extra check.

Not strongly opinionated though

[ashb]

My thought was that this is more likely a bug than desired behaviour -- for instance someone copy-and-pasted a task group but forgot to update the name

[dimberman]

I agree with @casassg here. It would be very hard to enforce this. I'm especially imagining situations where users import taskgroups from libraries and therefore do not have control over the task_group_id. That said, it would make sense for us to allow users to override the task_group_id via a kwarg in case they import two functions with the same id name.

[dimberman]

might be nicer to call it "@task_group" instead of "@TaskGroup" since there's no capitalization

[dimberman]

Can you unwind this? I wound't want to encourage people to nest like this.

[dimberman]

Hi @VBhojawala are you interested in continuing this PR? I'm glad to either finish it for you or help you finish :)

[ashb]

Hi @dimberman ,

Kindly review below PR and help me to finish it :
#13405

related: #8970

@VBhojawala commenting on another unrelated PR asking for a review is considered poor etiquette - please don't do this.

[ashb]

Closing in favour of #15034 which continues this work.