[REVIEW]Expose rmm-maximum_pool_size argument

dask_cuda/cli/dask_cuda_worker.py

@@ -73,6 +73,17 @@
     .. note::
         This size is a per-worker configuration, and not cluster-wide.""",
 )
+@click.option(
+    "--rmm-maximum-pool-size",
+    default=None,
+    help="""When ``--rmm-pool-size`` is specified, this argument indicates the maximum pool size.
+        Can be an integer (bytes), string (like ``"5GB"`` or ``"5000M"``) or ``None``.
+        By default, the total available memory on the GPU is used.
+        When ``--rmm_pool_size`` is not specified, this argument is ignored if provided.
+
+        .. note::
+            This size is a per-worker configuration, and not cluster-wide.""",
+)
 @click.option(
     "--rmm-managed-memory/--no-rmm-managed-memory",
     default=False,
@@ -277,6 +288,7 @@ def main(
     memory_limit,
     device_memory_limit,
     rmm_pool_size,
+    rmm_maximum_pool_size,
     rmm_managed_memory,
     rmm_async,
     rmm_log_directory,
@@ -327,6 +339,7 @@ def main(
         memory_limit,
         device_memory_limit,
         rmm_pool_size,
+        rmm_maximum_pool_size,
         rmm_managed_memory,
         rmm_async,
         rmm_log_directory,

dask_cuda/cuda_worker.py

@@ -58,6 +58,7 @@ def __init__(
         memory_limit="auto",
         device_memory_limit="auto",
         rmm_pool_size=None,
+        rmm_maximum_pool_size=None,
         rmm_managed_memory=False,
         rmm_async=False,
         rmm_log_directory=None,
@@ -152,6 +153,9 @@ def del_pid_file():
                 )
             if rmm_pool_size is not None:
                 rmm_pool_size = parse_bytes(rmm_pool_size)
+                if rmm_maximum_pool_size is not None:
+                    rmm_maximum_pool_size = parse_bytes(rmm_maximum_pool_size)
+
         else:
             if enable_nvlink:
                 warnings.warn(
@@ -239,7 +243,11 @@ def del_pid_file():
                         get_cpu_affinity(nvml_device_index(i, cuda_visible_devices(i)))
                     ),
                     RMMSetup(
-                        rmm_pool_size, rmm_managed_memory, rmm_async, rmm_log_directory,
+                        rmm_pool_size,
+                        rmm_maximum_pool_size,
+                        rmm_managed_memory,
+                        rmm_async,
+                        rmm_log_directory,
                     ),
                 },
                 name=name if nprocs == 1 or name is None else str(name) + "-" + str(i),

dask_cuda/local_cuda_cluster.py

@@ -119,6 +119,16 @@ class LocalCUDACluster(LocalCluster):
 
         .. note::
             This size is a per-worker configuration, and not cluster-wide.
+
+    rmm_maximum_pool_size : int, str or None, default None
+        When ``rmm_pool_size`` is set, this argument indicates
+        the maximum pool size.
+        Can be an integer (bytes), string (like ``"5GB"`` or ``"5000M"``) or ``None``.
+        By default, the total available memory on the GPU is used.
+        When ``rmm_pool_size`` is ``None``, this argument is ignored if provided.
+        .. note::
+            This size is a per-worker configuration, and not cluster-wide.
+
     rmm_managed_memory : bool, default False
         Initialize each worker with RMM and set it to use managed memory. If disabled,
         RMM may still be used by specifying ``rmm_pool_size``.
@@ -196,6 +206,7 @@ def __init__(
         enable_rdmacm=None,
         ucx_net_devices=None,
         rmm_pool_size=None,
+        rmm_maximum_pool_size=None,
         rmm_managed_memory=False,
         rmm_async=False,
         rmm_log_directory=None,
@@ -230,6 +241,7 @@ def __init__(
         )
 
         self.rmm_pool_size = rmm_pool_size
+        self.rmm_maximum_pool_size = rmm_maximum_pool_size
         self.rmm_managed_memory = rmm_managed_memory
         self.rmm_async = rmm_async
         if rmm_pool_size is not None or rmm_managed_memory:
@@ -248,6 +260,8 @@ def __init__(
                 )
             if self.rmm_pool_size is not None:
                 self.rmm_pool_size = parse_bytes(self.rmm_pool_size)
+                if self.rmm_maximum_pool_size is not None:
+                    self.rmm_maximum_pool_size = parse_bytes(self.rmm_maximum_pool_size)
         else:
             if enable_nvlink:
                 warnings.warn(
@@ -397,6 +411,7 @@ def new_worker_spec(self):
                     ),
                     RMMSetup(
                         self.rmm_pool_size,
+                        self.rmm_maximum_pool_size,
                         self.rmm_managed_memory,
                         self.rmm_async,
                         self.rmm_log_directory,

dask_cuda/utils.py

@@ -45,8 +45,16 @@ def setup(self, worker=None):
 
 
 class RMMSetup:
-    def __init__(self, nbytes, managed_memory, async_alloc, log_directory):
-        self.nbytes = nbytes
+    def __init__(
+        self,
+        initial_pool_size,
+        maximum_pool_size,
+        managed_memory,
+        async_alloc,
+        log_directory,
+    ):
+        self.initial_pool_size = initial_pool_size
+        self.maximum_pool_size = maximum_pool_size
         self.managed_memory = managed_memory
         self.async_alloc = async_alloc
         self.logging = log_directory is not None
@@ -63,15 +71,16 @@ def setup(self, worker=None):
                         worker, self.logging, self.log_directory
                     )
                 )
-        elif self.nbytes is not None or self.managed_memory:
+        elif self.initial_pool_size is not None or self.managed_memory:
             import rmm
 
-            pool_allocator = False if self.nbytes is None else True
+            pool_allocator = False if self.initial_pool_size is None else True
 
             rmm.reinitialize(
                 pool_allocator=pool_allocator,
                 managed_memory=self.managed_memory,
-                initial_pool_size=self.nbytes,
+                initial_pool_size=self.initial_pool_size,
+                maximum_pool_size=self.maximum_pool_size,
                 logging=self.logging,
                 log_file_name=get_rmm_log_file_name(
                     worker, self.logging, self.log_directory