Apply memory backpressure to the shuffle

[madsbk]

Depend on #23833.

The shuffle now takes its device memory from reserve_memory() on both sides, the four ShuffleManager.Inserter methods going in and extract_chunk coming out, so a request that cannot be satisfied queues alongside the other actors' and is served by priority. Before this both let the C++ side reserve and spill internally. cudf-polars already reserves this way for scans in actor_graph/io.py.

Breaking change

The insert methods and extract_chunk are coroutines now, and LocalRepartitioner._iter_chunks is an async generator. That updates twenty-one call sites across groupby.py, over.py, sort.py, shuffle.py and the shuffler tests.

Notes

Reserving inside the insert methods rather than at the call sites keeps the accounting in one place. That matters for insert_hash_keys and insert_index, whose Python-side reorder allocates a full table copy that nothing reserved before. They reserve partition_and_pack_cost(), which covers a reorder plus a pack, consume the reorder's share once it lands, and hand the remainder to split_and_pack().

_iter_chunks reserves per piece, so the TODO there about batching pieces up to target_partition_size would cut the reservation traffic as well as the unpack overhead.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[madsbk]

/ok to test