google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
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
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
17
#
18
# Licensed under the Apache License, Version 2.0 (the "License");
19
# you may not use this file except in compliance with the License.
20
# You may obtain a copy of the License at
21
#
22
#     http://www.apache.org/licenses/LICENSE-2.0
23
#
24
# Unless required by applicable law or agreed to in writing, software
25
# distributed under the License is distributed on an "AS IS" BASIS,
26
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
27
# See the License for the specific language governing permissions and
28
# limitations under the License.
29
# ==============================================================================
30
# Usage: configure.py [--quiet] [--no-deps]
31
#
32
# Options:
33
#  --quiet  Give less output.
34
#  --no-deps  Don't install Python dependencies
35
"""Configures Gigamol to be built from source."""
36

37
import argparse
38
import logging
39
import os
40
import subprocess
41
import sys
42
from typing import List
43

44
_BAZELRC = ".bazelrc"
45
_BAZEL_QUERY = ".bazel-query.sh"
46

47

48
# Writes variables to bazelrc file
49
def write_to_bazelrc(line):
50
  with open(_BAZELRC, "a") as f:
51
    f.write(line + "\n")
52

53

54
def write_action_env(var_name, var):
55
  write_to_bazelrc('build --action_env %s="%s"' % (var_name, str(var)))
56
  with open(_BAZEL_QUERY, "a") as f:
57
    f.write('{}="{}" '.format(var_name, var))
58

59

60
def generate_shared_lib_name(tf_lflags):
61
  """Converts the linkflag namespec to the full shared library name.
62

63
  Args:
64
    tf_lflags: List of linkflag namespec. The first entry specifies the
65
      directory containing the TensorFlow framework library. The second entry
66
      specifies the name of the Tensorflow shared lib (For Linux,
67
      '-l:libtensorflow_framework.so.%s' % version).
68

69
  Returns:
70
    Name of the Tensorflow shared lib.
71
  """
72
  # Assume Linux for now
73
  return tf_lflags[1][3:]
74

75

76
def create_build_configuration():
77
  """Main function to create build configuration."""
78
  print()
79
  print("Configuring Gigamol to be built from source...")
80

81
  pip_install_options = ["--upgrade"]
82
  parser = argparse.ArgumentParser()
83
  parser.add_argument("--quiet", action="store_true", help="Give less output.")
84
  parser.add_argument(
85
      "--no-deps",
86
      action="store_true",
87
      help="Do not check and install Python dependencies.",
88
  )
89
  args = parser.parse_args()
90
  if args.quiet:
91
    pip_install_options.append("--quiet")
92

93
  python_path = sys.executable
94
  with open("requirements.txt") as f:
95
    required_packages = f.read().splitlines()
96

97
  print()
98
  if args.no_deps:
99
    print("> Using pre-installed Tensorflow.")
100
  else:
101
    print("> Installing", required_packages)
102
    install_cmd = [python_path, "-m", "pip", "install"]
103
    install_cmd.extend(pip_install_options)
104
    install_cmd.extend(required_packages)
105
    subprocess.check_call(install_cmd)
106

107
  if os.path.isfile(_BAZELRC):
108
    os.remove(_BAZELRC)
109
  if os.path.isfile(_BAZEL_QUERY):
110
    os.remove(_BAZEL_QUERY)
111

112
  logging.disable(logging.WARNING)
113

114
  import tensorflow.compat.v2 as tf  # pylint: disable=g-import-not-at-top
115

116
  # pylint: disable=invalid-name
117
  _TF_CFLAGS = tf.sysconfig.get_compile_flags()
118
  _TF_LFLAGS = tf.sysconfig.get_link_flags()
119
  _TF_CXX11_ABI_FLAG = tf.sysconfig.CXX11_ABI_FLAG
120

121
  _TF_SHARED_LIBRARY_NAME = generate_shared_lib_name(_TF_LFLAGS)
122
  _TF_HEADER_DIR = _TF_CFLAGS[0][2:]
123
  _TF_SHARED_LIBRARY_DIR = _TF_LFLAGS[0][2:]
124
  # pylint: enable=invalid-name
125

126
  write_action_env("TF_HEADER_DIR", _TF_HEADER_DIR)
127
  write_action_env("TF_SHARED_LIBRARY_DIR", _TF_SHARED_LIBRARY_DIR)
128
  write_action_env("TF_SHARED_LIBRARY_NAME", _TF_SHARED_LIBRARY_NAME)
129
  write_action_env("TF_CXX11_ABI_FLAG", _TF_CXX11_ABI_FLAG)
130

131
  write_to_bazelrc("build --spawn_strategy=standalone")
132
  write_to_bazelrc("build --strategy=Genrule=standalone")
133
  write_to_bazelrc("build -c opt")
134

135
  print()
136
  print("Build configurations successfully written to", _BAZELRC)
137
  print()
138

139
  with open(_BAZEL_QUERY, "a") as f:
140
    f.write('bazel query "$@"')
141

142

143
if __name__ == "__main__":
144
  create_build_configuration()
145