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:
parent
9211d97305
commit
884a14ac9a
89
third_party/gpus/check_cuda_libs.py
vendored
Normal file
89
third_party/gpus/check_cuda_libs.py
vendored
Normal 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()
|
||||
|
||||
141
third_party/gpus/cuda_configure.bzl
vendored
141
third_party/gpus/cuda_configure.bzl
vendored
@ -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():
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user