Convert build_info to dict format and expose it.

Since this module now generates a dictionary to expose in tf.config, it
doesn't make much sense to store only certain values in the build_info
dictionary and others as module variables. This obsoletes a lot of code
in gen_build_info.py and I've removed it.

I also updated all the in-code references I've found to the build_info
module. I think this may break whomever used to be using the build_info
library, but since it wasn't part of the API, there was no guarantee
that it would continue to be available.

tensorflow/python/framework/config.py

@@ -656,23 +656,26 @@ def disable_mlir_graph_optimization():
   context.context().enable_mlir_graph_optimization = False
 
 
-@tf_export('config.get_cuda_version_used_to_compile_tf')
-def get_cuda_version_used_to_compile_tf():
-  """Get the version of NVIDIA CUDA used to compile this TensorFlow release.
+@tf_export('config.get_build_info()')
+def get_build_info():
+  """Get a dictionary describing TensorFlow's build environment.
+
+  Values are generated when TensorFlow is compiled, and are static for each
+  TensorFlow package. This information is limited to a subset of the following
+  keys based on the platforms targeted by the package:
+
+    - cuda_version
+    - cudnn_version
+    - tensorrt_version
+    - nccl_version
+    - is_cuda_build
+    - is_rocm_build
+    - msvcp_dll_names
+    - nvcuda_dll_name
+    - cudart_dll_name
+    - cudnn_dll_name
 
   Returns:
-    String representation of CUDA version number (Major.Minor) if CUDA support
-    is included, otherwise None.
+    A Dictionary describing TensorFlow's build environment.
   """
   return build_info.cuda_build_info.get('cuda_version', None)
-
-
-@tf_export('config.get_cudnn_version_used_to_compile_tf')
-def get_cudnn_version_used_to_compile_tf():
-  """Get the version of NVIDIA cuDNN used to compile this TensorFlow release.
-
-  Returns:
-    String representation of cuDNN version number (Major only) if cuDNN support
-    is included, otherwise None.
-  """
-  return build_info.cuda_build_info.get('cudnn_version', None)

tensorflow/python/keras/layers/recurrent_v2.py

@@ -601,7 +601,7 @@ def gpu_gru(inputs, init_h, kernel, recurrent_kernel, bias, mask, time_major,
   # (6 * units)
   bias = array_ops.split(K.flatten(bias), 6)
 
-  if build_info.is_cuda_build:
+  if build_info.build_info["is_cuda_build"]:
     # Note that the gate order for CuDNN is different from the canonical format.
     # canonical format is [z, r, h], whereas CuDNN is [r, z, h]. The swap need
     # to be done for kernel, recurrent_kernel, input_bias, recurrent_bias.
@@ -1361,7 +1361,7 @@ def gpu_lstm(inputs, init_h, init_c, kernel, recurrent_kernel, bias, mask,
   # so that mathematically it is same as the canonical LSTM implementation.
   full_bias = array_ops.concat((array_ops.zeros_like(bias), bias), 0)
 
-  if build_info.is_rocm_build:
+  if build_info.build_info["is_rocm_build"]:
     # ROCm MIOpen's weight sequence for LSTM is different from both canonical
     # and Cudnn format
     # MIOpen: [i, f, o, c] Cudnn/Canonical: [i, f, c, o]

tensorflow/python/platform/build_info_test.py

@@ -25,8 +25,8 @@ from tensorflow.python.platform import test
 class BuildInfoTest(test.TestCase):
 
   def testBuildInfo(self):
-    self.assertEqual(build_info.is_rocm_build, test.is_built_with_rocm())
-    self.assertEqual(build_info.is_cuda_build, test.is_built_with_cuda())
+    self.assertEqual(build_info.build_info["is_rocm_build"], test.is_built_with_rocm())
+    self.assertEqual(build_info.build_info["is_cuda_build"], test.is_built_with_cuda())
 
 
 if __name__ == '__main__':

tensorflow/python/platform/self_check.py

@@ -20,6 +20,7 @@ from __future__ import print_function
 
 import os
 
+MSVCP_DLL_NAMES = "msvcp_dll_names"
 
 try:
   from tensorflow.python.platform import build_info
@@ -42,9 +43,9 @@ def preload_check():
     # we load the Python extension, so that we can raise an actionable error
     # message if they are not found.
     import ctypes  # pylint: disable=g-import-not-at-top
-    if hasattr(build_info, "msvcp_dll_names"):
+    if MSVCP_DLL_NAMES in build_info.build_info:
       missing = []
-      for dll_name in build_info.msvcp_dll_names.split(","):
+      for dll_name in build_info.build_info[MSVCP_DLL_NAMES].split(","):
         try:
           ctypes.WinDLL(dll_name)
         except OSError:

tensorflow/tensorflow.bzl

@@ -2593,6 +2593,10 @@ def tf_version_info_genrule(name, out):
         arguments = "--generate \"$@\" --git_tag_override=${GIT_TAG_OVERRIDE:-}",
     )
 
+def dict_to_kv(d):
+    """Convert a dictionary to a space-joined list of key=value pairs."""
+    return " " + " ".join(["%s=%s" % (k,v) for k, v in d.items()])
+
 def tf_py_build_info_genrule(name, out):
     _local_genrule(
         name = name,
@@ -2600,16 +2604,17 @@ def tf_py_build_info_genrule(name, out):
         exec_tool = "//tensorflow/tools/build_info:gen_build_info",
         arguments =
             "--raw_generate \"$@\" " +
-            " --is_config_cuda " + if_cuda("True", "False") +
-            " --is_config_rocm " + if_rocm("True", "False") +
-            " --key_value " +
-            if_cuda(" cuda_version_number=${TF_CUDA_VERSION:-} cudnn_version_number=${TF_CUDNN_VERSION:-} ", "") +
-            if_windows(" msvcp_dll_names=msvcp140.dll,msvcp140_1.dll ", "") +
-            if_windows_cuda(" ".join([
-                "nvcuda_dll_name=nvcuda.dll",
-                "cudart_dll_name=cudart64_$(echo $${TF_CUDA_VERSION:-} | sed \"s/\\.//\").dll",
-                "cudnn_dll_name=cudnn64_${TF_CUDNN_VERSION:-}.dll",
-            ]), ""),
+            " --key_value"
+                + " is_rocm_build=" + if_rocm("True", "False")
+                + " is_cuda_build=" + if_cuda("True", "False")
+                # TODO(angerson) Can we reliably load CUDA compute capabilities here?
+                if_windows(dict_to_kv({
+                    "msvcp_dll_names": "msvcp140.dll,msvcp140_1.dll"
+                }), "") + if_windows_cuda(dict_to_kv({
+                    "nvcuda_dll_name": "nvcuda.dll",
+                    "cudart_dll_name": "cudart64_$$(echo $${TF_CUDA_VERSION:-} | sed \"s/\\.//\").dll",
+                    "cudnn_dll_name": "cudnn64_$${TF_CUDNN_VERSION:-}.dll",
+            }), ""),
     )
 
 def cc_library_with_android_deps(

tensorflow/tools/build_info/gen_build_info.py

@@ -28,48 +28,28 @@ import six
 from third_party.gpus import find_cuda_config
 
 
-def write_build_info(filename, is_config_cuda, is_config_rocm, key_value_list):
+def write_build_info(filename, key_value_list):
   """Writes a Python that describes the build.
 
   Args:
     filename: filename to write to.
-    is_config_cuda: Whether this build is using CUDA.
-    is_config_rocm: Whether this build is using ROCm.
     key_value_list: A list of "key=value" strings that will be added to the
-      module as additional fields.
-
-  Raises:
-    ValueError: If `key_value_list` includes the key "is_cuda_build", which
-      would clash with one of the default fields.
+      module's "build_info" dictionary as additional entries.
   """
-  module_docstring = "\"\"\"Generates a Python module containing information "
-  module_docstring += "about the build.\"\"\""
 
-  build_config_rocm_bool = "False"
-  build_config_cuda_bool = "False"
-
-  if is_config_rocm == "True":
-    build_config_rocm_bool = "True"
-  elif is_config_cuda == "True":
-    build_config_cuda_bool = "True"
-
-  key_value_pair_stmts = []
-  if key_value_list:
-    for arg in key_value_list:
-      key, value = six.ensure_str(arg).split("=")
-      if key == "is_cuda_build":
-        raise ValueError("The key \"is_cuda_build\" cannot be passed as one of "
-                         "the --key_value arguments.")
-      if key == "is_rocm_build":
-        raise ValueError("The key \"is_rocm_build\" cannot be passed as one of "
-                         "the --key_value arguments.")
-      key_value_pair_stmts.append("%s = %r" % (key, value))
-  key_value_pair_content = "\n".join(key_value_pair_stmts)
+  build_info = {}
+  for arg in key_value_list:
+    key, value = six.ensure_str(arg).split("=")
+    if value.lower() == "true":
+      build_info[key] = True
+    elif value.lower() == "false":
+      build_info[key] = False
+    else:
+      build_info[key] = value
 
   # Generate cuda_build_info, a dict describing the CUDA component versions
   # used to build TensorFlow.
-  cuda_build_info = "{}"
-  if is_config_cuda == "True":
+  if build_info.get("is_cuda_build", False):
     libs = ["_", "cuda", "cudnn"]
     if platform.system() == "Linux":
       if os.environ.get("TF_NEED_TENSORRT", "0") == "1":
@@ -82,16 +62,15 @@ def write_build_info(filename, is_config_cuda, is_config_rocm, key_value_list):
     backup_argv = sys.argv
     sys.argv = libs
     cuda = find_cuda_config.find_cuda_config()
-    cuda_build_info = str({
-        "cuda_version": cuda["cuda_version"],
-        "cudnn_version": cuda["cudnn_version"],
-        "tensorrt_version": cuda.get("tensorrt_version", None),
-        "nccl_version": cuda.get("nccl_version", None),
-    })
+
+    build_info["cuda_version"] = cuda["cuda_version"]
+    build_info["cudnn_version"] = cuda["cudnn_version"]
+    build_info["tensorrt_version"] = cuda.get("tensorrt_version", None)
+    build_info["nccl_version"] = cuda.get("nccl_version", None)
     sys.argv = backup_argv
 
   contents = f"""
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -105,15 +84,14 @@ def write_build_info(filename, is_config_cuda, is_config_rocm, key_value_list):
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-{module_docstring}
+\"\"\"Auto-generated module providing information about the build.\"\"\""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-is_rocm_build = {build_config_rocm_bool}
-is_cuda_build = {build_config_cuda_bool}
-cuda_build_info = {cuda_build_info}
+from collections import namedtuple
 
+build_info = {build_info}
 """
   open(filename, "w").write(contents)
 
@@ -121,16 +99,6 @@ cuda_build_info = {cuda_build_info}
 parser = argparse.ArgumentParser(
     description="""Build info injection into the PIP package.""")
 
-parser.add_argument(
-    "--is_config_cuda",
-    type=str,
-    help="'True' for CUDA GPU builds, 'False' otherwise.")
-
-parser.add_argument(
-    "--is_config_rocm",
-    type=str,
-    help="'True' for ROCm GPU builds, 'False' otherwise.")
-
 parser.add_argument("--raw_generate", type=str, help="Generate build_info.py")
 
 parser.add_argument(
@@ -138,10 +106,9 @@ parser.add_argument(
 
 args = parser.parse_args()
 
-if (args.raw_generate is not None) and (args.is_config_cuda is not None) and (
-    args.is_config_rocm is not None):
-  write_build_info(args.raw_generate, args.is_config_cuda, args.is_config_rocm,
-                   args.key_value)
+if args.raw_generate:
+  print(args.key_value)
+  write_build_info(args.raw_generate, args.key_value)
 else:
   raise RuntimeError(
-      "--raw_generate, --is_config_cuda and --is_config_rocm must be used")
+      "--raw_generate must be used.")

tensorflow/tools/pip_package/setup.py

@@ -86,7 +86,7 @@ REQUIRED_PACKAGES = [
 ]
 
 GPU_DESCRIPTION = ''
-if build_info.is_cuda_build:
+if build_info.build_info['is_cuda_build']:
   gpu_header = (f'\nTensorFlow {_VERSION} for NVIDIA GPUs was built with these '
                 'platform and library versions:\n\n  - ')
   cbi = build_info.cuda_build_info