Skip to content

Simplify memory-resource handling in Dask integration#172

Merged
rapids-bot[bot] merged 2 commits intorapidsai:branch-25.06from
rjzamora:use-existing-mr-dask
Mar 26, 2025
Merged

Simplify memory-resource handling in Dask integration#172
rapids-bot[bot] merged 2 commits intorapidsai:branch-25.06from
rjzamora:use-existing-mr-dask

Conversation

@rjzamora
Copy link
Member

@rjzamora rjzamora commented Mar 26, 2025

This is a follow-up to #150

Closes #157

The goal of this PR is to simplify memory-resource creation by avoiding it in rapidsmp altogether. Since the user can just use LocalCUDACluster to deploy a Dask cluster, they can also use existing options/utilities to create a memory pool on each worker. When rapidsmp.integrations.dask.bootstrap_dask_cluster is called, each worker only needs to wrap the current memory resource in a StatisticsResourceAdaptor.

This is technically "breaking", because it removes the pool_size argument from bootstrap_dask_cluster. However, we are only using that option in rapidsai/cudf#18335 (which is still experimental - and can be easily changed).

@rjzamora rjzamora added breaking Introduces a breaking change improvement Improves an existing functionality labels Mar 26, 2025
@rjzamora rjzamora self-assigned this Mar 26, 2025
@rjzamora rjzamora requested a review from a team as a code owner March 26, 2025 19:07
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: The current memory resource doesn't need to be a pool for spilling to work.

@TomAugspurger
Copy link
Contributor

I think this closes #157.

Does this require that the dask-cuda worker be created with any special arguments? I think the answer is no: rmm.mr.get_current_device_resource will always return something, regardless of whether the user created the LocalCUDACluster with something like rmm_pool_size.

@rjzamora rjzamora changed the title remove pool creation, and use existing mr on each process Simplify memory-resource handling in Dask integration Mar 26, 2025
@rjzamora
Copy link
Member Author

Does this require that the dask-cuda worker be created with any special arguments? I think the answer is no: rmm.mr.get_current_device_resource will always return something, regardless of whether the user created the LocalCUDACluster with something like rmm_pool_size.

Right - It doesn't require the user to pass in an rmm_pool_size argument, but they should do this if they want good performance.

@TomAugspurger
Copy link
Contributor

In that case, it might be good to add that note to our quickstart example: https://github.com/rapidsai/rapids-multi-gpu/blob/branch-25.06/docs/source/quickstart.md#dask-cudf-example.

@rjzamora
Copy link
Member Author

In that case, it might be good to add that note to our quickstart example

Good suggestion. I tweaked the code and added a brief comment.

Copy link
Member

@pentschev pentschev left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, thanks @rjzamora .

@rjzamora
Copy link
Member Author

/merge

@rapids-bot rapids-bot bot merged commit b1ca99b into rapidsai:branch-25.06 Mar 26, 2025
25 checks passed
@rjzamora rjzamora deleted the use-existing-mr-dask branch March 26, 2025 22:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking Introduces a breaking change improvement Improves an existing functionality

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Let rapidsmp use the RMM pools from a dask_cuda.LocalCUDACluster

3 participants

Comments