cuda_configure: make find_libs() compatible with remote execution

This change moves the logic of _find_cuda_lib() to check_cuda_libs.py. Instead of invoking _find_cuda_lib() once per library we now invoke check_cuda_libs.py once with a list of all libraries to look for as arguments.

For Example:
python check_cuda_libs.py /usr/local/cuda/lib64/libcudart.so.10.1 True /usr/local/cuda/lib64/libcudart_static.a False

PiperOrigin-RevId: 295765176
Change-Id: I743770ff640d009272f62c4ed5a89044b5343972
This commit is contained in:
Jakob Buchgraber 2020-02-18 10:29:11 -08:00 committed by TensorFlower Gardener
parent 9211d97305
commit 884a14ac9a
2 changed files with 153 additions and 77 deletions

89
third_party/gpus/check_cuda_libs.py vendored Normal file
View File

@ -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()

View File

@ -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():