# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates. # SPDX-License-Identifier: Apache-2.0 import argparse import git import os import re import setuptools import setuptools.command.build_ext import shutil import sys import time TOSA_CHECKER_VERSION = "0.2.0" DEFAULT_TENSORFLOW_VERSION = "2.13.0" # Get the TensorFlow™ source directory passed to the command line (if any). # If none is given the sources are pulled from the official TF repository. 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 ) argparser.add_argument( "--tosa_checker_copt", help="Build tosa_checker with additional copt (comma separated string)", default="", required=False ) argparser.add_argument( "--nightly", help="Build tosa_checker as a nightly wheel", action="store_true" ) args, unknown = argparser.parse_known_args() sys.argv = [sys.argv[0]] + unknown def increment_version(version_string): version_numbers = version_string.split(".") increment = int(version_numbers[1]) + 1 version_numbers[1] = str(increment) return ".".join(version_numbers) def get_package_version(nightly=False): if nightly: return "{}.dev{}".format( increment_version(TOSA_CHECKER_VERSION), time.strftime("%Y%m%d") ) else: return TOSA_CHECKER_VERSION def get_repo_version(repo_directory): r = git.repo.Repo(repo_directory) tag = r.git.tag('--points-at') if tag: return tag else: return r.head.commit.hexsha class BazelExtensionModule(setuptools.Extension): def __init__( self, py_module_name, library_name, bazel_target, bazel_shared_lib_output, tensorflow_version, ): super().__init__(py_module_name, sources=[]) self.library_name = library_name self.bazel_target = bazel_target self.bazel_shared_lib_output = bazel_shared_lib_output self.tensorflow_version = tensorflow_version class BazelBuildExtension(setuptools.command.build_ext.build_ext): """Override build_extension to build the library with bazel and copying it beforehand.""" def build_extension(self, ext): tensorflow_src_dir = args.tensorflow_src_dir if not tensorflow_src_dir: tensorflow_src_dir = os.path.join(self.build_temp, "tensorflow") self._clone_tf_repository( tensorflow_src_dir, ext.tensorflow_version, ) commands = [ "bazel", "build" ] if args.sanitizer: commands += [ "--config={}".format(args.sanitizer) ] if args.tosa_checker_copt: commands += [ "--per_file_copt=tosa_checker/tosa_checker.*@{}".format(args.tosa_checker_copt) ] 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) package_dir = os.path.join(shared_lib_dest_dir, ext.library_name) os.makedirs(shared_lib_dest_dir, exist_ok=True) os.makedirs(package_dir, exist_ok=True) shutil.copyfile(ext.bazel_shared_lib_output, shared_lib_dest_path) # Get the TensorFlow version this is built with tf_version = get_repo_version(tensorflow_src_dir) with open(os.path.join(ext.library_name, "__init__.py"), "r") as f: module_init_file = f.read() module_init_file += "__tensorflow_version__ = \"{}\"\n".format(tf_version) with open(os.path.join(package_dir, "__init__.py"), "w") as f: f.write(module_init_file) super().build_extension(ext) def _clone_tf_repository(self, tensorflow_src_dir, tensorflow_version): if os.path.exists(tensorflow_src_dir): return tensorflow_repo = "https://github.com/tensorflow/tensorflow.git" self.spawn( [ "git", "clone", "--depth=1", "--branch", "v" + tensorflow_version, tensorflow_repo, tensorflow_src_dir, ] ) def get_long_description(nightly): # Read the contents of README.md file this_directory = os.path.abspath(os.path.dirname(__file__)) with open(os.path.join(this_directory, "README.md"), encoding="utf-8") as f: long_description = f.read() if not nightly: # Replace relative links to existing files with absolute links to https://review.mlplatform.org url = f"https://review.mlplatform.org/plugins/gitiles/tosa/tosa_checker/+/refs/tags/{TOSA_CHECKER_VERSION}/" # Find all markdown links that match the format: [text](link) for match, link in re.findall(r"(\[.+?\]\((.+?)\))", long_description): # If the link is a file that exists, replace it with the web link to the file instead if os.path.exists(os.path.join(this_directory, link)): url_link = re.sub(r"\((.+?)\)", rf"({url}{link})", match) long_description = long_description.replace(match, url_link) return long_description setuptools.setup( name="tosa-checker", version=get_package_version(args.nightly), description="Tool to check if a ML model is compatible with the TOSA specification", long_description=get_long_description(args.nightly), long_description_content_type="text/markdown", author="Arm Limited", url="https://git.mlplatform.org/tosa/tosa_checker.git/", license="Apache-2.0", python_requires=">=3.8", cmdclass={"build_ext": BazelBuildExtension}, ext_modules=[ BazelExtensionModule( py_module_name="_tosa_checker_wrapper", library_name="tosa_checker", bazel_target="//tosa_checker:tosa_checker", bazel_shared_lib_output="bazel-bin/tosa_checker/_tosa_checker_wrapper.so", tensorflow_version=DEFAULT_TENSORFLOW_VERSION, ), ], classifiers=[ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3", "Topic :: Utilities", ], )