Adjust rmm pool handling in PDSH benchmarks

[TomAugspurger]

Description

This updates the handling of the RMM Pool argument for the PDSH CLI. Previously, --rmm-pool-size was used for both sync and async RMM memory resources.

As described in rapidsai/rmm#2060, the meaning of the pool size is different for these two types of memory resources, and rapidsai/dask-cuda#1563 is deprecating passing a pool size to dask_cuda.LocalCUDACluster with rmm_async=True.

No real difference in the timings (maybe a small improvement, but more likely just noise):

# New default
python python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py --path /datasets/toaugspurger/tpch/scale-100/ --no-print-results  --executor streaming --scheduler distributed --iterations 3 --rmm-async 1,2,3,4,5
# Total mean time across all queries: 45.6239 seconds

# Previous default
python python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py --path /datasets/toaugspurger/tpch/scale-100/ --no-print-results  --executor streaming --scheduler distributed --iterations 3 --rmm-async --rmm-pool-size 0.5 1,2,3,4,5
# Total mean time across all queries: 47.8799 seconds

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[rjzamora]

Thanks Tom! Any reason this is marked as Draft? Seems like this should go in before rapidsai/dask-cuda#1563, correct?

[bdice]

This looks fine to me -- can you help advise on #20129?

[TomAugspurger]

/merge