Added further automatic IBVERB lib and header check before enabling THD/Gloo IB support (#5264)

* Added further automatic IBVERB lib and header check before enabling THD/Gloo IB support

* Refectoring and addressed comments

5 files changed, 71 insertions, 42 deletions

diff --git a/tools/setup_helpers/cuda.py b/tools/setup_helpers/cuda.py
index 085c4eac17..60c957b1c6 100644
--- a/tools/setup_helpers/cuda.py
+++ b/tools/setup_helpers/cuda.py
@@ -1,19 +1,14 @@
 import os
 import glob
 import re
-import platform
 import ctypes.util
 from subprocess import Popen, PIPE
 
-from .env import check_env_flag
+from .env import IS_WINDOWS, IS_LINUX, IS_DARWIN, check_env_flag
 
 LINUX_HOME = '/usr/local/cuda'
 WINDOWS_HOME = glob.glob('C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*')
 
-IS_WINDOWS = platform.system() == 'Windows'
-IS_LINUX = platform.system() == 'Linux'
-IS_DARWIN = platform.system() == 'Darwin'
-
 
 def find_nvcc():
     if IS_WINDOWS:
diff --git a/tools/setup_helpers/cudnn.py b/tools/setup_helpers/cudnn.py
index 87efcfda3b..1cc48a7c11 100644
--- a/tools/setup_helpers/cudnn.py
+++ b/tools/setup_helpers/cudnn.py
@@ -1,21 +1,10 @@
 import os
-import platform
-import sys
 import glob
-from itertools import chain
 
-from .env import check_env_flag
+from .env import IS_WINDOWS, IS_CONDA, CONDA_DIR, check_env_flag, gather_paths
 from .cuda import WITH_CUDA, CUDA_HOME
 
 
-def gather_paths(env_vars):
-    return list(chain(*(os.getenv(v, '').split(':') for v in env_vars)))
-
-
-IS_WINDOWS = (platform.system() == 'Windows')
-IS_CONDA = 'conda' in sys.version or 'Continuum' in sys.version
-CONDA_DIR = os.path.join(os.path.dirname(sys.executable), '..')
-
 WITH_CUDNN = False
 CUDNN_LIB_DIR = None
 CUDNN_INCLUDE_DIR = None
diff --git a/tools/setup_helpers/dist_check.py b/tools/setup_helpers/dist_check.py
index fd1ba728ad..cc54f8310f 100644
--- a/tools/setup_helpers/dist_check.py
+++ b/tools/setup_helpers/dist_check.py
@@ -1,13 +1,12 @@
 import os
-import platform
 import subprocess
+import glob
 
-from .env import check_env_flag
+from .env import IS_CONDA, IS_WINDOWS, CONDA_DIR, check_env_flag, gather_paths
 
-IS_WINDOWS = (platform.system() == 'Windows')
 
-WITH_DISTRIBUTED = not check_env_flag('NO_DISTRIBUTED') and not IS_WINDOWS
-WITH_DISTRIBUTED_MW = WITH_DISTRIBUTED and check_env_flag('WITH_DISTRIBUTED_MW')
+WITH_DISTRIBUTED = not check_env_flag("NO_DISTRIBUTED") and not IS_WINDOWS
+WITH_DISTRIBUTED_MW = WITH_DISTRIBUTED and check_env_flag("WITH_DISTRIBUTED_MW")
 WITH_GLOO_IBVERBS = False
 
 IB_DEVINFO_CMD = "ibv_devinfo"
@@ -34,12 +33,16 @@ def should_build_ib():
     Helper function that detects the system's IB support and returns if we
     should build with IB support.
     """
+    ib_util_found = False
+    ib_lib_found = False
+    ib_header_found = False
+
     try:
         # If the command doesn't exist, we can directly return instead of
         # making a subprocess call
         full_cmd_path = get_command_path(IB_DEVINFO_CMD)
         if not full_cmd_path:
-            return False
+            ib_util_found = False
         subprocess.check_output([full_cmd_path, "--list"])
         # Here we just would like to simply run the command to test if IB
         # related tools / lib are installed without parsing the output. We
@@ -58,11 +61,46 @@ def should_build_ib():
         #   mlx5_2
         #   mlx5_1
         #   mlx5_0
-        return True
+        ib_util_found = True
     except Exception:
         # We just take all the exceptions here without affecting the build
-        return False
+        ib_util_found = False
+
+    lib_paths = list(filter(bool, [
+        "/usr/lib/",
+        "/usr/lib/x86_64-linux-gnu/",
+        "/usr/lib/powerpc64le-linux-gnu/",
+        "/usr/lib/aarch64-linux-gnu/",
+    ] + gather_paths([
+        "LIBRARY_PATH",
+    ]) + gather_paths([
+        "LD_LIBRARY_PATH",
+    ])))
+
+    include_paths = [
+        "/usr/include/",
+    ]
+
+    if IS_CONDA:
+        lib_paths.append(os.path.join(CONDA_DIR, "lib"))
+        include_paths.append(os.path.join(CONDA_DIR, "include"))
+
+    for path in lib_paths:
+        if path is None or not os.path.exists(path):
+            continue
+        ib_libraries = sorted(glob.glob(os.path.join(path, "libibverbs*")))
+        if ib_libraries:
+            ib_lib_found = True
+            break
+
+    for path in include_paths:
+        if path is None or not os.path.exists(path):
+            continue
+        if os.path.exists(os.path.join(path, "infiniband/verbs.h")):
+            ib_header_found = True
+            break
 
+    return ib_util_found and ib_lib_found and ib_lib_found
 
-WITH_GLOO_IBVERBS = WITH_DISTRIBUTED and (should_build_ib() or
-                                          check_env_flag("WITH_GLOO_IBVERBS"))
+if WITH_DISTRIBUTED:
+    WITH_GLOO_IBVERBS = should_build_ib() or check_env_flag("WITH_GLOO_IBVERBS")
diff --git a/tools/setup_helpers/env.py b/tools/setup_helpers/env.py
index 6c3dbc389d..f768bb6505 100644
--- a/tools/setup_helpers/env.py
+++ b/tools/setup_helpers/env.py
@@ -1,5 +1,21 @@
 import os
+import platform
+import sys
+from itertools import chain
+
+
+IS_WINDOWS = (platform.system() == 'Windows')
+IS_DARWIN = (platform.system() == 'Darwin')
+IS_LINUX = (platform.system() == 'Linux')
+
+
+IS_CONDA = 'conda' in sys.version or 'Continuum' in sys.version
+CONDA_DIR = os.path.join(os.path.dirname(sys.executable), '..')
 
 
 def check_env_flag(name):
     return os.getenv(name, '').upper() in ['ON', '1', 'YES', 'TRUE', 'Y']
+
+
+def gather_paths(env_vars):
+    return list(chain(*(os.getenv(v, '').split(':') for v in env_vars)))
diff --git a/tools/setup_helpers/nccl.py b/tools/setup_helpers/nccl.py
index 7599312549..fcb68a798c 100644
--- a/tools/setup_helpers/nccl.py
+++ b/tools/setup_helpers/nccl.py
@@ -1,22 +1,13 @@
 import os
-import sys
 import glob
-import platform
 import warnings
 from itertools import chain
 
-from .env import check_env_flag
-from .cuda import WITH_CUDA, CUDA_HOME
-
+from .env import IS_WINDOWS, IS_DARWIN, IS_CONDA, CONDA_DIR, check_env_flag, \
+    gather_paths
 
-def gather_paths(env_vars):
-    return list(chain(*(os.getenv(v, '').split(':') for v in env_vars)))
-
-is_conda = 'conda' in sys.version or 'Continuum' in sys.version
-conda_dir = os.path.join(os.path.dirname(sys.executable), '..')
+from .cuda import WITH_CUDA, CUDA_HOME
 
-IS_WINDOWS = (platform.system() == 'Windows')
-IS_DARWIN = (platform.system() == 'Darwin')
 
 WITH_NCCL = WITH_CUDA and not IS_DARWIN and not IS_WINDOWS
 WITH_SYSTEM_NCCL = False
@@ -52,8 +43,8 @@ if WITH_CUDA and not check_env_flag('NO_SYSTEM_NCCL'):
         '/usr/include'
     ]))
 
-    if is_conda:
-        lib_paths.append(os.path.join(conda_dir, 'lib'))
+    if IS_CONDA:
+        lib_paths.append(os.path.join(CONDA_DIR, 'lib'))
     for path in lib_paths:
         path = os.path.expanduser(path)
         if path is None or not os.path.exists(path):