1 #!/usr/bin/env python 2 # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 3 # 4 # Licensed under the Apache License, Version 2.0 (the "License"); 5 # you may not use this file except in compliance with the License. 6 # You may obtain a copy of the License at 7 # 8 # http://www.apache.org/licenses/LICENSE-2.0 9 # 10 # Unless required by applicable law or agreed to in writing, software 11 # distributed under the License is distributed on an "AS IS" BASIS, 12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 # See the License for the specific language governing permissions and 14 # limitations under the License. 15 # ============================================================================== 16 """Generates a Python module containing information about the build.""" 17 from __future__ import absolute_import 18 from __future__ import division 19 from __future__ import print_function 20 import argparse 21 22 23 def write_build_info(filename, build_config, key_value_list): 24 """Writes a Python that describes the build. 25 26 Args: 27 filename: filename to write to. 28 build_config: A string that represents the config used in this build (e.g. 29 "cuda"). 30 key_value_list: A list of "key=value" strings that will be added to the 31 module as additional fields. 32 33 Raises: 34 ValueError: If `key_value_list` includes the key "is_cuda_build", which 35 would clash with one of the default fields. 36 """ 37 module_docstring = "\"\"\"Generates a Python module containing information " 38 module_docstring += "about the build.\"\"\"" 39 if build_config == "cuda": 40 build_config_bool = "True" 41 else: 42 build_config_bool = "False" 43 44 key_value_pair_stmts = [] 45 if key_value_list: 46 for arg in key_value_list: 47 key, value = arg.split("=") 48 if key == "is_cuda_build": 49 raise ValueError("The key \"is_cuda_build\" cannot be passed as one of " 50 "the --key_value arguments.") 51 key_value_pair_stmts.append("%s = %r" % (key, value)) 52 key_value_pair_content = "\n".join(key_value_pair_stmts) 53 54 contents = """ 55 # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 56 # 57 # Licensed under the Apache License, Version 2.0 (the "License"); 58 # you may not use this file except in compliance with the License. 59 # You may obtain a copy of the License at 60 # 61 # http://www.apache.org/licenses/LICENSE-2.0 62 # 63 # Unless required by applicable law or agreed to in writing, software 64 # distributed under the License is distributed on an "AS IS" BASIS, 65 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 66 # See the License for the specific language governing permissions and 67 # limitations under the License. 68 # ============================================================================== 69 %s 70 from __future__ import absolute_import 71 from __future__ import division 72 from __future__ import print_function 73 74 is_cuda_build = %s 75 76 %s 77 """ % (module_docstring, build_config_bool, key_value_pair_content) 78 open(filename, "w").write(contents) 79 80 81 parser = argparse.ArgumentParser( 82 description="""Build info injection into the PIP package.""") 83 84 parser.add_argument( 85 "--build_config", 86 type=str, 87 help="Either 'cuda' for GPU builds or 'cpu' for CPU builds.") 88 89 parser.add_argument("--raw_generate", type=str, help="Generate build_info.py") 90 91 parser.add_argument("--key_value", type=str, nargs="*", 92 help="List of key=value pairs.") 93 94 args = parser.parse_args() 95 96 if args.raw_generate is not None and args.build_config is not None: 97 write_build_info(args.raw_generate, args.build_config, args.key_value) 98 else: 99 raise RuntimeError("--raw_generate and --build_config must be used") 100
