Fix import error when distributed.rmm.pool-size is setted
Fix import error when distributed.rmm.pool-size is set
@@ -125,11 +125,19 @@ 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)
+
+ if pool_size_str is not None:
+ pool_size = parse_bytes(pool_size_str)
+ rmm.reinitialize(
+ pool_allocator=True, managed_memory=False, initial_pool_size=pool_size
+ )
except ImportError:
try:
import numba.cuda
@@ -140,6 +148,7 @@ def init_once():
return a
device_array = numba_device_array
+
except ImportError:
def device_array(n):
@@ -147,11 +156,10 @@ def init_once():
"In order to send/recv CUDA arrays, Numba or RMM is required"
)
- 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(
- pool_allocator=True, managed_memory=False, initial_pool_size=pool_size
+ warnings.warn(
+ "Initial RMM pool size defined, but RMM is not available. "
+ "Please consider installing RMM or removing the pool size option."
)
