build: update build.py to pass versions as input parameter and convert version map to dictionary

build.py

@@ -49,7 +49,7 @@
 # information.
 #
 # The TRITON_VERSION file indicates the Triton version and
-# TRITON_VERSION_MAP is used to determine the corresponding container
+# DEFAULT_TRITON_VERSION_MAP is used to determine the corresponding container
 # version and upstream container version (upstream containers are
 # dependencies required by Triton). These versions may be overridden.
 
@@ -69,17 +69,17 @@
 # different versions are used then one backend or the other will
 # incorrectly load the other version of the openvino libraries.
 #
-TRITON_VERSION_MAP = {
-    "2.53.0dev": (
-        "24.12dev",  # triton container
-        "24.10",  # upstream container
-        "1.19.2",  # ORT
-        "2024.4.0",  # ORT OpenVINO
-        "2024.4.0",  # Standalone OpenVINO
-        "3.2.6",  # DCGM version
-        "0.5.5",  # vLLM version
-        "3.12.3",  # RHEL Python version
-    )
+
+DEFAULT_TRITON_VERSION_MAP = {
+    "release_version": "2.53.0dev",
+    "triton_container_version": "24.12dev",
+    "upstream_container_version": "24.10",
+    "ort_version": "1.19.2",
+    "ort_openvino_version": "2024.4.0",
+    "standalone_openvino_version": "2024.4.0",
+    "dcgm_version": "3.2.6",
+    "vllm_version": "0.5.5",
+    "rhel_py_version": "3.12.3",
 }
 
 CORE_BACKENDS = ["ensemble"]
@@ -142,13 +142,9 @@ def target_machine():
 
 def container_versions(version, container_version, upstream_container_version):
     if container_version is None:
-        if version not in TRITON_VERSION_MAP:
-            fail("container version not known for {}".format(version))
-        container_version = TRITON_VERSION_MAP[version][0]
+        container_version = FLAGS.triton_container_version
     if upstream_container_version is None:
-        if version not in TRITON_VERSION_MAP:
-            fail("upstream container version not known for {}".format(version))
-        upstream_container_version = TRITON_VERSION_MAP[version][1]
+        upstream_container_version = FLAGS.upstream_container_version
     return container_version, upstream_container_version
 
 
@@ -665,7 +661,7 @@ def onnxruntime_cmake_args(images, library_paths):
             None,
             os.getenv("TRITON_BUILD_ONNXRUNTIME_VERSION")
             if os.getenv("TRITON_BUILD_ONNXRUNTIME_VERSION")
-            else TRITON_VERSION_MAP[FLAGS.version][2],
+            else FLAGS.ort_version,
         )
     ]
 
@@ -701,15 +697,15 @@ def onnxruntime_cmake_args(images, library_paths):
                     "onnxruntime",
                     "TRITON_BUILD_CONTAINER_VERSION",
                     None,
-                    TRITON_VERSION_MAP[FLAGS.version][1],
+                    FLAGS.triton_container_version,
                 )
             )
 
         # TODO: TPRD-333 OpenVino extension is not currently supported by our manylinux build
         if (
             (target_machine() != "aarch64")
             and (target_platform() != "rhel")
-            and (TRITON_VERSION_MAP[FLAGS.version][3] is not None)
+            and (FLAGS.ort_openvino_version is not None)
         ):
             cargs.append(
                 cmake_backend_enable(
@@ -721,7 +717,7 @@ def onnxruntime_cmake_args(images, library_paths):
                     "onnxruntime",
                     "TRITON_BUILD_ONNXRUNTIME_OPENVINO_VERSION",
                     None,
-                    TRITON_VERSION_MAP[FLAGS.version][3],
+                    FLAGS.ort_openvino_version,
                 )
             )
 
@@ -744,7 +740,7 @@ def openvino_cmake_args():
             "openvino",
             "TRITON_BUILD_OPENVINO_VERSION",
             None,
-            TRITON_VERSION_MAP[FLAGS.version][4],
+            FLAGS.standalone_openvino_version,
         )
     ]
     if target_platform() == "windows":
@@ -767,7 +763,7 @@ def openvino_cmake_args():
                     "openvino",
                     "TRITON_BUILD_CONTAINER_VERSION",
                     None,
-                    TRITON_VERSION_MAP[FLAGS.version][1],
+                    FLAGS.upstream_container_version,
                 )
             )
     return cargs
@@ -822,7 +818,7 @@ def fil_cmake_args(images):
                 "fil",
                 "TRITON_BUILD_CONTAINER_VERSION",
                 None,
-                TRITON_VERSION_MAP[FLAGS.version][1],
+                FLAGS.upstream_container_version,
             )
         )
 
@@ -968,7 +964,7 @@ def create_dockerfile_buildbase_rhel(ddir, dockerfile_name, argmap):
             wget
 """
     # Requires openssl-devel to be installed first for pyenv build to be successful
-    df += change_default_python_version_rhel(TRITON_VERSION_MAP[FLAGS.version][7])
+    df += change_default_python_version_rhel(FLAGS.rhel_py_version)
     df += """
 
 RUN pip3 install --upgrade pip \\
@@ -1230,7 +1226,6 @@ def create_dockerfile_linux(
 LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 COPY --chown=1000:1000 docker/sagemaker/serve /usr/bin/.
 """
-
     # This is required since libcublasLt.so is not present during the build
     # stage of the PyTorch backend
     if not FLAGS.enable_gpu and ("pytorch" in backends):
@@ -1414,9 +1409,7 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
         readline-devel
 """
             # Requires openssl-devel to be installed first for pyenv build to be successful
-            df += change_default_python_version_rhel(
-                TRITON_VERSION_MAP[FLAGS.version][7]
-            )
+            df += change_default_python_version_rhel(FLAGS.rhel_py_version)
             df += """
 RUN pip3 install --upgrade pip \\
     && pip3 install --upgrade \\
@@ -1455,7 +1448,7 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
 # vLLM needed for vLLM backend
 RUN pip3 install vllm=={}
 """.format(
-            TRITON_VERSION_MAP[FLAGS.version][6]
+            FLAGS.vllm_version
         )
 
     if "dali" in backends:
@@ -1487,7 +1480,6 @@ def dockerfile_prepare_container_linux(argmap, backends, enable_gpu, target_mach
 """.format(
         argmap["NVIDIA_BUILD_ID"], argmap["NVIDIA_BUILD_ID"], argmap["NVIDIA_BUILD_REF"]
     )
-
     return df
 
 
@@ -1647,9 +1639,7 @@ def create_build_dockerfiles(
         "TRITON_VERSION": FLAGS.version,
         "TRITON_CONTAINER_VERSION": FLAGS.container_version,
         "BASE_IMAGE": base_image,
-        "DCGM_VERSION": ""
-        if FLAGS.version is None or FLAGS.version not in TRITON_VERSION_MAP
-        else TRITON_VERSION_MAP[FLAGS.version][5],
+        "DCGM_VERSION": FLAGS.dcgm_version,
     }
 
     # For CPU-only image we need to copy some cuda libraries and dependencies
@@ -2027,9 +2017,7 @@ def backend_build(
     # installed within the container will not be picked up by the python backend stub process pybind
     # bindings. It must instead must be installed via pyenv. We package it here for better usability.
     if target_platform() == "rhel" and be == "python":
-        major_minor_version = ".".join(
-            (TRITON_VERSION_MAP[FLAGS.version][7]).split(".")[:2]
-        )
+        major_minor_version = ".".join((FLAGS.rhel_py_version).split(".")[:2])
         version_matched_files = "/usr/lib64/libpython" + major_minor_version + "*"
         cmake_script.cp(
             version_matched_files, os.path.join(repo_install_dir, "backends", be)
@@ -2479,12 +2467,6 @@ def enable_all():
         required=False,
         help="The Triton container version to build. If not specified the container version will be chosen automatically based on --version value.",
     )
-    parser.add_argument(
-        "--upstream-container-version",
-        type=str,
-        required=False,
-        help="The upstream container version to use for the build. If not specified the upstream container version will be chosen automatically based on --version value.",
-    )
     parser.add_argument(
         "--container-prebuild-command",
         type=str,
@@ -2632,7 +2614,60 @@ def enable_all():
         required=False,
         help="Override specified backend CMake argument in the build as <backend>:<name>=<value>. The argument is passed to CMake as -D<name>=<value>. This flag only impacts CMake arguments that are used by build.py. To unconditionally add a CMake argument to the backend build use --extra-backend-cmake-arg.",
     )
-
+    parser.add_argument(
+        "--release-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["release_version"],
+        help="This flag sets the release version for Triton Inference Server to be built. Default: the latest released version.",
+    )
+    parser.add_argument(
+        "--triton-container-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["triton_container_version"],
+        help="This flag sets the container version for Triton Inference Server to be built. Default: the latest released version.",
+    )
+    parser.add_argument(
+        "--upstream-container-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["upstream_container_version"],
+        help="This flag sets the upstream container version for Triton Inference Server to be built. Default: the latest released version.",
+    )
+    parser.add_argument(
+        "--ort-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["ort_version"],
+        help="This flag sets the ORT version for Triton Inference Server to be built. Default: the latest supported version.",
+    )
+    parser.add_argument(
+        "--ort-openvino-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["ort_openvino_version"],
+        help="This flag sets the OpenVino version for Triton Inference Server to be built. Default: the latest supported version.",
+    )
+    parser.add_argument(
+        "--standalone-openvino-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["standalone_openvino_version"],
+        help="This flag sets the standalon OpenVino version for Triton Inference Server to be built. Default: the latest supported version.",
+    )
+    parser.add_argument(
+        "--dcgm-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["dcgm_version"],
+        help="This flag sets the DCGM version for Triton Inference Server to be built. Default: the latest supported version.",
+    )
+    parser.add_argument(
+        "--vllm-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["vllm_version"],
+        help="This flag sets the vLLM version for Triton Inference Server to be built. Default: the latest supported version.",
+    )
+    parser.add_argument(
+        "--rhel-py-version",
+        required=False,
+        default=DEFAULT_TRITON_VERSION_MAP["rhel_py_version"],
+        help="This flag sets the Python version for RHEL platform of Triton Inference Server to be built. Default: the latest supported version.",
+    )
     FLAGS = parser.parse_args()
 
     if FLAGS.image is None:
@@ -2689,8 +2724,7 @@ def enable_all():
     # Determine the versions. Start with Triton version, if --version
     # is not explicitly specified read from TRITON_VERSION file.
     if FLAGS.version is None:
-        with open(os.path.join(THIS_SCRIPT_DIR, "TRITON_VERSION"), "r") as vfile:
-            FLAGS.version = vfile.readline().strip()
+        FLAGS.version = DEFAULT_TRITON_VERSION_MAP["release_version"]
 
     if FLAGS.build_parallel is None:
         FLAGS.build_parallel = multiprocessing.cpu_count() * 2
@@ -2708,15 +2742,9 @@ def enable_all():
     # explicitly. For release branches we use the release branch as
     # the default, otherwise we use 'main'.
     default_repo_tag = "main"
-    cver = FLAGS.container_version
+    cver = FLAGS.upstream_container_version
     if cver is None:
-        if FLAGS.version not in TRITON_VERSION_MAP:
-            fail(
-                "unable to determine default repo-tag, container version not known for {}".format(
-                    FLAGS.version
-                )
-            )
-        cver = TRITON_VERSION_MAP[FLAGS.version][0]
+        cver = FLAGS.triton_container_version
     if not cver.endswith("dev"):
         default_repo_tag = "r" + cver
     log("default repo-tag: {}".format(default_repo_tag))