diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py index 54a85279ffd..bc48e1a6e73 100644 --- a/distributed/comm/ucx.py +++ b/distributed/comm/ucx.py @@ -101,10 +101,15 @@ def device_array(n): ) pool_size_str = dask.config.get("rmm.pool-size") + enable_logging = dask.config.get("rmm.logging") 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 + pool_allocator=True, + managed_memory=False, + initial_pool_size=pool_size, + logging=enable_logging, + log_file_name="", # Something more useful here ;) ) diff --git a/distributed/distributed-schema.yaml b/distributed/distributed-schema.yaml index 1a1f2566a63..40c95d0a406 100644 --- a/distributed/distributed-schema.yaml +++ b/distributed/distributed-schema.yaml @@ -770,6 +770,10 @@ properties: - "null" description: The size of the memory pool in bytes + logging: + type: boolean + description: + Whether to enable logging with RMM ucx: type: object description: | diff --git a/distributed/distributed.yaml b/distributed/distributed.yaml index cbf66dc220d..a86afc9edf3 100644 --- a/distributed/distributed.yaml +++ b/distributed/distributed.yaml @@ -172,6 +172,7 @@ distributed: pdb-on-err: False # enter debug mode on scheduling error rmm: pool-size: null + logging: False ucx: tcp: null # enable tcp nvlink: null # enable cuda_ipc