Apologies for the difficulties #5977 caused.

because I'm pretty sure all the partitions are now read from partd before any further computation is done.

I don't think that's the case, at least it wasn't the intention. It's been a while, bu5t my hope was that https://github.com/dask/dask/pull/5977/files#diff-f1e38476dc4a4713fecd130660c4c21dd8122aa3bef55d05c1452a25961f3a86R420-R425 would create a task graph where the later tasks depend on just a small result indicating that the shuffle stage was finished.

@TomAugspurger I could easily be misunderstanding how the task graph works. (First time contributor!)

Looking some more at what _noop returns (it looks like the key of the relevant collect-1 task?), is the idea that the key gets resolved into the partition again later somehow? But if that's supposed to happen after the partd directory has been deleted, where can the partition come from if not somewhere in memory?

More broadly:

How can we cleanup the partd directory midway through a computation without already having read everything from the partd file into memory?

Like it would make sense to me if the cleanup task appeared at the end of the below task graph (from my sort example) after the write_partition tasks. But if the write_partition tasks all indirectly depend on cleanup, doesn't that mean that no writes can happen before the partd file is removed, and therefore that all reads from the partd file are forced to happen before all writes?

Just to add some more info, here is the ordering of tasks for my sort / set_index test script using 2021.2.0:

I believe the shuffle-collect tasks reading from partd are tasks 12 and 14 (while tasks 17 and 22 are the noops). Notice how both read tasks are executed before the write tasks, which are tasks 21 adnd 26 I think. Also note that the cleanup task, which deletes the on disk partd data, is task 16. So the partd directory gets deleted before any writes happen. Which is why I think everything must be in memory.

Here's the ordering of tasks on this branch:

Here, one of the writes happens as task 16, and is then followed by the collect / read from partd step for the second partition as task 17. So only one partition has to be in memory.

Thanks for the additional info! I won't really have time to dig into this, but I'm sufficiently convinced.

Can you try running this a bunch of times locally on this branch to see if you can surfaced the original failure in #5867? I don't recall if I could reproduce it with pytest-repeat, or if it needed a completely clean slate.

cc @jrbourbeau, this could probably go in before the release if @sinclairtarget has a chance to stress test it locally.

The release has been pushed out to next week, so that should take so pressure off.

@TomAugspurger I've run the test that surfaced that original error 100 times with:

#! /bin/bash

set -e

for i in {1..100}
do
    pytest dask/dataframe/tests/test_shuffle.py -k test_rearrange[processes-disk]
done

Pretty low-tech but saw no errors. Let me know if there's anything else I can do.

I also checked out 2.13.0 (which I think was the last release before #5977 went in) and ran the same test 100 times, but sadly didn't see an error.

Ok it seems like we can merge this and keep an eye out for that flaky test. Does that sound reasonable @TomAugspurger?

Thanks @sinclairtarget!