Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[POC][DNM] Experimental to/from_worker_storage API #1299

Draft
wants to merge 66 commits into
base: branch-24.02
Choose a base branch
from

Conversation

rjzamora
Copy link
Member

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

raydouglass and others added 30 commits October 17, 2019 15:03
* Add ucx-py dependency to CI
[REVIEW] Backport TLS files missing fix
raydouglass and others added 27 commits June 7, 2022 11:44
Copy link

copy-pr-bot bot commented Dec 13, 2023

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.

@github-actions github-actions bot added the python python code needed label Dec 13, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
python python code needed
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants