diff --git a/third_party/gpus/check_cuda_libs.py b/third_party/gpus/check_cuda_libs.py new file mode 100644 index 00000000000..b7b36e6466e --- /dev/null +++ b/third_party/gpus/check_cuda_libs.py @@ -0,0 +1,89 @@ +# Copyright 2020 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Verifies that a list of libraries is installed on the system. + +Takes a a list of arguments with every two subsequent arguments being a logical +tuple of (path, check_soname). The path to the library and either True or False +to indicate whether to check the soname field on the shared library. + +Example Usage: +./check_cuda_libs.py /path/to/lib1.so True /path/to/lib2.so False +""" +import os +import os.path +import platform +import subprocess +import sys + +# pylint: disable=g-import-not-at-top,g-importing-member +try: + from shutil import which +except ImportError: + from distutils.spawn import find_executable as which +# pylint: enable=g-import-not-at-top,g-importing-member + + +class ConfigError(Exception): + pass + + +def _is_windows(): + return platform.system() == "Windows" + + +def check_cuda_lib(path, check_soname=True): + """Tests if a library exists on disk and whether its soname matches the filename. + + Args: + path: the path to the library. + check_soname: whether to check the soname as well. + + Raises: + ConfigError: If the library does not exist or if its soname does not match + the filename. + """ + if not os.path.isfile(path): + raise ConfigError("No library found under: " + path) + objdump = which("objdump") + if check_soname and objdump is not None and not _is_windows(): + # Decode is necessary as in py3 the return type changed from str to bytes + output = subprocess.check_output([objdump, "-p", path]).decode("ascii") + output = [line for line in output.splitlines() if "SONAME" in line] + sonames = [line.strip().split(" ")[-1] for line in output] + if not any([soname == os.path.basename(path) for soname in sonames]): + raise ConfigError("None of the libraries match their SONAME: " + path) + + +def main(): + try: + args = [argv for argv in sys.argv[1:]] + if len(args) % 2 == 1: + raise ConfigError("Expected even number of arguments") + checked_paths = [] + for i in range(0, len(args), 2): + path = args[i] + check_cuda_lib(path, check_soname=args[i + 1] == "True") + checked_paths.append(path) + # pylint: disable=superfluous-parens + print(os.linesep.join(checked_paths)) + # pylint: enable=superfluous-parens + except ConfigError as e: + sys.stderr.write(str(e)) + sys.exit(1) + + +if __name__ == "__main__": + main() + diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl index c15f3c08189..6fbe306457f 100644 --- a/third_party/gpus/cuda_configure.bzl +++ b/third_party/gpus/cuda_configure.bzl @@ -40,6 +40,7 @@ load( load( "//third_party/remote_config:common.bzl", "err_out", + "execute", "get_bash_bin", "get_cpu_value", "get_python_bin", @@ -447,67 +448,46 @@ def lib_name(base_name, cpu_value, version = None, static = False): else: auto_configure_fail("Invalid cpu_value: %s" % cpu_value) -def find_lib(repository_ctx, paths, check_soname = True): - """ - Finds a library among a list of potential paths. - - Args: - paths: List of paths to inspect. - - Returns: - Returns the first path in paths that exist. - """ - objdump = repository_ctx.which("objdump") - mismatches = [] - for path in [repository_ctx.path(path) for path in paths]: - if not path.exists: - continue - if check_soname and objdump != None and not is_windows(repository_ctx): - output = raw_exec(repository_ctx, [objdump, "-p", str(path)]).stdout - output = [line for line in output.splitlines() if "SONAME" in line] - sonames = [line.strip().split(" ")[-1] for line in output] - if not any([soname == path.basename for soname in sonames]): - mismatches.append(str(path)) - continue - return str(path) - if mismatches: - auto_configure_fail( - "None of the libraries match their SONAME: " + ", ".join(mismatches), - ) - auto_configure_fail("No library found under: " + ", ".join(paths)) - -def _find_cuda_lib( - lib, - repository_ctx, - cpu_value, - basedir, - version, - static = False): - """Finds the given CUDA or cuDNN library on the system. - - Args: - lib: The name of the library, such as "cudart" - repository_ctx: The repository context. - cpu_value: The name of the host operating system. - basedir: The install directory of CUDA or cuDNN. - version: The version of the library. - static: True if static library, False if shared object. - - Returns: - Returns the path to the library. - """ +def _lib_path(lib, cpu_value, basedir, version, static): file_name = lib_name(lib, cpu_value, version, static) - return find_lib( - repository_ctx, - ["%s/%s" % (basedir, file_name)], - check_soname = version and not static, + return "%s/%s" % (basedir, file_name) + +def _should_check_soname(version, static): + return version and not static + +def _check_cuda_lib_params(lib, cpu_value, basedir, version, static = False): + return ( + _lib_path(lib, cpu_value, basedir, version, static), + _should_check_soname(version, static), ) -def _find_libs(repository_ctx, cuda_config): +def _check_cuda_libs(repository_ctx, script_path, libs): + python_bin = get_python_bin(repository_ctx) + contents = repository_ctx.read(script_path).splitlines() + + cmd = "from os import linesep;" + cmd += "f = open('script.py', 'w');" + for line in contents: + cmd += "f.write('%s' + linesep);" % line + cmd += "f.close();" + cmd += "from os import system;" + args = " ".join([path + " " + str(check) for path, check in libs]) + cmd += "system('%s script.py %s');" % (python_bin, args) + + all_paths = [path for path, _ in libs] + checked_paths = execute(repository_ctx, [python_bin, "-c", cmd]).stdout.splitlines() + if all_paths != checked_paths: + auto_configure_fail("Error with installed CUDA libs. Expected '%s'. Actual '%s'." % (all_paths, checked_paths)) + +def _find_libs(repository_ctx, check_cuda_libs_script, cuda_config): """Returns the CUDA and cuDNN libraries on the system. + Also, verifies that the script actually exist. + Args: repository_ctx: The repository context. + check_cuda_libs_script: The path to a script verifying that the cuda + libraries exist on the system. cuda_config: The CUDA config as returned by _get_cuda_config Returns: @@ -515,80 +495,86 @@ def _find_libs(repository_ctx, cuda_config): """ cpu_value = cuda_config.cpu_value stub_dir = "" if is_windows(repository_ctx) else "/stubs" - return { - "cuda": _find_cuda_lib( + + check_cuda_libs_params = { + "cuda": _check_cuda_lib_params( "cuda", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"] + stub_dir, - None, + version = None, + static = False, ), - "cudart": _find_cuda_lib( + "cudart": _check_cuda_lib_params( "cudart", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"], cuda_config.cuda_version, + static = False, ), - "cudart_static": _find_cuda_lib( + "cudart_static": _check_cuda_lib_params( "cudart_static", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"], cuda_config.cuda_version, static = True, ), - "cublas": _find_cuda_lib( + "cublas": _check_cuda_lib_params( "cublas", - repository_ctx, cpu_value, cuda_config.config["cublas_library_dir"], cuda_config.cuda_lib_version, + static = False, ), - "cusolver": _find_cuda_lib( + "cusolver": _check_cuda_lib_params( "cusolver", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"], cuda_config.cuda_lib_version, + static = False, ), - "curand": _find_cuda_lib( + "curand": _check_cuda_lib_params( "curand", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"], cuda_config.cuda_lib_version, + static = False, ), - "cufft": _find_cuda_lib( + "cufft": _check_cuda_lib_params( "cufft", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"], cuda_config.cuda_lib_version, + static = False, ), - "cudnn": _find_cuda_lib( + "cudnn": _check_cuda_lib_params( "cudnn", - repository_ctx, cpu_value, cuda_config.config["cudnn_library_dir"], cuda_config.cudnn_version, + static = False, ), - "cupti": _find_cuda_lib( + "cupti": _check_cuda_lib_params( "cupti", - repository_ctx, cpu_value, cuda_config.config["cupti_library_dir"], cuda_config.cuda_version, + static = False, ), - "cusparse": _find_cuda_lib( + "cusparse": _check_cuda_lib_params( "cusparse", - repository_ctx, cpu_value, cuda_config.config["cuda_library_dir"], cuda_config.cuda_lib_version, + static = False, ), } + # Verify that the libs actually exist at their locations. + _check_cuda_libs(repository_ctx, check_cuda_libs_script, check_cuda_libs_params.values()) + + paths = {filename: v[0] for (filename, v) in check_cuda_libs_params.items()} + return paths + def _cudart_static_linkopt(cpu_value): """Returns additional platform-specific linkopts for cudart.""" return "" if cpu_value == "Darwin" else "\"-lrt\"," @@ -924,7 +910,8 @@ def _create_local_cuda_repository(repository_ctx): ], )) - cuda_libs = _find_libs(repository_ctx, cuda_config) + check_cuda_libs_script = repository_ctx.path(Label("@org_tensorflow//third_party/gpus:check_cuda_libs.py")) + cuda_libs = _find_libs(repository_ctx, check_cuda_libs_script, cuda_config) cuda_lib_srcs = [] cuda_lib_outs = [] for path in cuda_libs.values():