From a2bf28fc746ace4b22fba31f3c2cb1df8fb379be Mon Sep 17 00:00:00 2001 From: Vibhu Jawa Date: Mon, 7 Oct 2024 17:32:03 -0700 Subject: [PATCH] Allow passing in a list Signed-off-by: Vibhu Jawa --- dask_cuda/local_cuda_cluster.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/dask_cuda/local_cuda_cluster.py b/dask_cuda/local_cuda_cluster.py index ed95ba3b..bce6ed64 100644 --- a/dask_cuda/local_cuda_cluster.py +++ b/dask_cuda/local_cuda_cluster.py @@ -289,9 +289,12 @@ def __init__( self.rmm_managed_memory = rmm_managed_memory self.rmm_async = rmm_async self.rmm_release_threshold = rmm_release_threshold - if rmm_allocator_external_lib_list is not None: - rmm_allocator_external_lib_list = [s.strip() for s in - rmm_allocator_external_lib_list.split(',')] + if rmm_allocator_external_lib_list is not None and isinstance( + rmm_allocator_external_lib_list, str + ): + rmm_allocator_external_lib_list = [ + s.strip() for s in rmm_allocator_external_lib_list.split(",") + ] self.rmm_allocator_external_lib_list = rmm_allocator_external_lib_list if rmm_pool_size is not None or rmm_managed_memory or rmm_async: @@ -447,7 +450,7 @@ def new_worker_spec(self): release_threshold=self.rmm_release_threshold, log_directory=self.rmm_log_directory, track_allocations=self.rmm_track_allocations, - external_lib_list=self.rmm_allocator_external_lib_list + external_lib_list=self.rmm_allocator_external_lib_list, ), PreImport(self.pre_import), CUDFSetup(self.enable_cudf_spill, self.cudf_spill_stats),