[POC][DNM] Experimental to/from_worker_storage API

[rjzamora]

I ran into a few road blocks while exploring "p2p" shuffling with RAPIDS, and ultimately decided it was worthwhile to explore a dedicated API for persisting/loading a DataFrame collection to/from distributed local storage.

Note on Shuffling

In addition to adding a new to_worker_storage/from_worker_storage API, this PR also adds an explicit-comms mechanism to shuffle a DataFrame collection directly to distributed/local disk. This shuffle_to_parquet function shuffles and writes the data to storage in rounds, so that you can process significantly larger-than-memory data. The performance of this new function is not good compared to in-memory shuffling, but it is certainly much faster than "p2p".

This API can not be embedded within a lazy dask.dataframe query like the "tasks" or "p2p" shuffle can. However, this API can be used to explicitly compose workflows that are known to require large shuffle operations (e.g. NeMo Data Curator). E.g.

from dask_cuda.explicit_comms.dataframe.utils import to_worker_storage, from_worker_storage

ddf = ...

local_path = "/raid/persisted_ddf"
to_worker_storage(ddf, local_path, shuffle_on="id")
shuffled_ddf = from_worker_storage(local_path)

TODO:

• Significant cleanup needed (this is very rough)
• Better design description/discussion needed
• Better API description/discussion needed
• Significant testing needed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.