Skip to content

Commit c2283f7

Browse files
Adjust rmm pool handling in PDSH benchmarks (#20138)
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 ``` Authors: - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Richard (Rick) Zamora (https://github.com/rjzamora) - Bradley Dice (https://github.com/bdice) URL: #20138
1 parent b7e7da8 commit c2283f7

File tree

1 file changed

+11
-3
lines changed
  • python/cudf_polars/cudf_polars/experimental/benchmarks

1 file changed

+11
-3
lines changed

python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -669,11 +669,12 @@ def parse_args(
669669
)
670670
parser.add_argument(
671671
"--rmm-pool-size",
672-
default=0.5,
672+
default=None,
673673
type=float,
674674
help=textwrap.dedent("""\
675675
Fraction of total GPU memory to allocate for RMM pool.
676-
Default: 0.5 (50%% of GPU memory)"""),
676+
Default: 0.5 (50%% of GPU memory) when --no-rmm-async,
677+
None when --rmm-async"""),
677678
)
678679
parser.add_argument(
679680
"--rmm-release-threshold",
@@ -777,7 +778,14 @@ def parse_args(
777778
default=False,
778779
help="Enable statistics planning.",
779780
)
780-
return parser.parse_args(args)
781+
782+
parsed_args = parser.parse_args(args)
783+
784+
if parsed_args.rmm_pool_size is None and not parsed_args.rmm_async:
785+
# The default rmm pool size depends on the rmm_async flag
786+
parsed_args.rmm_pool_size = 0.5
787+
788+
return parsed_args
781789

782790

783791
def run_polars(

0 commit comments

Comments
 (0)