From a7365b37854bd543ec7bdffe5f68ace1c5c9b96a Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 20 Sep 2021 11:56:15 -0700 Subject: [PATCH] Handle `ucp` import error during `initialize()` --- dask_cuda/initialize.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/dask_cuda/initialize.py b/dask_cuda/initialize.py index 8f159cab7..275e98552 100644 --- a/dask_cuda/initialize.py +++ b/dask_cuda/initialize.py @@ -17,7 +17,13 @@ def _create_cuda_context(): try: # Added here to ensure the parent `LocalCUDACluster` process creates the CUDA # context directly from the UCX module, thus avoiding a similar warning there. - distributed.comm.ucx.init_once() + try: + distributed.comm.ucx.init_once() + except ModuleNotFoundError: + # UCX intialization has to be delegated to Distributed, it will take care + # of setting correct environment variables and importing `ucp` after that. + # Therefore if ``import ucp`` fails we can just continue here. + pass cuda_visible_device = int( os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]