Clean up build.py

.gitignore

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

tools/ci_build/build.py

@@ -14,115 +14,40 @@
 import sys
 from pathlib import Path
 
-
-def version_to_tuple(version: str) -> tuple:
-    v = []
-    for s in version.split("."):
-        with contextlib.suppress(ValueError):
-            v.append(int(s))
-    return tuple(v)
-
-
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 REPO_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", ".."))
 
 sys.path.insert(0, os.path.join(REPO_DIR, "tools", "python"))
-
-
+# The Following utility packages are from onnxruntime/tools/python/util
 import util.android as android  # noqa: E402
-from util import get_logger, is_linux, is_macOS, is_windows, run  # noqa: E402
+from util import (
+    get_logger,
+    is_linux,
+    is_macOS,
+    is_windows,
+    run,
+    version_to_tuple,
+    check_python_version,
+    str_to_bool,
+    openvino_verify_device_type,
+    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,
+    BuildError,
+    BaseError,
+    setup_cuda_vars,
+    setup_tensorrt_vars,
+    setup_migraphx_vars,
+    setup_cann_vars,
+)  # noqa: E402
 
 log = get_logger("build")
-
-
-class BaseError(Exception):
-    """Base class for errors originating from build.py."""
-
-
-class BuildError(BaseError):
-    """Error from running build steps."""
-
-    def __init__(self, *messages):
-        super().__init__("\n".join(messages))
-
-
-class UsageError(BaseError):
-    """Usage related error."""
-
-    def __init__(self, message):
-        super().__init__(message)
-
-
-def _check_python_version():
-    required_minor_version = 8
-    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()]
-
-
-_check_python_version()
-
-
-def _openvino_verify_device_type(device_read):
-    choices = ["CPU", "GPU", "NPU"]
-
-    choices1 = [
-        "CPU_NO_PARTITION",
-        "GPU_NO_PARTITION",
-        "NPU_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:", "MULTI:", "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", "NPU"]
-        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','NPU'] \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", "MULTI", "AUTO"))):
-            invalid_hetero_build()
-        sys.exit("Wrong Build Type selected")
-
-    if status_hetero is False:
-        invalid_hetero_build()
-
-    return device_read
+check_python_version()
 
 
 def parse_arguments():
@@ -201,7 +126,7 @@
     parser.add_argument("--mpi_home", help="Path to MPI installation dir")
     parser.add_argument("--nccl_home", help="Path to NCCL installation dir")
     parser.add_argument(
-        "--use_mpi", nargs="?", default=False, const=True, type=_str_to_bool, help="Disabled by default."
+        "--use_mpi", nargs="?", default=False, const=True, type=str_to_bool, help="Disabled by default."
     )
 
     # enable ONNX tests
@@ -539,7 +464,7 @@
         "--use_openvino",
         nargs="?",
         const="CPU",
-        type=_openvino_verify_device_type,
+        type=openvino_verify_device_type,
         help="Build with OpenVINO for specific hardware.",
     )
     parser.add_argument(
@@ -789,35 +714,6 @@
     return args
 
 
-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_linux_distro():
-    try:
-        with open("/etc/os-release") as f:
-            dist_info = dict(line.strip().split("=", 1) for line in f.readlines())
-        return dist_info.get("NAME", "").strip('"'), dist_info.get("VERSION", "").strip('"')
-    except (OSError, ValueError):
-        return "", ""
-
-
-def get_config_build_dir(build_dir, config):
-    # build directory per configuration
-    return os.path.join(build_dir, config)
-
-
 def run_subprocess(
     args,
     cwd=None,
@@ -901,69 +797,6 @@
             os.symlink(source_onnx_model_dir, src_model_dir, target_is_directory=True)
 
 
-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 or args.visionos) 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
-    try:
-        import psutil
-
-        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={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 generate_build_tree(
     cmake_path,
     source_dir,
@@ -1731,94 +1564,6 @@
         run_subprocess(cmd_args, env=env)
 
 
-def add_dir_if_exists(directory, dir_list):
-    if os.path.isdir(directory):
-        dir_list.append(directory)
-
-
-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.",
-                f"cuda_home='{cuda_home}' valid={cuda_home_valid}. cudnn_home='{cudnn_home}' valid={cudnn_home_valid}",
-            )
-
-    return cuda_home, cudnn_home
-
-
-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_dml_build(args, cmake_path, build_dir, configs):
     if not args.use_dml:
         return

tools/python/util/__init__.py

@@ -5,6 +5,26 @@
 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 .build_errors import BaseError, BuildError, UsageError
+from .open_vino_utils import openvino_verify_device_type
 
 try:
     import flatbuffers  # noqa: F401