cc @randerzander @ayushdg @VibhuJawa

Thanks @madsbk . This is fun to see.

Two questions:

1. Can I ask you to compare against https://docs.dask.org/en/latest/futures.html#submit-tasks-from-tasks ? When are times when we would to use the approach here instead?
2. It looks like you have found a 400x speedup on some workload. Can I ask you to share more about that workload, what was slow before, and how this helps so much?

1. Can I ask you to compare against https://docs.dask.org/en/latest/futures.html#submit-tasks-from-tasks ? When are times when we would to use the approach here instead?

Sorry, I should have made that more clear. Using get_client().submit() and secede() works great in many cases but in cases of shuffle, this approach requires extra memory because the inputs to the shuffle jobs must be in-memory until the jobs completes. We can handle this problem in the regular shuffle implementation by prioritizing the shuffle-split tasks and in this PR by setting the "priority" of the newly inserted tasks.

2. It looks like you have found a 400x speedup on some workload. Can I ask you to share more about that workload, what was slow before, and how this helps so much?

This is the regular shuffle implementation without staging where max_branch == npartitions and it creates n**2 number of tasks!

this approach requires extra memory because the inputs to the shuffle jobs must be in-memory until the jobs completes

Can you expand on this? What I think you're saying here is that it's hard for a task to create new tasks and then return. That the creating task needs to remain active until its sub-tasks finish. Is that right? If so, does fire_and_forget get you past this?

This is the regular shuffle implementation without staging where max_branch == npartitions and it creates n**2 number of tasks!

Ah, that makes sense. Thanks for the extra context.

I am comparing it here because the new shuffle also doesn't use staging at the moment. I am working on a staging version.

Hrm, but then you still have n**2 tasks, right? It's nice that you can create those tasks on the fly, but fundamentally this still seems like it'll quickly become a bottleneck as you increase the number of partitions. I may not understand how your new shuffle algorithm works though.

Can you expand on this? What I think you're saying here is that it's hard for a task to create new tasks and then return. That the creating task needs to remain active until its sub-tasks finish. Is that right? If so, does fire_and_forget get you past this?

What I am doing in this PR is like fire_and_forget() but with dependency support.

Hrm, but then you still have n**2 tasks, right? It's nice that you can create those tasks on the fly, but fundamentally this still seems like it'll quickly become a bottleneck as you increase the number of partitions. I may not understand how your new shuffle algorithm works though.

True, but using insert_tasks() we don't have to create all tasks at once, we can organize the communication as we see fit.
Staging is properly preferable when partition sizes are small because it limits the splitting to max_branchs but with large enough partitions I expect non-staging to be preferable.

What I am doing in this PR is like fire_and_forget() but with dependency support.

Sorry to keep asking questions, but I'm not sure what this means.

It might help if you were to provide some very simple usage examples of what kinds of problems this solves. Right now I understand that you're focused on dataframe shuffles, but if we want to change the scheduler we will probably need to write more general documentation around a feature like this. This would also help reviewers to understand the motivations of the work more generally.

I enjoyed reading your docstring (thanks for including it in the header comment). I think that the applications mentioned there would be good examples to bring up, both in docs and tests. I am especially curious where the current tasks-from-tasks approach falls short and how this approach can make things cleaner.

Once I know more about what you're doing here we can bring in the right other people who use the tasks-from-tasks approach to get their feedback.

True, but using insert_tasks() we don't have to create all tasks at once, we can organize the communication as we see fit.
Staging is properly preferable when partition sizes are small because it limits the splitting to max_branchs but with large enough partitions I expect non-staging to be preferable.

OK, it seems like there are two simultaneous conversations happening here.

1. Dataframe shuffles
2. Tasks inserting new tasks.

In the interest of other readers I suggest that we move the dataframe shuffles conversation to another issue (I suspect that such an issue exists somewhere already) and lets focus on the tasks-inserting-tasks feature here.

Hrm, after a brief search I couldn't find a good general "dataframe shuffles are slow at scale" issue. I've written one up here. Hopefully we can centralize conversation there: dask/dask#6314

In order to explain this PR, I will walk through a shuffle example that hopefully makes it clear how insert_tasks() works.

Initially, we create a task graph that, after fusion, looks like (npartitions=2):

Instead of having the getitem and concat tasks in the graph from the start, we have rearguard_dyshuffle and dynshuffle.

Now, when one of the dynshuffle tasks is executed, it will first call shuffle_group() and then call insert_tasks() to create the getitem and concat tasks:

Similarly, when the other dynshuffle task is executed, the graph becomes:

At this point, the graph is identical with the original shuffle approach, which makes is possible to re-calculate the whole calculation and thus support resilient.

I am still working on getting some performance numbers on large datasets, which is where this PR will makes a significant difference hopefully :)

With regards to accelerating dataframe shuffle operations, I think that we should consider how this would be used to batch together many local partitions when splitting. If there is a clean way of doing that that allowed for orders of magnitude scale increases and roughly O(n) tasks then I think that that would be very exciting.

I often come back to this issue when thinking about how to track dependencies for dynamic tasks. I don't exactly need resilience; I am more interested in propagating information about all tasks executed during a computation.

The more ergonomic way that would happen is to have some hook that would be triggered with each '.compute()' call. Then I would be able to maintain metadata on TaskState tracking these dynamic internal dependencies.

Internal is the key word here; by task dependencies dask means inputs, while for my use case I need to consider keys for which compute was called as dependencies.

insert_tasks wouldn't suit me, since I do want to modify the current task by adding instances of internal dependencies, ie dependencies on compute calls.

There is a possible issue with resilience here, since re-execution may result in a simplified graph, or in a resource depletion as all internal dependencies are computed at once. This doesn't apply in my use case, since resilience is not a concern.

I hope this can fit here, as I am interested/using tasks within tasks. I will make a separate issue for my use case once I have something readable.