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

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

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):
+def _lib_path(lib, cpu_value, basedir, version, static):
-    """
-      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.
-      """
     file_name = lib_name(lib, cpu_value, version, static)
-    return find_lib(
+    return "%s/%s" % (basedir, file_name)
-        repository_ctx,
+
-        ["%s/%s" % (basedir, file_name)],
+def _should_check_soname(version, static):
-        check_soname = version and not 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():