|
| 1 | +# Copyright 2019 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 | + |
| 16 | +from __future__ import print_function |
| 17 | + |
| 18 | +import os |
| 19 | +import re |
| 20 | +import sys |
| 21 | +import tensorflow as tf |
| 22 | + |
| 23 | +def write_config(): |
| 24 | + """Retrive compile and link information from tensorflow and write to .bazelrc.""" |
| 25 | + |
| 26 | + cflags = tf.sysconfig.get_compile_flags() |
| 27 | + |
| 28 | + inc_regex = re.compile("^-I") |
| 29 | + opt_regex = re.compile("^-D") |
| 30 | + |
| 31 | + include_list = [] |
| 32 | + opt_list = [] |
| 33 | + |
| 34 | + for arg in cflags: |
| 35 | + if inc_regex.match(arg): |
| 36 | + include_list.append(arg) |
| 37 | + elif opt_regex.match(arg): |
| 38 | + opt_list.append(arg) |
| 39 | + else: |
| 40 | + print("WARNING: Unexpected cflag item {}".format(arg)) |
| 41 | + |
| 42 | + |
| 43 | + if len(include_list) != 1: |
| 44 | + print("ERROR: Expected a single include directory in " + |
| 45 | + "tf.sysconfig.get_compile_flags()") |
| 46 | + exit(1) |
| 47 | + |
| 48 | + |
| 49 | + library_regex = re.compile("^-l") |
| 50 | + libdir_regex = re.compile("^-L") |
| 51 | + |
| 52 | + library_list = [] |
| 53 | + libdir_list = [] |
| 54 | + |
| 55 | + lib = tf.sysconfig.get_link_flags() |
| 56 | + |
| 57 | + for arg in lib: |
| 58 | + if library_regex.match(arg): |
| 59 | + library_list.append(arg) |
| 60 | + elif libdir_regex.match(arg): |
| 61 | + libdir_list.append(arg) |
| 62 | + else: |
| 63 | + print("WARNING: Unexpected link flag item {}".format(arg)) |
| 64 | + |
| 65 | + if len(library_list) != 1 or len(libdir_list) != 1: |
| 66 | + print("ERROR: Expected exactly one lib and one libdir in" + |
| 67 | + "tf.sysconfig.get_link_flags()") |
| 68 | + exit(1) |
| 69 | + |
| 70 | + try: |
| 71 | + |
| 72 | + with open(".bazelrc", "w") as bazel_rc: |
| 73 | + for opt in opt_list: |
| 74 | + bazel_rc.write('build --copt="{}"\n'.format(opt)) |
| 75 | + |
| 76 | + bazel_rc.write('build --action_env TF_HEADER_DIR="{}"\n' |
| 77 | + .format(include_list[0][2:])) |
| 78 | + |
| 79 | + bazel_rc.write('build --action_env TF_SHARED_LIBRARY_DIR="{}"\n' |
| 80 | + .format(libdir_list[0][2:])) |
| 81 | + library_name = library_list[0][2:] |
| 82 | + if library_name.startswith(":"): |
| 83 | + library_name = library_name[1:] |
| 84 | + else: |
| 85 | + library_name = "lib" + library_name + ".so" |
| 86 | + bazel_rc.write('build --action_env TF_SHARED_LIBRARY_NAME="{}"\n' |
| 87 | + .format(library_name)) |
| 88 | + bazel_rc.close() |
| 89 | + except OSError: |
| 90 | + print("ERROR: Writing .bazelrc") |
| 91 | + exit(1) |
| 92 | + |
| 93 | + |
| 94 | +def compile_bazel(): |
| 95 | + write_config() |
| 96 | + |
| 97 | + if os.system('rm -f tensorflow_gcs_config/*.so && bazel build -c dbg //tensorflow_gcs_config:_gcs_config_ops.so && cp bazel-bin/tensorflow_gcs_config/_gcs_config_ops.so tensorflow_gcs_config/') != 0: |
| 98 | + raise Exception('Failed to build C extension.') |
0 commit comments