Clean up build.py

.gitignore

@@ -195,4 +195,4 @@ Package.pins
 Package.resolved
 .build/
 .swiftpm/
-repros/
+repros/

tools/python/util/__init__.py

@@ -5,6 +5,28 @@
 from .logger import get_logger
 from .platform_helpers import is_linux, is_macOS, is_windows  # noqa: F401
 from .run import run  # noqa: F401
+from .build_helpers import (
+    version_to_tuple,
+    check_python_version,
+    str_to_bool,
+    is_reduced_ops_build,
+    resolve_executable_path,
+    get_config_build_dir,
+    use_dev_mode,
+    add_default_definition,
+    normalize_arg_list,
+    number_of_parallel_jobs,
+    number_of_nvcc_threads,
+    setup_cann_vars,
+    setup_tensorrt_vars,
+    setup_migraphx_vars,
+    setup_cuda_vars,
+
+)
+from .base_error import BaseError
+from .usage_error import UsageError
+from .build_error import BuildError
+from .open_vino_utils import openvino_verify_device_type
 
 try:
     import flatbuffers  # noqa: F401

tools/python/util/base_error.py

@@ -0,0 +1,6 @@
+#  // Copyright (c) Microsoft Corporation. All rights reserved.
+#  // Licensed under the MIT License.
+class BaseError(Exception):
+    """Base class for errors originating from build.py."""
+
+    pass

tools/python/util/build_error.py

@@ -0,0 +1,10 @@
+#  // Copyright (c) Microsoft Corporation. All rights reserved.
+#  // Licensed under the MIT License.
+from .base_error import BaseError
+
+
+class BuildError(BaseError):
+    """Error from running build steps."""
+
+    def __init__(self, *messages):
+        super().__init__("\n".join(messages))

tools/python/util/build_helpers.py

@@ -0,0 +1,203 @@
+#  // Copyright (c) Microsoft Corporation. All rights reserved.
+#  // Licensed under the MIT License.
+import contextlib
+import sys
+import shutil
+import os
+from .UsageError import UsageError
+from .BuildError import BuildError
+from .platform_helpers import is_windows, is_macOS
+
+
+def version_to_tuple(version: str) -> tuple:
+    v = []
+    for s in version.split("."):
+        with contextlib.suppress(ValueError):
+            v.append(int(s))
+    return tuple(v)
+
+
+def check_python_version():
+    required_minor_version = 7
+    if (sys.version_info.major, sys.version_info.minor) < (3, required_minor_version):
+        raise UsageError(
+            f"Invalid Python version. At least Python 3.{required_minor_version} is required. "
+            f"Actual Python version: {sys.version}"
+        )
+
+
+def str_to_bool(s):
+    """Convert string to bool (in argparse context)."""
+    if s.lower() not in ["true", "false"]:
+        raise ValueError("Need bool; got %r" % s)
+    return {"true": True, "false": False}[s.lower()]
+
+
+def is_reduced_ops_build(args):
+    return args.include_ops_by_config is not None
+
+
+def resolve_executable_path(command_or_path):
+    """Returns the absolute path of an executable."""
+    if command_or_path and command_or_path.strip():
+        executable_path = shutil.which(command_or_path)
+        if executable_path is None:
+            raise BuildError(f"Failed to resolve executable path for '{command_or_path}'.")
+        return os.path.abspath(executable_path)
+    else:
+        return None
+
+
+def get_config_build_dir(build_dir, config):
+    # build directory per configuration
+    return os.path.join(build_dir, config)
+
+
+def use_dev_mode(args):
+    if args.compile_no_warning_as_error:
+        return False
+    if args.use_acl:
+        return False
+    if args.use_armnn:
+        return False
+    if args.ios and is_macOS():
+        return False
+    SYSTEM_COLLECTIONURI = os.getenv("SYSTEM_COLLECTIONURI")  # noqa: N806
+    if SYSTEM_COLLECTIONURI and SYSTEM_COLLECTIONURI != "https://dev.azure.com/onnxruntime/":
+        return False
+    return True
+
+
+def add_default_definition(definition_list, key, default_value):
+    for x in definition_list:
+        if x.startswith(key + "="):
+            return definition_list
+    definition_list.append(key + "=" + default_value)
+
+
+def normalize_arg_list(nested_list):
+    return [i for j in nested_list for i in j] if nested_list else []
+
+
+def number_of_parallel_jobs(args):
+    return os.cpu_count() if args.parallel == 0 else args.parallel
+
+
+def number_of_nvcc_threads(args):
+    if args.nvcc_threads >= 0:
+        return args.nvcc_threads
+
+    nvcc_threads = 1
+    import psutil
+    try:
+        available_memory = psutil.virtual_memory().available
+        if isinstance(available_memory, int) and available_memory > 0:
+            if available_memory > 60 * 1024 * 1024 * 1024:
+                # When available memory is large enough, chance of OOM is small.
+                nvcc_threads = 4
+            else:
+                # NVCC need a lot of memory to compile 8 flash attention cu files in Linux or 4 cutlass fmha cu files
+                # in Windows. Here we select number of threads to ensure each thread has enough memory (>= 4 GB). For
+                # example, Standard_NC4as_T4_v3 has 4 CPUs and 28 GB memory. When parallel=4 and nvcc_threads=2,
+                # total nvcc threads is 4 * 2, which is barely able to build in 28 GB memory so we will use
+                # nvcc_threads=1.
+                memory_per_thread = 4 * 1024 * 1024 * 1024
+                fmha_cu_files = 4 if is_windows() else 16
+                fmha_parallel_jobs = min(fmha_cu_files, number_of_parallel_jobs(args))
+                nvcc_threads = max(1, int(available_memory / (memory_per_thread * fmha_parallel_jobs)))
+                print(
+                    f"nvcc_threads={nvcc_threads} to ensure memory per thread >= 4GB for available_memory="
+                    f"{available_memory} and fmha_parallel_jobs={fmha_parallel_jobs}"
+                )
+    except ImportError:
+        print(
+            "Failed to import psutil. Please `pip install psutil` for better estimation of nvcc threads. Use "
+            "nvcc_threads=1"
+        )
+
+    return nvcc_threads
+
+
+def setup_cann_vars(args):
+    cann_home = ""
+
+    if args.use_cann:
+        cann_home = args.cann_home if args.cann_home else os.getenv("ASCEND_HOME_PATH")
+
+        cann_home_valid = cann_home is not None and os.path.exists(cann_home)
+
+        if not cann_home_valid:
+            raise BuildError(
+                "cann_home paths must be specified and valid.",
+                f"cann_home='{cann_home}' valid={cann_home_valid}.",
+            )
+
+    return cann_home
+
+
+def setup_tensorrt_vars(args):
+    tensorrt_home = ""
+    if args.use_tensorrt:
+        tensorrt_home = args.tensorrt_home if args.tensorrt_home else os.getenv("TENSORRT_HOME")
+        tensorrt_home_valid = tensorrt_home is not None and os.path.exists(tensorrt_home)
+        if not tensorrt_home_valid:
+            raise BuildError(
+                "tensorrt_home paths must be specified and valid.",
+                f"tensorrt_home='{tensorrt_home}' valid={tensorrt_home_valid}.",
+            )
+
+        # Set maximum workspace size in byte for
+        # TensorRT (1GB = 1073741824 bytes).
+        os.environ["ORT_TENSORRT_MAX_WORKSPACE_SIZE"] = "1073741824"
+
+        # Set maximum number of iterations to detect unsupported nodes
+        # and partition the models for TensorRT.
+        os.environ["ORT_TENSORRT_MAX_PARTITION_ITERATIONS"] = "1000"
+
+        # Set minimum subgraph node size in graph partitioning
+        # for TensorRT.
+        os.environ["ORT_TENSORRT_MIN_SUBGRAPH_SIZE"] = "1"
+
+        # Set FP16 flag
+        os.environ["ORT_TENSORRT_FP16_ENABLE"] = "0"
+
+    return tensorrt_home
+
+
+def setup_migraphx_vars(args):
+    migraphx_home = None
+
+    if args.use_migraphx:
+        print(f"migraphx_home = {args.migraphx_home}")
+        migraphx_home = args.migraphx_home or os.getenv("MIGRAPHX_HOME") or None
+
+        migraphx_home_not_valid = migraphx_home and not os.path.exists(migraphx_home)
+
+        if migraphx_home_not_valid:
+            raise BuildError(
+                "migraphx_home paths must be specified and valid.",
+                f"migraphx_home='{migraphx_home}' valid={migraphx_home_not_valid}.",
+            )
+    return migraphx_home or ""
+
+
+def setup_cuda_vars(args):
+    cuda_home = ""
+    cudnn_home = ""
+
+    if args.use_cuda:
+        cuda_home = args.cuda_home if args.cuda_home else os.getenv("CUDA_HOME")
+        cudnn_home = args.cudnn_home if args.cudnn_home else os.getenv("CUDNN_HOME")
+
+        cuda_home_valid = cuda_home is not None and os.path.exists(cuda_home)
+        cudnn_home_valid = cudnn_home is not None and os.path.exists(cudnn_home)
+
+        if not cuda_home_valid or (not is_windows() and not cudnn_home_valid):
+            raise BuildError(
+                "cuda_home and cudnn_home paths must be specified and valid.",
+                "cuda_home='{}' valid={}. cudnn_home='{}' valid={}".format(
+                    cuda_home, cuda_home_valid, cudnn_home, cudnn_home_valid
+                ),
+            )
+
+    return cuda_home, cudnn_home

tools/python/util/open_vino_utils.py

@@ -0,0 +1,89 @@
+# Copyright (C) 2022-2023 Intel Corporation
+# Licensed under the MIT License
+
+import os
+import site
+import sys
+
+
+def add_openvino_libs_to_path() -> None:
+    """Adds OpenVINO libraries to the PATH environment variable on Windows."""
+    if sys.platform == "win32":
+        # Installer, pip installs openvino dlls to the different directories
+        # and those paths need to be visible to the openvino-ep modules
+        #
+        # If you're using a custom installation of openvino,
+        # add the location of openvino dlls to your system PATH.
+        openvino_libs = []
+        # looking for the libs in the pip installation path.
+        if os.path.isdir(os.path.join(site.getsitepackages()[1], "openvino", "libs")):
+            openvino_libs.append(os.path.join(site.getsitepackages()[1], "openvino", "libs"))
+        else:
+            # setupvars.bat script set all libs paths to OPENVINO_LIB_PATHS environment variable.
+            openvino_libs_installer = os.getenv("OPENVINO_LIB_PATHS")
+            if openvino_libs_installer:
+                openvino_libs.extend(openvino_libs_installer.split(";"))
+            else:
+                sys.exit(
+                    "Error: Please set the OPENVINO_LIB_PATHS environment variable. "
+                    "If you use an install package, please, run setupvars.bat"
+                )
+        for lib in openvino_libs:
+            lib_path = os.path.join(os.path.dirname(__file__), lib)
+            if os.path.isdir(lib_path):
+                os.environ["PATH"] = os.path.abspath(lib_path) + ";" + os.environ["PATH"]
+                os.add_dll_directory(os.path.abspath(lib_path))
+
+
+def openvino_verify_device_type(device_read):
+    choices = ["CPU_FP32", "CPU_FP16", "GPU_FP32", "GPU_FP16"]
+
+    choices1 = [
+        "CPU_FP32_NO_PARTITION",
+        "CPU_FP16_NO_PARTITION",
+        "GPU_FP32_NO_PARTITION",
+        "GPU_FP16_NO_PARTITION",
+    ]
+    status_hetero = True
+    res = False
+    if device_read in choices:
+        res = True
+    elif device_read in choices1:
+        res = True
+    elif device_read.startswith("HETERO:") or device_read.startswith("MULTI:") or device_read.startswith("AUTO:"):
+        res = True
+        comma_separated_devices = device_read.split(":")
+        comma_separated_devices = comma_separated_devices[1].split(",")
+        if len(comma_separated_devices) < 2:
+            print("At least two devices required in Hetero/Multi/Auto Mode")
+            status_hetero = False
+        dev_options = ["CPU", "GPU"]
+        for dev in comma_separated_devices:
+            if dev not in dev_options:
+                status_hetero = False
+                break
+
+    def invalid_hetero_build():
+        print("\nIf trying to build Hetero/Multi/Auto, specify the supported devices along with it.\n")
+        print("specify the keyword HETERO or MULTI or AUTO followed by the devices ")
+        print("in the order of priority you want to build\n")
+        print("The different hardware devices that can be added in HETERO or MULTI or AUTO")
+        print("are ['CPU','GPU'] \n")
+        print("An example of how to specify the hetero build type. Ex: HETERO:GPU,CPU \n")
+        print("An example of how to specify the MULTI build type. Ex: MULTI:GPU,CPU \n")
+        print("An example of how to specify the AUTO build type. Ex: AUTO:GPU,CPU \n")
+        sys.exit("Wrong Build Type selected")
+
+    if res is False:
+        print("\nYou have selected wrong configuration for the build.")
+        print("pick the build type for specific Hardware Device from following options: ", choices)
+        print("(or) from the following options with graph partitioning disabled: ", choices1)
+        print("\n")
+        if not (device_read.startswith("HETERO") or device_read.startswith("MULTI") or device_read.startswith("AUTO")):
+            invalid_hetero_build()
+        sys.exit("Wrong Build Type selected")
+
+    if status_hetero is False:
+        invalid_hetero_build()
+
+    return device_read

tools/python/util/usage_error.py

@@ -0,0 +1,10 @@
+#  // Copyright (c) Microsoft Corporation. All rights reserved.
+#  // Licensed under the MIT License.
+from .base_error import BaseError
+
+
+class UsageError(BaseError):
+    """Usage related error."""
+
+    def __init__(self, message):
+        super().__init__(message)