Disk-spilling shuffler

[madsbk]

Shuffle extraction is often the OOM point. It acts as a barrier that waits on all data before continuing, so at that point a rank holds its entire local share of the shuffle. In RapidsMPF the normal answer is to spill device memory to host, but on very large datasets host memory is limiting too, and there is no tier below it.

We propose an option for the existing Shuffler to spill shuffled data to disk, and to unspill it back to host or device only when the user extracts it.

Proposal

Add a DiskBuffer, a file-backed handle to a single spilled chunk. Disk becomes a tier the shuffler can move a chunk into, and it stays entirely outside the MemoryType taxonomy.

Concretely, rapidsmpf::shuffler::detail::Chunk gains a disk-resident state, ReceivedChunks::spill gets a disk fallback instead of a throw, and Shuffler::extract rehydrates. Nothing in BufferResource, Buffer, MemoryType, or the general spill machinery changes.

Host limit

The Shuffler takes a host limit, a byte count or percentage mirroring the existing spill_device_limit. Device pressure spills device to host exactly as today, and crossing the host limit is what triggers host to disk. Disk is therefore entered on a threshold the user sets, rather than only once host allocation has already failed.

Why not MemoryType::DISK

A general MemoryType::DISK is something we may well want later, but it is a much bigger change than the enum makes it look, and most of the cost lands downstream.

Within RapidsMPF, Buffer cannot represent disk. Its storage is a std::variant of a device and a host buffer (buffer.hpp:370), and data(), write_access(), and exclusive_data_access() all return a raw pointer that a file-backed buffer does not have. Every cross-tier move funnels through buffer_copy() (buffer.cpp:134), one cuda_memcpy_async between two pointers against a single reservation, so disk cannot join that funnel without turning it into a tier dispatch that affects every BufferResource user. On top of that, once DISK is in MEMORY_TYPES it becomes reachable by accident from the sites that pass the whole array as a preference list.

Downstream is the larger cost. cudf_streaming has its own exhaustive switch over rapidsmpf::MemoryType with a default: RAPIDSMPF_FAIL, reached on the spill path, and it sizes an array by rapidsmpf::MEMORY_TYPES.size().

[madsbk]

Remember, @nirandaperera, all of this work should stay an implementation detail of the shuffler. We do not want any of it leaking outside the Shuffler for now.

It is fine to introduce new data structures such as DiskBuffer if needed, but they should remain private to the shuffler.

Only once we have evaluated the design and performance of the disk shuffler should we take the next step and consider library-wide disk spilling.

[nirandaperera]

I think I am happy with the status of bench_shuffle now. All the disk-shuffler work is in #1197. Benchmark is semantically similar to libcudf streaming bench_shuffle. Main difference is no partitioning/ concatenation using cudf (purely with PackedData).
TL;DR : Now we can shuffle with device memory limited to input data size (and without any host memory).

• Disk spilling/ unspilling in the shuffler is limited to device buffers (Disk buffer impl is generic though, can store/ restore from any ptr. But enabling this throughout spilling paths require many more changes as outlined by @madsbk here

The disk shuffler work introduces the following env vars,

• RAPIDSMPF_SPILL_HOST_LIMIT - An upper limit to all host buffer allocations (via BR, doesnt limit new, malloc etc)
• RAPIDSMPF_DISK_SPILL_DIR - Spilling directory. BR resource will use RAPIDSMPF_DISK_SPILL_DIR/<proc_id> as the spill dir (this disambiguates spill paths for parallel processes).

What I have tested (in umb B200 machine with 8 ranks):

• Shuffling without host memory (extreme case)

  • Disable host entirely -> buffers will reside in DEVICE/ DISK.
  • Limit RAPIDSMPF_SPILL_DEVICE_LIMIT to total input batch size (nranks * output_partitions_per_rank * payload_size * nbatches). Once bench inputs are generated, device memory is exhausted.
  • This relies on the fact that, local receives are immediately spillable. Hence at least payload_size becomes available for receive buffers.

• Shuffling with device memory for only a single input batch (more extreme case)

  • Limit host to total input batch size (nranks * output_partitions_per_rank * payload_size * nbatches ) & device to a single batch size (nranks * output_partitions_per_rank * payload_size_) --> Bench inputs are now spread across both device and host memory. Host is fully exhausted.
  • This is a lower limit than the previous case, but this lets spilling into both host and disk.

To get this to work, I had to make some changes to Metadata-Payload-exchange (defer output buffer allocation when reservations are unavailable, etc). Maybe @pentschev & @wence- can take a look at it.

The benchmark PR may not reviewable in its current state. But its stacked as follows

1. Add host memory limits Add memory limiting capability to host memory resource #1193
2. Add a KvikIO-backed DiskResource for blocking disk spill I/O Add a KvikIO-backed DiskResource for blocking disk spill I/O #1186
3. Add a file-backed DiskBuffer Add a file-backed DiskBuffer #1195
4. Add shuffler disk residency for received chunks Add shuffler disk residency for received chunks #1196
5. Bench PR Add a pre-partitioned shuffle benchmark #1197