celery: call a task after all subtasks of all subtasks have run

Question

I have two main tasks: main_A and main_B.

extract_new spawns a large number of sub tasks, which itself spawns the subsub1 and subsub2 tasks.

I want to run main_B after all subsubA and subsubB tasks have run. In this diagram, tasks depend on the tasks above:

                  main_A
                  /    \
                 /      \
                /        \
               /          \
              sub          sub
              /\           / \
             /  \         /   \
          sub1   sub2    sub1  sub2
            \     |      /      /
             \    |     /      /
              \   |    /      /
               \  |   /      /

                  main_B

How can I do that? If I set main_B as the callback for main_A, it will be run juste after main_A has spawned the sub tasks. If I run sub tasks as a chord with main_B as a callback, main_B will be run before the subsub1 and subsub2 have finished.

Both of those are not options, since main_B depends on data created by the sub, subsub1 and subsub2 tasks.

Answer 1

If you require full non-blocking behaviour then you can refactor your code so that instead of spawning sub-tasks your tasks simply build a canvas object and pass it back up the call chain

my_chain = mainA.s((params,...)) | mainA_subtasks.s(params) | main_B(params)

@app.task
def mainA(params):
    # in here put the main work done by mainA

def mainA_subtasks(params):

    my_group = []
    for i in some_loop():
        tasks = get_sub_tasks(params, i)
        my_group.extend(tasks)

    return group(my_group)

# and so on

Answer 2

you should be able to use AsyncResult.collect() to manage this - see docs

@app.task
def runner(x, y, z):
    result = main_a.apply_async((x, y,z))
    result.collect()

    main_B.apply_async()

you'll need to check for errors etc also