Partial rechunks within P2P

[hendrikmakait]

Closes #8326
Blocked by #8207

This PR currently splits on entirely independent subset of the rechunking. A more intricate (yet better) implementation would create independent shuffles even for cases where an input would go into more than one shuffle.

This PR splits any outputs that allows us to cull more inputs, i.e., two outputs along an axis are separated if they do not end in the same input chunk.

• Tests added / passed
• Passes pre-commit run --all-files

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       27 files  ±  0         27 suites  ±0   11h 57m 44s ⏱️ + 21m 4s
  3 950 tests +  2    3 838 ✔️ +  4     109 💤  -   1  3 ❌  - 1 
49 683 runs  +42  47 392 ✔️ +64  2 288 💤  - 19  3 ❌  - 3 

For more details on these failures, see this check.

Results for commit 1f554ef. ± Comparison against base commit 9273186.

This pull request removes 5 and adds 7 tests. Note that renamed tests count towards both.

distributed.shuffle.tests.test_rechunk ‑ test_rechunk_with_single_output_chunk_raises
distributed.shuffle.tests.test_rechunk ‑ test_worker_for_homogeneous_distribution[1]
distributed.shuffle.tests.test_rechunk ‑ test_worker_for_homogeneous_distribution[2]
distributed.shuffle.tests.test_rechunk ‑ test_worker_for_homogeneous_distribution[41]
distributed.shuffle.tests.test_rechunk ‑ test_worker_for_homogeneous_distribution[50]

distributed.shuffle.tests.test_rechunk ‑ test_cull_p2p_rechunk_independent_partitions
distributed.shuffle.tests.test_rechunk ‑ test_cull_p2p_rechunk_overlapping_partitions
distributed.shuffle.tests.test_rechunk ‑ test_partial_rechunk_homogeneous_distribution
distributed.shuffle.tests.test_rechunk ‑ test_pick_worker_homogeneous_distribution[1]
distributed.shuffle.tests.test_rechunk ‑ test_pick_worker_homogeneous_distribution[2]
distributed.shuffle.tests.test_rechunk ‑ test_pick_worker_homogeneous_distribution[41]
distributed.shuffle.tests.test_rechunk ‑ test_pick_worker_homogeneous_distribution[50]

♻️ This comment has been updated with latest results.

[hendrikmakait]

There's some stuff we still calculate for the individual n-dimensional partials (scaling as a product) but that could be computed for the partials per axis (scaling as a sum). I'll take a another look at some quick refactoring here but this might very well be premature optimization.

[hendrikmakait]

@crusaderky: Looking at my old benchmark results, you are in fact reproducing them. However, test_tiles_to_rows has significantly improved in performance since then:

[hendrikmakait]

Searching through the commit history, this is likely because #8207 has been merged, which this branch also included at the time.

[crusaderky]

The side-by-side dashboard makes it pretty obvious why there's such a dramatic slowdown.

test_adjacent_groups[0.3-128 MiB-p2p-disk], 180 input/output chunks.
This PR picks even workers only for the output chunks! (worker 0 is always at the bottom)

peek.mp4

[hendrikmakait]

Quick summary of an offline discussion I had with @crusaderky: In the example above, each partial rechunk has five output tasks (180 inputs / 36 barriers). As a result, the deterministic worker assignment logic will always pick the same five workers to store those outputs which underutilizes the cluster. We should be able to assess the impact of better worker assignment by randomizing it.

[crusaderky]

Adding a naive randomization of the workers (d9116fb) fixes most problems:

• there are runtime regressions in some tests with 8 MiB chunks. I'm oriented towards accepting them for the greater good.
• there is a runtime regression in test_swap_axes[1-128 MiB-p2p-disk]. It makes no sense to me as it should not be impacted by this PR. Investigating.
• the peak memory regressions are false positives, please ignore: A/B plot is misleading for p2p-disk memory usage and tick measures coiled/benchmarks#1240
• worker randomization breaks TPCH query 7: p2p worker selection is brittle for TPCH query 7 #8424

[hendrikmakait]

Regarding the regression of test_swap_axes, I've seen that happen before when running benchmarks for this PR. In the past, they disappeared if I test that one in isolation. I'm not sure what's going on exactly.

[crusaderky]

Regarding the regression of test_swap_axes, I've seen that happen before when running benchmarks for this PR. In the past, they disappeared if I test that one in isolation. I'm not sure what's going on exactly.

Reran it in isolation. You're right, these results don't make sense for me

[crusaderky]

Waiting for an alternative suggestion to the worker randomization from @fjetter, as discussed offline.
This PR is otherwise green to go for me.

[hendrikmakait]

I think what @fjetter and I both have in mind would look like this: hendrikmakait#14

FWIW, I'd not block on that, a follow-up with another benchmark run is trivial.

[crusaderky]

Running a final A/B test. Results out tomorrow morning.

[crusaderky]

Looks good. Comment is the same as in the previous posts.

[fjetter]

I think what @fjetter and I both have in mind would look like this: hendrikmakait#14

essentially, yes. I would've not shifted by one but by len(workers) but the gist is the same. I think shifting by len(workers) generates more homogeneous distributions for small numbers of shuffles

[hendrikmakait]

essentially, yes. I would've not shifted by one but by len(workers) but the gist is the same. I think shifting by len(workers) generates more homogeneous distributions for small numbers of shuffles

I'm confused, wouldn't shifting by len(workers) have no effect whatsoever in a cluster without fluctuating workers?

The problem the shift solves is that static range partitioning if len(partitions) < len(workers) skips some workers (e.g., only picks every other worker if len(partitions) / len(workers) = 0.5). By shifting by one, this should be as homogeneous as possible for any range-partitioned pattern as long as len(partitions) remains constant. For non-constant len(partitions), I can't come up with anything else that's better unless we start counting assignments per worker.

[hendrikmakait]

Note that we could potentially investigate switching assignment to round-robin and then shift by the len(partitions) % len(workers). IIRC, the lack of buffering in our disk buffers currently prohibits that though.

[crusaderky]

🥳