setup.py

# -*- coding: utf-8 -*-
# Copyright (C) 2018-2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import os.path
import sys
import errno
import subprocess  # nosec
import typing
import platform
import multiprocessing
from fnmatch import fnmatchcase
from pathlib import Path
from shutil import copyfile, rmtree
from setuptools import setup, find_namespace_packages, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.build_clib import build_clib
from setuptools.command.install import install
from distutils.command.build import build
from distutils.command.clean import clean
from distutils.errors import DistutilsSetupError
from distutils.file_util import copy_file
from distutils import log

WHEEL_LIBS_INSTALL_DIR = os.path.join("openvino", "libs")
WHEEL_LIBS_PACKAGE = "openvino.libs"
PYTHON_VERSION = f"python{sys.version_info.major}.{sys.version_info.minor}"

LIBS_DIR = "bin" if platform.system() == "Windows" else "lib"
CONFIG = "Release" if platform.system() in {"Windows", "Darwin"} else ""

machine = platform.machine()
if machine == "x86_64" or machine == "AMD64":
    ARCH = "intel64"
elif machine == "X86" or machine == "i686":
    ARCH = "ia32"
elif machine == "arm" or machine == "armv7l":
    ARCH = "arm"
elif machine == "aarch64" or machine == "arm64" or machine == "ARM64":
    ARCH = "arm64"

# The following variables can be defined in environment or .env file
SCRIPT_DIR = Path(__file__).resolve().parents[0]
CMAKE_BUILD_DIR = os.getenv("CMAKE_BUILD_DIR", ".")
OPENVINO_BUILD_DIR = os.getenv("OPENVINO_BUILD_DIR", CMAKE_BUILD_DIR)
OPENVINO_PYTHON_BUILD_DIR = os.getenv("OPENVINO_PYTHON_BUILD_DIR", CMAKE_BUILD_DIR)
OV_RUNTIME_LIBS_DIR = os.getenv("OV_RUNTIME_LIBS_DIR", f"runtime/{LIBS_DIR}/{ARCH}/{CONFIG}")
TBB_LIBS_DIR = os.getenv("TBB_LIBS_DIR", f"runtime/3rdparty/tbb/{LIBS_DIR}")
PUGIXML_LIBS_DIR = os.getenv("PUGIXML_LIBS_DIR", f"runtime/3rdparty/pugixml/{LIBS_DIR}")
PY_PACKAGES_DIR = os.getenv("PY_PACKAGES_DIR", f"python/{PYTHON_VERSION}")
LIBS_RPATH = "$ORIGIN" if sys.platform == "linux" else "@loader_path"

LIB_INSTALL_CFG = {
    "ie_libs": {
        "name": "core",
        "prefix": "libs.core",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "hetero_plugin": {
        "name": "hetero",
        "prefix": "libs.hetero",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "gpu_plugin": {
        "name": "gpu",
        "prefix": "libs.gpu",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "cpu_plugin": {
        "name": "cpu",
        "prefix": "libs.cpu",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "multi_plugin": {
        "name": "multi",
        "prefix": "libs.multi",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "batch_plugin": {
        "name": "batch",
        "prefix": "libs.batch",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "tbb_libs": {
        "name": "tbb",
        "prefix": "libs.tbb",
        "install_dir": TBB_LIBS_DIR,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "pugixml_libs": {
        "name": "pugixml",
        "prefix": "libs.pugixml",
        "install_dir": PUGIXML_LIBS_DIR,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "ir_libs": {
        "name": "ir",
        "prefix": "libs.ir",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "paddle_libs": {
        "name": "paddle",
        "prefix": "libs.paddle",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "pytorch_libs": {
        "name": "pytorch",
        "prefix": "libs.pytorch",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "onnx_libs": {
        "name": "onnx",
        "prefix": "libs.onnx",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "tensorflow_libs": {
        "name": "tensorflow",
        "prefix": "libs.tensorflow",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
    "tensorflow_lite_libs": {
        "name": "tensorflow_lite",
        "prefix": "libs.tensorflow_lite",
        "install_dir": OV_RUNTIME_LIBS_DIR,
        "rpath": LIBS_RPATH,
        "binary_dir": OPENVINO_BUILD_DIR,
    },
}

PY_INSTALL_CFG = {
    "ie_py": {
        "name": f"pyie_{PYTHON_VERSION}",
        "prefix": "site-packages",
        "install_dir": PY_PACKAGES_DIR,
        "binary_dir": OPENVINO_PYTHON_BUILD_DIR,
    },
    "ngraph_py": {
        "name": f"pyngraph_{PYTHON_VERSION}",
        "prefix": "site-packages",
        "install_dir": PY_PACKAGES_DIR,
        "binary_dir": OPENVINO_PYTHON_BUILD_DIR,
    },
    "pyopenvino": {
        "name": f"pyopenvino_{PYTHON_VERSION}",
        "prefix": "site-packages",
        "install_dir": PY_PACKAGES_DIR,
        "binary_dir": OPENVINO_PYTHON_BUILD_DIR,
    },
}


class PrebuiltExtension(Extension):
    """Initialize Extension."""

    def __init__(self, name, sources, *args, **kwargs):
        if len(sources) != 1:
            nln = "\n"
            raise DistutilsSetupError(f"PrebuiltExtension can accept only one source, but got: {nln}{nln.join(sources)}")
        super().__init__(name, sources, *args, **kwargs)


class CustomBuild(build):
    """Custom implementation of build_clib."""

    cmake_build_types = ["Release", "Debug", "RelWithDebInfo", "MinSizeRel"]
    user_options = [
        ("config=", None, "Build configuration [{types}].".format(types="|".join(cmake_build_types))),
        ("jobs=", None, "Specifies the number of jobs to use with make."),
        ("cmake-args=", None, "Additional options to be passed to CMake."),
    ]

    def initialize_options(self):
        """Set default values for all the options that this command supports."""
        super().initialize_options()
        self.build_base = "build"
        self.config = None
        self.jobs = None
        self.cmake_args = None

    def finalize_options(self):
        """Set final values for all the options that this command supports."""
        super().finalize_options()

        if not self.config:
            if self.debug:
                self.config = "Debug"
            else:
                self.announce("Set default value for CMAKE_BUILD_TYPE = Release.", level=4)
                self.config = "Release"
        else:
            build_types = [item.lower() for item in self.cmake_build_types]
            try:
                i = build_types.index(str(self.config).lower())
                self.config = self.cmake_build_types[i]
                self.debug = True if "Debug" == self.config else False
            except ValueError:
                self.announce("Unsupported CMAKE_BUILD_TYPE value: " + self.config, level=4)
                self.announce("Supported values: {types}".format(types=", ".join(self.cmake_build_types)), level=4)
                sys.exit(1)
        if self.jobs is None and os.getenv("MAX_JOBS") is not None:
            self.jobs = os.getenv("MAX_JOBS")
        self.jobs = multiprocessing.cpu_count() if self.jobs is None else int(self.jobs)

    def run(self):
        global OPENVINO_BUILD_DIR
        self.jobs = multiprocessing.cpu_count()
        plat_specifier = ".{0}-{1}.{2}".format(self.plat_name, *sys.version_info[:2])
        self.build_temp = os.path.join(self.build_base, "temp" + plat_specifier, self.config)
        self.announce(f"Create build directory: {self.build_temp}", level=3)

        # if setup.py is directly called use CMake to build product
        if OPENVINO_BUILD_DIR == ".":
            # set path to the root of OpenVINO CMakeList file
            openvino_root_dir = Path(__file__).resolve().parents[4]
            self.announce(f"Configuring cmake project: {openvino_root_dir}", level=3)
            self.spawn(["cmake", "-S" + str(openvino_root_dir),
                                 "-B" + self.build_temp,
                                 "-DCMAKE_BUILD_TYPE={type}".format(type=self.config),
                                 "-DENABLE_PYTHON=ON",
                                 "-DENABLE_NCC_STYLE=OFF",
                                 "-DENABLE_CPPLINT=OFF",
                                 "-DENABLE_TEMPLATE=OFF",
                                 "-DENABLE_SAMPLES=OFF"])

            self.announce("Building binaries", level=3)
            self.spawn(["cmake", "--build", self.build_temp,
                                 "--config", self.config,
                                 "-j", str(self.jobs)])
            OPENVINO_BUILD_DIR = self.build_temp
        self.run_command("build_clib")

        build.run(self)
        # Copy extra package_data content filtered by find_packages
        dst = Path(self.build_lib)
        src = Path(get_package_dir(PY_INSTALL_CFG))
        exclude = ignore_patterns("*ez_setup*", "*__pycache__*", "*.egg-info*")
        for path in src.glob("**/*"):
            if path.is_dir() or exclude(str(path)):
                continue
            path_rel = path.relative_to(src)
            (dst / path_rel.parent).mkdir(exist_ok=True, parents=True)
            copyfile(path, dst / path_rel)


class PrepareLibs(build_clib):
    """Prepare prebuilt libraries."""

    def run(self):
        self.configure(LIB_INSTALL_CFG)
        self.configure(PY_INSTALL_CFG)
        self.generate_package(get_install_dirs_list(LIB_INSTALL_CFG))

    def configure(self, install_cfg):
        """Collect prebuilt libraries. Install them to the temp directories, set rpath."""
        for comp, comp_data in install_cfg.items():
            install_prefix = comp_data.get("prefix")
            install_dir = comp_data.get("install_dir")
            binary_dir = comp_data.get("binary_dir")
            if install_dir and not os.path.isabs(install_dir):
                self.announce(f"Installing {comp}", level=3)
                self.spawn(["cmake", "--install", binary_dir,
                                     "--prefix", install_prefix,
                                     "--config", "Release",
                                     "--strip",
                                     "--component", comp_data.get("name")])
                install_dir = os.path.join(install_prefix, install_dir)
            # set rpath if applicable
            if sys.platform != "win32" and comp_data.get("rpath"):
                for path in filter(
                    lambda x: any(item in ([".so"] if sys.platform == "linux" else [".dylib", ".so"])
                                  for item in x.suffixes), Path(install_dir).glob("*"),
                ):
                    set_rpath(comp_data["rpath"], os.path.realpath(path))

    def generate_package(self, src_dirs):
        """Collect package data files from preinstalled dirs and put all runtime libraries to the subpackage."""
        # additional blacklist filter, just to fix cmake install issues
        blacklist = [".lib", ".pdb", "_debug.dll", "_debug.dylib"]
        package_dir = os.path.join(get_package_dir(PY_INSTALL_CFG), WHEEL_LIBS_INSTALL_DIR)

        for src_dir in src_dirs:
            local_base_dir = Path(src_dir)

            # skip symlinks of higher level like libX.so or libX.dylib
            for symlink in local_base_dir.rglob("*"):
                if symlink.is_symlink():
                    file_name = os.readlink(symlink)
                    if not os.path.isabs(file_name):
                        file_name = os.path.join(os.path.dirname(symlink), file_name)
                    if Path(file_name).is_symlink():
                        self.announce(f"Unlink symlink {symlink}, use {file_name} instead", level=3)
                        os.unlink(symlink)

            # transform libX.so.Y / libX.Y.dylib symlinks to real files
            for symlink in local_base_dir.rglob("*"):
                if symlink.is_symlink():
                    file_name = os.readlink(symlink)
                    if not os.path.isabs(file_name):
                        file_name = os.path.join(os.path.dirname(symlink), file_name)

                    os.unlink(symlink)
                    os.rename(file_name, symlink)
                    self.announce(f"Resolved symlink {symlink} as {file_name}", level=3)

            # copy so / dylib files to WHEEL_LIBS_INSTALL_DIR
            for file_path in local_base_dir.rglob("*"):
                file_name = os.path.basename(file_path)
                if file_path.is_symlink():
                    sys.exit(f"Wheel package content must not contain symlinks {file_path}")
                if file_path.is_file() and not any(file_name.endswith(ext) for ext in blacklist):
                    dst_file = os.path.join(package_dir, os.path.relpath(file_path, local_base_dir))
                    os.makedirs(os.path.dirname(dst_file), exist_ok=True)
                    copyfile(file_path, dst_file)
                    self.announce(f"Copy {file_path} to {dst_file}", level=3)

        if Path(package_dir).exists():
            self.announce(f"Adding {WHEEL_LIBS_PACKAGE} package", level=3)
            packages.append(WHEEL_LIBS_PACKAGE)
            package_data.update({WHEEL_LIBS_PACKAGE: ["*"]})


class CopyExt(build_ext):
    """Copy extension files to the build directory."""

    def run(self):
        if len(self.extensions) == 1:
            self.run_command("build_clib")
            self.extensions = []
            self.extensions = find_prebuilt_extensions(get_install_dirs_list(PY_INSTALL_CFG))
        for extension in self.extensions:
            if not isinstance(extension, PrebuiltExtension):
                raise DistutilsSetupError(f"copy_ext can accept PrebuiltExtension only, but got {extension.name}")
            src = extension.sources[0]
            dst = self.get_ext_fullpath(extension.name)
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            # setting relative path to find dlls
            if sys.platform != "win32":
                rpath = os.path.relpath(get_package_dir(PY_INSTALL_CFG), os.path.dirname(src))
                rpath = os.path.join(LIBS_RPATH, rpath, WHEEL_LIBS_INSTALL_DIR)
                set_rpath(rpath, os.path.realpath(src))

            copy_file(src, dst, verbose=self.verbose, dry_run=self.dry_run)


class CustomInstall(install):
    """Enable build_clib during the installation."""

    def run(self):
        self.run_command("build")
        install.run(self)


class CustomClean(clean):
    """Clean up staging directories."""

    def clean(self, install_cfg):
        for comp, comp_data in install_cfg.items():
            install_prefix = comp_data.get("prefix")
            self.announce(f"Cleaning {comp}: {install_prefix}", level=3)
            if os.path.exists(install_prefix):
                rmtree(install_prefix)

    def run(self):
        self.clean(LIB_INSTALL_CFG)
        self.clean(PY_INSTALL_CFG)
        clean.run(self)


def ignore_patterns(*patterns):
    """Filter names by given patterns."""
    return lambda name: any(fnmatchcase(name, pat=pat) for pat in patterns)


def is_tool(name):
    """Check if the command-line tool is available."""
    try:
        devnull = subprocess.DEVNULL
        subprocess.Popen([name], stdout=devnull, stderr=devnull).communicate()  # nosec
    except OSError as error:
        if error.errno == errno.ENOENT:
            return False
    return True


def remove_rpath(file_path):
    """Remove rpath from binaries.

    :param file_path: binary path
    :type file_path: pathlib.Path
    """
    if sys.platform == "darwin":
        cmd = (
            f"otool -l {file_path} "  # noqa: P103
            f"| grep LC_RPATH -A3 "
            f"| grep -o 'path.*' "
            f"| cut -d ' ' -f2 "
            f"| xargs -I{{}} install_name_tool -delete_rpath {{}} {file_path}"
        )
        if os.WEXITSTATUS(os.system(cmd)) != 0:  # nosec
            sys.exit(f"Could not remove rpath for {file_path}")
    else:
        sys.exit(f"Unsupported platform: {sys.platform}")


def set_rpath(rpath, executable):
    """Setting rpath for linux and macOS libraries."""
    print(f"Setting rpath {rpath} for {executable}")  # noqa: T001, T201
    cmd = []
    rpath_tool = ""

    if sys.platform == "linux":
        with open(os.path.realpath(executable), "rb") as file:
            if file.read(1) != b"\x7f":
                log.warn(f"WARNING: {executable}: missed ELF header")
                return
        rpath_tool = "patchelf"
        cmd = [rpath_tool, "--set-rpath", rpath, executable]
    elif sys.platform == "darwin":
        rpath_tool = "install_name_tool"
        cmd = [rpath_tool, "-add_rpath", rpath, executable]
    else:
        sys.exit(f"Unsupported platform: {sys.platform}")

    if is_tool(rpath_tool):
        if sys.platform == "darwin":
            remove_rpath(executable)
        ret_info = subprocess.run(cmd, check=True, shell=False)  # nosec
        if ret_info.returncode != 0:
            sys.exit(f"Could not set rpath: {rpath} for {executable}")
    else:
        sys.exit(f"Could not found {rpath_tool} on the system, " f"please make sure that this tool is installed")


def find_prebuilt_extensions(search_dirs):
    """Collect prebuilt python extensions."""
    extensions = []
    ext_pattern = ""
    if sys.platform == "linux":
        ext_pattern = "**/*.so"
    elif sys.platform == "win32":
        ext_pattern = "**/*.pyd"
    elif sys.platform == "darwin":
        ext_pattern = "**/*.so"
    for base_dir in search_dirs:
        for path in Path(base_dir).glob(ext_pattern):
            if path.match("openvino/libs/*") or path.match(f"openvino/libs/openvino-{OPENVINO_VERSION}/*"):
                continue
            relpath = path.relative_to(base_dir)
            if relpath.parent != ".":
                package_names = str(relpath.parent).split(os.path.sep)
            else:
                package_names = []
            package_names.append(path.name.split(".", 1)[0])
            name = ".".join(package_names)
            extensions.append(PrebuiltExtension(name, sources=[str(path)]))
    if not extensions:
        extensions.append(PrebuiltExtension("openvino", sources=[str("setup.py")]))
    return extensions


def get_description(desc_file_path):
    """Read description from README.md."""
    with open(desc_file_path, "r", encoding="utf-8") as fstream:
        description = fstream.read()
    return description


def get_dependencies(requirements_file_path):
    """Read dependencies from requirements.txt."""
    with open(requirements_file_path, "r", encoding="utf-8") as fstream:
        dependencies = fstream.read()
    return dependencies


def get_install_dirs_list(install_cfg):
    """Collect all available directories with libs or python packages."""
    dirs = []
    for comp_info in install_cfg.values():
        cfg_prefix = comp_info.get("prefix")
        cfg_dir = comp_info.get("install_dir")
        if cfg_dir:
            if not os.path.isabs(cfg_dir):
                cfg_dir = os.path.join(cfg_prefix, cfg_dir)
            if cfg_dir not in dirs:
                dirs.append(cfg_dir)
    return dirs


def get_package_dir(install_cfg):
    """Get python package path based on config. All the packages should be located in one directory."""
    py_package_path = ""
    dirs = get_install_dirs_list(install_cfg)
    if len(dirs) != 0:
        # setup.py support only one package directory, all modules should be located there
        py_package_path = dirs[0]
    return py_package_path


def concat_files(output_file, input_files):
    """Concatenates multuple input files to a single output file."""
    with open(output_file, "w", encoding="utf-8") as outfile:
        for filename in input_files:
            with open(filename, "r", encoding="utf-8") as infile:
                content = infile.read()
                outfile.write(content)
    return output_file


platforms = ["linux", "win32", "darwin"]
if not any(pl in sys.platform for pl in platforms):
    sys.exit(f"Unsupported platform: {sys.platform}, expected: linux, win32, darwin")

OPENVINO_VERSION = WHEEL_VERSION = os.getenv("WHEEL_VERSION", "0.0.0")
# copy license file into the build directory
package_license = os.getenv("WHEEL_LICENSE", SCRIPT_DIR.parents[3] / "LICENSE")
if os.path.exists(package_license):
    copyfile(package_license, "LICENSE")

packages = find_namespace_packages(get_package_dir(PY_INSTALL_CFG))
package_data: typing.Dict[str, list] = {}
pkg_name = os.getenv("WHEEL_PACKAGE_NAME", "openvino")
ext_modules = find_prebuilt_extensions(get_install_dirs_list(PY_INSTALL_CFG)) if pkg_name == "openvino" else []

description_md = SCRIPT_DIR.parents[3] / "docs" / "install_guides" / "pypi-openvino-rt.md"
md_files = [description_md, SCRIPT_DIR.parents[3] / "docs" / "install_guides" / "pre-release-note.md"]
docs_url = "https://docs.openvino.ai/latest/index.html"

if (os.getenv("CI_BUILD_DEV_TAG")):
    output = Path.cwd() / "build" / "pypi-openvino-rt.md"
    output.parent.mkdir(exist_ok=True)
    description_md = concat_files(output, md_files)
    docs_url = "https://docs.openvino.ai/nightly/index.html"
    OPENVINO_VERSION = WHEEL_VERSION[0:8]

setup(
    version=WHEEL_VERSION,
    build=os.getenv("WHEEL_BUILD", "000"),
    author_email=os.getenv("WHEEL_AUTHOR_EMAIL", "openvino_pushbot@intel.com"),
    name=pkg_name,
    license=os.getenv("WHEEL_LICENCE_TYPE", "OSI Approved :: Apache Software License"),
    author=os.getenv("WHEEL_AUTHOR", "Intel(R) Corporation"),
    description=os.getenv("WHEEL_DESC", "OpenVINO(TM) Runtime"),
    install_requires=get_dependencies(os.getenv("WHEEL_REQUIREMENTS", SCRIPT_DIR.parents[0] / "requirements.txt")),
    long_description=get_description(os.getenv("WHEEL_OVERVIEW", description_md)),
    long_description_content_type="text/markdown",
    download_url=os.getenv("WHEEL_DOWNLOAD_URL", "https://github.com/openvinotoolkit/openvino/tags"),
    url=os.getenv("WHEEL_URL", docs_url),
    cmdclass={
        "build": CustomBuild,
        "install": CustomInstall,
        "build_clib": PrepareLibs,
        "build_ext": CopyExt,
        "clean": CustomClean,
    },
    ext_modules=ext_modules,
    packages=packages,
    package_dir={"": get_package_dir(PY_INSTALL_CFG)},
    package_data=package_data,
    zip_safe=False,
)