From 4157562c37c76902c79e7eca275951f3a4b1ef78 Mon Sep 17 00:00:00 2001 From: Teng Li Date: Thu, 15 Feb 2018 23:09:08 -0800 Subject: 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 --- tools/setup_helpers/cuda.py | 7 +---- tools/setup_helpers/cudnn.py | 13 +-------- tools/setup_helpers/dist_check.py | 58 ++++++++++++++++++++++++++++++++------- tools/setup_helpers/env.py | 16 +++++++++++ tools/setup_helpers/nccl.py | 19 ++++--------- 5 files changed, 71 insertions(+), 42 deletions(-) (limited to 'tools') 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): -- cgit v1.2.3