Refactor for distributed.rmm.pool-size warning

dask · Jun 1, 2022 · 6575251 · 6575251
1 parent 2467b8c
commit 6575251
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
@@ -125,13 +125,14 @@ def init_once():
 
     ucp.init(options=ucx_config, env_takes_precedence=True)
 
+    pool_size_str = dask.config.get("distributed.rmm.pool-size")
+
     # Find the function, `cuda_array()`, to use when allocating new CUDA arrays
     try:
         import rmm
 
         device_array = lambda n: rmm.DeviceBuffer(size=n)
 
-        pool_size_str = dask.config.get("distributed.rmm.pool-size")
         if pool_size_str is not None:
             pool_size = parse_bytes(pool_size_str)
             rmm.reinitialize(
@@ -148,19 +149,19 @@ def numba_device_array(n):
 
             device_array = numba_device_array
 
-            pool_size_str = dask.config.get("distributed.rmm.pool-size")
-            if pool_size_str is not None:
-                warnings.warn(
-                    "Initial RMM pool size defined, but RMM is not available. "
-                    "Please consider installing RMM or removing the pool size option."
-                )
         except ImportError:
 
             def device_array(n):
                 raise RuntimeError(
                     "In order to send/recv CUDA arrays, Numba or RMM is required"
                 )
 
+        if pool_size_str is not None:
+            warnings.warn(
+                "Initial RMM pool size defined, but RMM is not available. "
+                "Please consider installing RMM or removing the pool size option."
+            )
+
 
 def _close_comm(ref):
     """Callback to close Dask Comm when UCX Endpoint closes or errors