diff options
Diffstat (limited to 'setup.py')
| -rw-r--r-- | setup.py | 43 |
1 files changed, 29 insertions, 14 deletions
@@ -18,6 +18,12 @@ argparser = argparse.ArgumentParser() argparser.add_argument( "--tensorflow_src_dir", help="TensorFlow source directory path", required=False ) +argparser.add_argument( + "--sanitizer", + help="Build using a sanitizer (choose from asan or ubsan)", + choices=["asan", "ubsan"], + required=False +) args, unknown = argparser.parse_known_args() sys.argv = [sys.argv[0]] + unknown @@ -51,21 +57,30 @@ class BazelBuildExtension(setuptools.command.build_ext.build_ext): ext.tensorflow_version, ) - self.spawn( - [ - "bazel", - "build", - "-c", - "opt", - # FIXME Some of the Bazel targets dependencies we use have - # a 'friends' visibility, check if our Bazel target can be added - # to the 'friends' list. - "--check_visibility=false", - "--override_repository=org_tensorflow=" - + os.path.abspath(tensorflow_src_dir), - ext.bazel_target, + commands = [ + "bazel", + "build" + ] + + if args.sanitizer: + commands += [ + "--config={}".format(args.sanitizer) ] - ) + + commands += [ + # FIXME Some of the Bazel targets dependencies we use have + # a 'friends' visibility, check if our Bazel target can be added + # to the 'friends' list. + "-c", + "opt", + "--check_visibility=false", + "--override_repository=org_tensorflow={}".format( + os.path.abspath(tensorflow_src_dir) + ), + ext.bazel_target + ] + + self.spawn(commands) shared_lib_dest_path = self.get_ext_fullpath(ext.name) shared_lib_dest_dir = os.path.dirname(shared_lib_dest_path) |