From 76800bc768c5757869314ec1928b01ec49f32ee1 Mon Sep 17 00:00:00 2001
From: Rick Ratzel <3039903+rlratzel@users.noreply.github.com>
Date: Tue, 12 Jul 2022 12:27:32 -0400
Subject: [PATCH] [HOTFIX] Generalize ModuleNotFoundError exception handling to
 ImportError for Amazon SMS libnuma.so.1 bug (#2385)

Generalize ModuleNotFoundError exception handling to ImportError for Amazon SMS libnuma.so.1 bug
https://github.com/rapidsai/cugraph/issues/2113

This allows cugraph to be imported in a SageMaker environment without having to remove `ucx-py`

Authors:
   - Dylan Chima-Sanchez (https://github.com/betochimas)
   - Rick Ratzel (https://github.com/rlratzel)

Approvers:
   - Brad Rees (https://github.com/BradReesWork)
---
 conda/recipes/libcugraph/meta.yaml                | 2 +-
 python/cugraph/cugraph/dask/common/input_utils.py | 6 ++++--
 python/cugraph/cugraph/dask/common/mg_utils.py    | 6 ++++--
 python/cugraph/cugraph/dask/comms/comms.py        | 6 ++++--
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/conda/recipes/libcugraph/meta.yaml b/conda/recipes/libcugraph/meta.yaml
index 76bed2d0517..e56fd9bed6e 100644
--- a/conda/recipes/libcugraph/meta.yaml
+++ b/conda/recipes/libcugraph/meta.yaml
@@ -41,7 +41,7 @@ requirements:
     - libraft-headers {{ minor_version }}.*
     - libcugraphops {{ minor_version }}.*
     - librmm {{ minor_version }}.*
-    - cudf {{ minor_version }}.*
+    - libcudf {{ minor_version }}.*
     - boost-cpp {{ boost_cpp_version }}
     - nccl {{ nccl_version }}
     - ucx-proc=*=gpu
diff --git a/python/cugraph/cugraph/dask/common/input_utils.py b/python/cugraph/cugraph/dask/common/input_utils.py
index ca2b45a0a26..9db3964c021 100644
--- a/python/cugraph/cugraph/dask/common/input_utils.py
+++ b/python/cugraph/cugraph/dask/common/input_utils.py
@@ -25,8 +25,10 @@
 from cugraph.dask.common.read_utils import MissingUCXPy
 try:
     from raft.dask.common.utils import get_client
-except ModuleNotFoundError as err:
-    if err.name == "ucp":
+except ImportError as err:
+    # FIXME: Generalize since err.name is arr when
+    # libnuma.so.1 is not available
+    if err.name == "ucp" or err.name == "arr":
         get_client = MissingUCXPy()
     else:
         raise
diff --git a/python/cugraph/cugraph/dask/common/mg_utils.py b/python/cugraph/cugraph/dask/common/mg_utils.py
index 4ac472e36e5..99d80b3c659 100644
--- a/python/cugraph/cugraph/dask/common/mg_utils.py
+++ b/python/cugraph/cugraph/dask/common/mg_utils.py
@@ -23,8 +23,10 @@
 from cugraph.dask.common.read_utils import MissingUCXPy
 try:
     from raft.dask.common.utils import default_client
-except ModuleNotFoundError as err:
-    if err.name == "ucp":
+except ImportError as err:
+    # FIXME: Generalize since err.name is arr when
+    # libnuma.so.1 is not available
+    if err.name == "ucp" or err.name == "arr":
         default_client = MissingUCXPy()
     else:
         raise
diff --git a/python/cugraph/cugraph/dask/comms/comms.py b/python/cugraph/cugraph/dask/comms/comms.py
index 4837628dec6..c1f3ac304d5 100644
--- a/python/cugraph/cugraph/dask/comms/comms.py
+++ b/python/cugraph/cugraph/dask/comms/comms.py
@@ -17,8 +17,10 @@
 try:
     from raft.dask.common.comms import Comms as raftComms
     from raft.dask.common.comms import get_raft_comm_state
-except ModuleNotFoundError as err:
-    if err.name == "ucp":
+except ImportError as err:
+    # FIXME: Generalize since err.name is arr when
+    # libnuma.so.1 is not available
+    if err.name == "ucp" or err.name == "arr":
         raftComms = MissingUCXPy()
         get_raft_comm_state = MissingUCXPy()
     else: