Now use dist_check to check all distributed related and addressed commments

teng-li · teng-li · commit b45de3a1281a · 2018-01-24T11:29:57.000-08:00
diff --git a/setup.py b/setup.py
@@ -25,16 +25,15 @@
 from tools.setup_helpers.split_types import split_types
 from tools.setup_helpers.generate_code import generate_code
 from tools.setup_helpers.ninja_builder import NinjaBuilder, ninja_build_ext
-from tools.setup_helpers.ib_detect import WITH_IB_DEVICES
+from tools.setup_helpers.dist_check import WITH_DISTRIBUTED, \
+    WITH_DISTRIBUTED_MW, WITH_GLOO_IBVERBS
 
 DEBUG = check_env_flag('DEBUG')
 
 IS_WINDOWS = (platform.system() == 'Windows')
 IS_DARWIN = (platform.system() == 'Darwin')
 IS_LINUX = (platform.system() == 'Linux')
 
-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_SCALARS = check_env_flag('WITH_SCALARS')
 
@@ -139,8 +138,7 @@ def build_libs(libs):
         my_env["CUDNN_LIBRARY"] = CUDNN_LIBRARY
         my_env["CUDNN_INCLUDE_DIR"] = CUDNN_INCLUDE_DIR
 
-    if WITH_DISTRIBUTED and (WITH_IB_DEVICES or
-                             check_env_flag("WITH_GLOO_IBVERBS")):
+    if WITH_GLOO_IBVERBS:
         build_libs_cmd += ['--with-gloo-ibverbs']
 
     if subprocess.call(build_libs_cmd + libs, env=my_env) != 0:
diff --git a/tools/setup_helpers/dist_check.py b/tools/setup_helpers/dist_check.py
@@ -0,0 +1,68 @@
+import os
+import platform
+import subprocess
+
+from .env import check_env_flag
+
+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_GLOO_IBVERBS = False
+
+IB_DEVINFO_CMD = "ibv_devinfo"
+
+
+def get_command_path(command):
+    """
+    Helper function that checks if the command exists in the path and gets the
+    full path of a given linux command if it exists.
+    """
+    def excutable(command_path):
+        return os.path.isfile(command_path) and os.access(command_path, os.X_OK)
+
+    for path in os.environ["PATH"].split(os.pathsep):
+        command_path = os.path.join(path, command)
+        if excutable(command_path):
+            return command_path
+
+    return None
+
+
+def should_build_ib():
+    """
+    Helper function that detects the system's IB support and returns if we
+    should build with IB support.
+    """
+    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
+        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
+        # will enable IB build as long as the command runs successfully.
+        #
+        # The output should look like either:
+        #
+        # > ibv_devinfo --list
+        # 0 HCAs founds:
+        #
+        # or
+        #
+        # > ibv_devinfo --list
+        # 4 HCAs found:
+        #   mlx5_3
+        #   mlx5_2
+        #   mlx5_1
+        #   mlx5_0
+        return True
+    except Exception:
+        # We just take all the exceptions here without affecting the build
+        return False
+
+
+WITH_GLOO_IBVERBS = WITH_DISTRIBUTED and (should_build_ib() or
+                                          check_env_flag("WITH_GLOO_IBVERBS"))
diff --git a/tools/setup_helpers/ib_detect.py b/tools/setup_helpers/ib_detect.py
