diff --git a/dask_cuda/benchmarks/utils.py b/dask_cuda/benchmarks/utils.py index 5aeaa9b0f..32d92158e 100644 --- a/dask_cuda/benchmarks/utils.py +++ b/dask_cuda/benchmarks/utils.py @@ -1,5 +1,4 @@ import argparse -import os from dask.distributed import SSHCluster @@ -187,7 +186,6 @@ def get_scheduler_workers(dask_scheduler=None): def setup_memory_pool(pool_size=None, disable_pool=False): import cupy - 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 f9a82f686..28d7c0d35 100644 --- a/dask_cuda/cuda_worker.py +++ b/dask_cuda/cuda_worker.py @@ -72,6 +72,10 @@ def __init__( net_devices=None, **kwargs, ): + # Required by RAPIDS libraries (e.g., cuDF) to ensure no context + # initialization happens before we can set CUDA_VISIBLE_DEVICES + os.environ["RAPIDS_NO_INITIALIZE"] = "True" + enable_proctitle_on_current() enable_proctitle_on_children() @@ -133,7 +137,6 @@ def del_pid_file(): if rmm_pool_size is not None or rmm_managed_memory: try: - 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 6fc425484..e53c10de0 100644 --- a/dask_cuda/initialize.py +++ b/dask_cuda/initialize.py @@ -24,7 +24,6 @@ about Dask configuration. """ import logging -import os import click import numba.cuda @@ -47,7 +46,6 @@ def initialize( ): if create_cuda_context: try: - os.environ["RAPIDS_NO_INITIALIZE"] = "True" numba.cuda.current_context() except Exception: logger.error("Unable to start CUDA Context", exc_info=True) @@ -107,7 +105,6 @@ def dask_setup( ): if create_cuda_context: try: - 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 2c771c6a4..6598b52eb 100644 --- a/dask_cuda/local_cuda_cluster.py +++ b/dask_cuda/local_cuda_cluster.py @@ -150,6 +150,10 @@ def __init__( rmm_managed_memory=False, **kwargs, ): + # Required by RAPIDS libraries (e.g., cuDF) to ensure no context + # initialization happens before we can set CUDA_VISIBLE_DEVICES + os.environ["RAPIDS_NO_INITIALIZE"] = "True" + if CUDA_VISIBLE_DEVICES is None: CUDA_VISIBLE_DEVICES = cuda_visible_devices(0) if isinstance(CUDA_VISIBLE_DEVICES, str): @@ -166,7 +170,6 @@ def __init__( self.rmm_managed_memory = rmm_managed_memory if rmm_pool_size is not None or rmm_managed_memory: try: - 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 637712386..72c7736c1 100644 --- a/dask_cuda/utils.py +++ b/dask_cuda/utils.py @@ -35,7 +35,6 @@ 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["RAPIDS_NO_INITIALIZE"] = "True" import rmm pool_allocator = False if self.nbytes is None else True