diff --git a/CMakeLists.txt b/CMakeLists.txt index bf374cd..082aa0b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -95,6 +95,7 @@ set(TRITON_BUILD_CONTAINER "" CACHE STRING "Triton container to use a base for b set(TRITON_BUILD_CONTAINER_VERSION "" CACHE STRING "Triton container version to target") set(TRITON_BUILD_ONNXRUNTIME_VERSION "" CACHE STRING "ONNXRuntime version to build") set(TRITON_BUILD_ONNXRUNTIME_OPENVINO_VERSION "" CACHE STRING "ONNXRuntime OpenVINO version to build") +set(TRITON_BUILD_ONNXRUNTIME_PARALLEL "" CACHE STRING "Limit ONNX Runtime internal build parallelism") set(TRITON_BUILD_TARGET_PLATFORM "" CACHE STRING "Target platform for ONNXRuntime build") set(TRITON_BUILD_CUDA_VERSION "" CACHE STRING "Version of CUDA install") set(TRITON_BUILD_CUDA_HOME "" CACHE PATH "Path to CUDA install") @@ -324,6 +325,9 @@ endif() # TRITON_ENABLE_ONNXRUNTIME_OPENVINO # if(TRITON_ONNXRUNTIME_DOCKER_BUILD) set(_GEN_FLAGS "") + if(NOT "${TRITON_BUILD_ONNXRUNTIME_PARALLEL}" STREQUAL "") + set(_GEN_FLAGS ${_GEN_FLAGS} "--ort-build-parallel=${TRITON_BUILD_ONNXRUNTIME_PARALLEL}") + endif() # TRITON_BUILD_ONNXRUNTIME_PARALLEL if(NOT ${TRITON_BUILD_TARGET_PLATFORM} STREQUAL "") set(_GEN_FLAGS ${_GEN_FLAGS} "--target-platform=${TRITON_BUILD_TARGET_PLATFORM}") endif() # TRITON_BUILD_TARGET_PLATFORM diff --git a/tools/gen_ort_dockerfile.py b/tools/gen_ort_dockerfile.py index 17ec6a3..f82664f 100755 --- a/tools/gen_ort_dockerfile.py +++ b/tools/gen_ort_dockerfile.py @@ -382,12 +382,16 @@ def dockerfile_for_linux(output_file): else: cuda_archs = "75-real;80-real;86-real;90-real;100f;110f;120f" + ort_parallel_arg = "--parallel" + if FLAGS.ort_build_parallel is not None: + ort_parallel_arg += f" {FLAGS.ort_build_parallel}" + df += """ WORKDIR /workspace/onnxruntime -ARG COMMON_BUILD_ARGS="--config ${{ONNXRUNTIME_BUILD_CONFIG}} --parallel --skip_submodule_sync --build_shared_lib \ +ARG COMMON_BUILD_ARGS="--config ${{ONNXRUNTIME_BUILD_CONFIG}} {} --skip_submodule_sync --build_shared_lib \ --compile_no_warning_as_error --build_dir /workspace/build --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES='{}' --cmake_extra_defines CMAKE_POLICY_VERSION_MINIMUM=3.5 --build_wheel" """.format( - cuda_archs + ort_parallel_arg, cuda_archs ) df += """ @@ -541,6 +545,13 @@ def preprocess_gpu_flags(): choices=["Debug", "Release", "RelWithDebInfo"], help="ORT build configuration.", ) + parser.add_argument( + "--ort-build-parallel", + type=int, + required=False, + default=None, + help="Limit ONNX Runtime internal build parallelism.", + ) parser.add_argument( "--target-platform", required=False,