diff --git a/dask_cuda/benchmarks/utils.py b/dask_cuda/benchmarks/utils.py index 40b969a61..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 0afd9928e..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 d98849d5b..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 6f08519da..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 b65507ace..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