[FEA] Expose rmm maximum_pool_size to LocalCUDACluster and dask-cuda-worker API

[VibhuJawa]

We should expose rmm's maximum_pool_size argument (See docs) to the LocalCUDACluster and dask-cuda-worker CLI API .

Why
By default, the total available memory on the GPU is used by RMM . This can cause problems for workflows where it actually grows to the total available device memory and we need some memory outside the POOL .

Why we may need room for other allocations:

1. Competing Process: Common case is Client & Worker on the same GPU.

2. Competing pool/libraries: Some libraries might need to allocate some memory outside of the pool.
   Like Pytorch (Even Nccl/Raft might need some room.)

Giving the user the ability to set maximum_pool_size easily can circumvent those issues.

Workflow Context:
I ran into this while working on a workflow where the pool grew to 32501MiB (out of cards total 32510MiB ) and it left very little memory for cuML to do some non POOL allocation leading to failure.

Current WorkAround:
Using the RMM reinitialize API.

cluster = LocalCUDACluster(
                          jit_unspill=True,
                          rmm_pool_size=None,
                          local_directory='/raid3/vjawa/')
                          
client = Client(cluster)
## Set pool
_ = client.run(rmm.reinitialize,pool_allocator=True,initial_pool_size=28*(2**30), maximum_pool_size=(28)*(2**30));

[VibhuJawa]

Happy to take this on if we decide to do this.

[mmccarty]

Would this allow me to run multiple workers per GPU?

[pentschev]

Exposing this as rmm_maximum_pool_size/--rmm-maximum-pool-size sounds good to me @VibhuJawa . I believe the only thing we have to be careful about is to make this parameter a no-op if --rmm-pool-size isn't specified as that's what we use today to decide whether to create an RMM pool at all or not. Let me know if you want to work on it, otherwise I can do that too, it shouldn't take too long.

[pentschev]

Would this allow me to run multiple workers per GPU?

@mmccarty I don't quite get the direct relation between this and having multiple workers per GPU. Today you can do that manually but I don't believe that would necessarily beneficial performance-wise. What's your use case for multiple workers per GPU?

[VibhuJawa]

Exposing this as rmm_maximum_pool_size/--rmm-maximum-pool-size sounds good to me @VibhuJawa . I believe the only thing we have to be careful about is to make this parameter a no-op if --rmm-pool-size isn't specified as that's what we use today to decide whether to create an RMM pool at all or not. Let me know if you want to work on it, otherwise I can do that too, it shouldn't take too long.

Thanks @pentschev . Will push a PR keeping your suggestion in mind soon.

Would like to contribute to dask-cuda to just know the code base slightly better.

[mmccarty]

@pentschev Happy to talk about the use case. Is there a better issue? or should I just create a new one?

[pentschev]

@mmccarty I think you should open a new issue, it may be covered somewhere but I can't tell for sure without more details.

[jakirkham]

There was issue ( #571 ). Though yeah hard to tell whether that is related without knowing more about the use case.