Add option to track RMM allocations

dask_cuda/cli/dask_cuda_worker.py

@@ -119,6 +119,14 @@
         Logging will only be enabled if ``--rmm-pool-size`` or ``--rmm-managed-memory``
         are specified.""",
 )
+@click.option(
+    "--rmm-track-allocations/--no-rmm-track-allocations",
+    default=False,
+    show_default=True,
+    help="""Track memory allocations made by RMM. If ``True``, wraps the memory
+    resource of each worker with a ``rmm.mr.TrackingResourceAdaptor`` that
+    allows querying the amount of memory allocated by RMM.""",
+)
 @click.option(
     "--pid-file", type=str, default="", help="File to write the process PID.",
 )
@@ -293,6 +301,7 @@ def main(
     rmm_managed_memory,
     rmm_async,
     rmm_log_directory,
+    rmm_track_allocations,
     pid_file,
     resources,
     dashboard,
@@ -344,6 +353,7 @@ def main(
         rmm_managed_memory,
         rmm_async,
         rmm_log_directory,
+        rmm_track_allocations,
         pid_file,
         resources,
         dashboard,

dask_cuda/cuda_worker.py

@@ -62,6 +62,7 @@ def __init__(
         rmm_managed_memory=False,
         rmm_async=False,
         rmm_log_directory=None,
+        rmm_track_allocations=False,
         pid_file=None,
         resources=None,
         dashboard=True,
@@ -248,6 +249,7 @@ def del_pid_file():
                         rmm_managed_memory,
                         rmm_async,
                         rmm_log_directory,
+                        rmm_track_allocations,
                     ),
                 },
                 name=name if nprocs == 1 or name is None else str(name) + "-" + str(i),

dask_cuda/local_cuda_cluster.py

@@ -210,6 +210,7 @@ def __init__(
         rmm_managed_memory=False,
         rmm_async=False,
         rmm_log_directory=None,
+        rmm_track_allocations=False,
         jit_unspill=None,
         log_spilling=False,
         worker_class=None,
@@ -272,6 +273,7 @@ def __init__(
                 )
 
         self.rmm_log_directory = rmm_log_directory
+        self.rmm_track_allocations = rmm_track_allocations
 
         if not kwargs.pop("processes", True):
             raise ValueError(
@@ -415,6 +417,7 @@ def new_worker_spec(self):
                         self.rmm_managed_memory,
                         self.rmm_async,
                         self.rmm_log_directory,
+                        self.rmm_track_allocations,
                     ),
                 },
             }

dask_cuda/tests/test_dask_cuda_worker.py

@@ -253,3 +253,34 @@ def test_cuda_visible_devices_uuid(loop):  # noqa: F811
 
                     result = client.run(lambda: os.environ["CUDA_VISIBLE_DEVICES"])
                     assert list(result.values())[0] == gpu_uuid
+
+
+def test_rmm_track_allocations(loop):  # noqa: F811
+    rmm = pytest.importorskip("rmm")
+    with popen(["dask-scheduler", "--port", "9369", "--no-dashboard"]):
+        with popen(
+            [
+                "dask-cuda-worker",
+                "127.0.0.1:9369",
+                "--host",
+                "127.0.0.1",
+                "--rmm-pool-size",
+                "2 GB",
+                "--no-dashboard",
+                "--rmm-track-allocations",
+            ]
+        ):
+            with Client("127.0.0.1:9369", loop=loop) as client:
+                assert wait_workers(client, n_gpus=get_n_gpus())
+
+                memory_resource_type = client.run(
+                    rmm.mr.get_current_device_resource_type
+                )
+                for v in memory_resource_type.values():
+                    assert v is rmm.mr.TrackingResourceAdaptor
+
+                memory_resource_upstream_type = client.run(
+                    lambda: type(rmm.mr.get_current_device_resource().upstream_mr)
+                )
+                for v in memory_resource_upstream_type.values():
+                    assert v is rmm.mr.PoolMemoryResource

dask_cuda/utils.py

@@ -52,6 +52,7 @@ def __init__(
         managed_memory,
         async_alloc,
         log_directory,
+        track_allocations,
     ):
         if initial_pool_size is None and maximum_pool_size is not None:
             raise ValueError(
@@ -65,10 +66,12 @@ def __init__(
         self.async_alloc = async_alloc
         self.logging = log_directory is not None
         self.log_directory = log_directory
+        self.rmm_track_allocations = track_allocations
 
     def setup(self, worker=None):
+        import rmm
+
         if self.async_alloc:
-            import rmm
 
             rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource())
             if self.logging:
@@ -78,7 +81,6 @@ def setup(self, worker=None):
                     )
                 )
         elif self.initial_pool_size is not None or self.managed_memory:
-            import rmm
 
             pool_allocator = False if self.initial_pool_size is None else True
 
@@ -92,6 +94,9 @@ def setup(self, worker=None):
                     worker, self.logging, self.log_directory
                 ),
             )
+        if self.rmm_track_allocations:
+            mr = rmm.mr.get_current_device_resource()
+            rmm.mr.set_current_device_resource(rmm.mr.TrackingResourceAdaptor(mr))
 
 
 def unpack_bitmask(x, mask_bits=64):