From 7389387cefb436ff6abf11e5ec0fe39039e0b503 Mon Sep 17 00:00:00 2001 From: John Kirkham Date: Fri, 14 Aug 2020 11:26:13 -0700 Subject: [PATCH 1/2] Run `black` on Dask-CUDA --- dask_cuda/benchmarks/utils.py | 2 +- dask_cuda/cuda_worker.py | 2 +- dask_cuda/initialize.py | 4 ++-- dask_cuda/local_cuda_cluster.py | 2 +- dask_cuda/utils.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/dask_cuda/benchmarks/utils.py b/dask_cuda/benchmarks/utils.py index 40b969a61..3d07474df 100644 --- a/dask_cuda/benchmarks/utils.py +++ b/dask_cuda/benchmarks/utils.py @@ -187,7 +187,7 @@ def get_scheduler_workers(dask_scheduler=None): def setup_memory_pool(pool_size=None, disable_pool=False): import cupy - os.environ['RMM_NO_INITIALIZE'] = 'True' + os.environ["RMM_NO_INITIALIZE"] = "True" import rmm rmm.reinitialize( diff --git a/dask_cuda/cuda_worker.py b/dask_cuda/cuda_worker.py index 0afd9928e..897e3a85d 100644 --- a/dask_cuda/cuda_worker.py +++ b/dask_cuda/cuda_worker.py @@ -133,7 +133,7 @@ def del_pid_file(): if rmm_pool_size is not None or rmm_managed_memory: try: - os.environ['RMM_NO_INITIALIZE'] = 'True' + os.environ["RMM_NO_INITIALIZE"] = "True" import rmm # noqa F401 except ImportError: raise ValueError( diff --git a/dask_cuda/initialize.py b/dask_cuda/initialize.py index d98849d5b..8d98f596b 100644 --- a/dask_cuda/initialize.py +++ b/dask_cuda/initialize.py @@ -47,7 +47,7 @@ def initialize( ): if create_cuda_context: try: - os.environ['RMM_NO_INITIALIZE'] = 'True' + os.environ["RMM_NO_INITIALIZE"] = "True" numba.cuda.current_context() except Exception: logger.error("Unable to start CUDA Context", exc_info=True) @@ -107,7 +107,7 @@ def dask_setup( ): if create_cuda_context: try: - os.environ['RMM_NO_INITIALIZE'] = 'True' + os.environ["RMM_NO_INITIALIZE"] = "True" numba.cuda.current_context() except Exception: logger.error("Unable to start CUDA Context", exc_info=True) diff --git a/dask_cuda/local_cuda_cluster.py b/dask_cuda/local_cuda_cluster.py index 6f08519da..a725ff99d 100644 --- a/dask_cuda/local_cuda_cluster.py +++ b/dask_cuda/local_cuda_cluster.py @@ -168,7 +168,7 @@ def __init__( self.rmm_managed_memory = rmm_managed_memory if rmm_pool_size is not None or rmm_managed_memory: try: - os.environ['RMM_NO_INITIALIZE'] = 'True' + os.environ["RMM_NO_INITIALIZE"] = "True" import rmm # noqa F401 except ImportError: raise ValueError( diff --git a/dask_cuda/utils.py b/dask_cuda/utils.py index b65507ace..6a56ac946 100644 --- a/dask_cuda/utils.py +++ b/dask_cuda/utils.py @@ -35,7 +35,7 @@ def __init__(self, nbytes, managed_memory): def setup(self, worker=None): if self.nbytes is not None or self.managed_memory is True: - os.environ['RMM_NO_INITIALIZE'] = 'True' + os.environ["RMM_NO_INITIALIZE"] = "True" import rmm pool_allocator = False if self.nbytes is None else True From 1eac90701ded0333755bb95a4d8f38ba547222dd Mon Sep 17 00:00:00 2001 From: John Kirkham Date: Fri, 14 Aug 2020 11:28:34 -0700 Subject: [PATCH 2/2] Replace RMM_NO_INITIALIZE w/ RAPIDS_NO_INITIALIZE --- dask_cuda/benchmarks/utils.py | 2 +- dask_cuda/cuda_worker.py | 2 +- dask_cuda/initialize.py | 4 ++-- dask_cuda/local_cuda_cluster.py | 2 +- dask_cuda/utils.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/dask_cuda/benchmarks/utils.py b/dask_cuda/benchmarks/utils.py index 3d07474df..5aeaa9b0f 100644 --- a/dask_cuda/benchmarks/utils.py +++ b/dask_cuda/benchmarks/utils.py @@ -187,7 +187,7 @@ def get_scheduler_workers(dask_scheduler=None): def setup_memory_pool(pool_size=None, disable_pool=False): import cupy - os.environ["RMM_NO_INITIALIZE"] = "True" + os.environ["RAPIDS_NO_INITIALIZE"] = "True" import rmm rmm.reinitialize( diff --git a/dask_cuda/cuda_worker.py b/dask_cuda/cuda_worker.py index 897e3a85d..f9a82f686 100644 --- a/dask_cuda/cuda_worker.py +++ b/dask_cuda/cuda_worker.py @@ -133,7 +133,7 @@ def del_pid_file(): if rmm_pool_size is not None or rmm_managed_memory: try: - os.environ["RMM_NO_INITIALIZE"] = "True" + os.environ["RAPIDS_NO_INITIALIZE"] = "True" import rmm # noqa F401 except ImportError: raise ValueError( diff --git a/dask_cuda/initialize.py b/dask_cuda/initialize.py index 8d98f596b..6fc425484 100644 --- a/dask_cuda/initialize.py +++ b/dask_cuda/initialize.py @@ -47,7 +47,7 @@ def initialize( ): if create_cuda_context: try: - os.environ["RMM_NO_INITIALIZE"] = "True" + os.environ["RAPIDS_NO_INITIALIZE"] = "True" numba.cuda.current_context() except Exception: logger.error("Unable to start CUDA Context", exc_info=True) @@ -107,7 +107,7 @@ def dask_setup( ): if create_cuda_context: try: - os.environ["RMM_NO_INITIALIZE"] = "True" + os.environ["RAPIDS_NO_INITIALIZE"] = "True" numba.cuda.current_context() except Exception: logger.error("Unable to start CUDA Context", exc_info=True) diff --git a/dask_cuda/local_cuda_cluster.py b/dask_cuda/local_cuda_cluster.py index a725ff99d..fd154326c 100644 --- a/dask_cuda/local_cuda_cluster.py +++ b/dask_cuda/local_cuda_cluster.py @@ -168,7 +168,7 @@ def __init__( self.rmm_managed_memory = rmm_managed_memory if rmm_pool_size is not None or rmm_managed_memory: try: - os.environ["RMM_NO_INITIALIZE"] = "True" + os.environ["RAPIDS_NO_INITIALIZE"] = "True" import rmm # noqa F401 except ImportError: raise ValueError( diff --git a/dask_cuda/utils.py b/dask_cuda/utils.py index 6a56ac946..637712386 100644 --- a/dask_cuda/utils.py +++ b/dask_cuda/utils.py @@ -35,7 +35,7 @@ def __init__(self, nbytes, managed_memory): def setup(self, worker=None): if self.nbytes is not None or self.managed_memory is True: - os.environ["RMM_NO_INITIALIZE"] = "True" + os.environ["RAPIDS_NO_INITIALIZE"] = "True" import rmm pool_allocator = False if self.nbytes is None else True