xref: /aosp_15_r20/external/tensorflow/tensorflow/tools/build_info/gen_build_info.py (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Generates a Python module containing information about the build."""
16import argparse
17
18# cuda.cuda is only valid in OSS
19try:
20  from cuda.cuda import cuda_config  # pylint: disable=g-import-not-at-top
21except ImportError:
22  cuda_config = None
23
24# tensorrt.tensorrt is only valid in OSS
25try:
26  from tensorrt.tensorrt import tensorrt_config  # pylint: disable=g-import-not-at-top
27except ImportError:
28  tensorrt_config = None
29
30
31def write_build_info(filename, key_value_list):
32  """Writes a Python that describes the build.
33
34  Args:
35    filename: filename to write to.
36    key_value_list: A list of "key=value" strings that will be added to the
37      module's "build_info" dictionary as additional entries.
38  """
39
40  build_info = {}
41
42  if cuda_config:
43    build_info.update(cuda_config.config)
44
45  if tensorrt_config:
46    build_info.update(tensorrt_config.config)
47
48  for arg in key_value_list:
49    key, value = arg.split("=")
50    if value.lower() == "true":
51      build_info[key] = True
52    elif value.lower() == "false":
53      build_info[key] = False
54    else:
55      build_info[key] = value.format(**build_info)
56
57  # Sort the build info to ensure deterministic output.
58  sorted_build_info_pairs = sorted(build_info.items())
59
60  contents = """
61# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
62#
63# Licensed under the Apache License, Version 2.0 (the "License");
64# you may not use this file except in compliance with the License.
65# You may obtain a copy of the License at
66#
67#     http://www.apache.org/licenses/LICENSE-2.0
68#
69# Unless required by applicable law or agreed to in writing, software
70# distributed under the License is distributed on an "AS IS" BASIS,
71# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
72# See the License for the specific language governing permissions and
73# limitations under the License.
74# ==============================================================================
75\"\"\"Auto-generated module providing information about the build.\"\"\"
76import collections
77
78build_info = collections.OrderedDict(%s)
79""" % sorted_build_info_pairs
80  open(filename, "w").write(contents)
81
82
83parser = argparse.ArgumentParser(
84    description="""Build info injection into the PIP package.""")
85
86parser.add_argument("--raw_generate", type=str, help="Generate build_info.py")
87
88parser.add_argument(
89    "--key_value", type=str, nargs="*", help="List of key=value pairs.")
90
91args = parser.parse_args()
92
93if args.raw_generate:
94  write_build_info(args.raw_generate, args.key_value)
95else:
96  raise RuntimeError("--raw_generate must be used.")
97