From b9eafd7e33237a86cfb14f86c8c02ca71777d28e Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Thu, 5 Jun 2025 20:52:47 -0700
Subject: [PATCH 1/8] infra: Support use wheels from Build stage for release
 image build

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 docker/Dockerfile.multi           |  20 +++---
 docker/Makefile                   |   2 +
 jenkins/Build.groovy              |   1 +
 jenkins/BuildDockerImage.groovy   |  20 +++++-
 scripts/get_wheel_from_package.py | 101 ++++++++++++++++++++++++++++++
 5 files changed, 135 insertions(+), 9 deletions(-)
 create mode 100644 scripts/get_wheel_from_package.py

diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi
index 95aa670a090e..042ae035b95f 100644
--- a/docker/Dockerfile.multi
+++ b/docker/Dockerfile.multi
@@ -128,8 +128,9 @@ ENV CCACHE_DIR=/root/.cache/ccache
 # Build the TRT-LLM wheel
 ARG GITHUB_MIRROR=""
 ARG BUILD_WHEEL_ARGS="--clean --benchmarks"
+ARG BUILD_WHEEL_SCRIPT="scripts/build_wheel.py"
 RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=${CCACHE_DIR} \
-    GITHUB_MIRROR=$GITHUB_MIRROR python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS}
+    GITHUB_MIRROR=$GITHUB_MIRROR python3 ${BUILD_WHEEL_SCRIPT} ${BUILD_WHEEL_ARGS}
 
 FROM ${DEVEL_IMAGE} AS release
 
@@ -137,7 +138,11 @@ FROM ${DEVEL_IMAGE} AS release
 RUN mkdir -p /root/.cache/pip
 
 WORKDIR /app/tensorrt_llm
-COPY --from=wheel /src/tensorrt_llm/build/tensorrt_llm*.whl .
+
+ARG SRC_DIR=/src/tensorrt_llm
+
+COPY --from=wheel ${SRC_DIR}/build/tensorrt_llm*.whl .
+
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install tensorrt_llm*.whl && \
     rm tensorrt_llm*.whl && \
@@ -155,14 +160,13 @@ RUN ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensor
     ldconfig && \
     ! ( ldd -v bin/executorWorker | grep tensorrt_llm | grep -q "not found" )
 
-ARG SRC_DIR=/src/tensorrt_llm
-COPY --from=wheel ${SRC_DIR}/benchmarks benchmarks
+COPY benchmarks benchmarks
 ARG CPP_BUILD_DIR=${SRC_DIR}/cpp/build
 COPY --from=wheel \
-     ${CPP_BUILD_DIR}/benchmarks/bertBenchmark \
-     ${CPP_BUILD_DIR}/benchmarks/gptManagerBenchmark \
-     ${CPP_BUILD_DIR}/benchmarks/disaggServerBenchmark \
-     benchmarks/cpp/
+    ${SRC_DIR}/build/benchmarks/bertBenchmark \
+    ${SRC_DIR}/build/benchmarks/gptManagerBenchmark \
+    ${SRC_DIR}/build/benchmarks/disaggServerBenchmark \
+    benchmarks/cpp/
 
 COPY examples examples
 RUN chmod -R a+w examples && \
diff --git a/docker/Makefile b/docker/Makefile
index 2b5022b1ee8e..dde0e461c6fc 100644
--- a/docker/Makefile
+++ b/docker/Makefile
@@ -39,6 +39,7 @@ PLATFORM           ?= $(shell uname -m | grep -q 'aarch64' && echo "arm64" || ec
 CUDA_ARCHS         ?= $(if $(filter arm64,$(PLATFORM)),'90-real;100-real;120-real',)
 BUILD_WHEEL_OPTS   ?=
 BUILD_WHEEL_ARGS   ?= $(shell grep '^ARG BUILD_WHEEL_ARGS=' Dockerfile.multi | grep -o '=.*' | tr -d '="')$(if $(CUDA_ARCHS), --cuda_architectures $(CUDA_ARCHS))$(if $(BUILD_WHEEL_OPTS), $(BUILD_WHEEL_OPTS))
+BUILD_WHEEL_SCRIPT ?=
 TORCH_INSTALL_TYPE ?= skip
 CUDA_VERSION       ?=
 CUDNN_VERSION      ?=
@@ -80,6 +81,7 @@ endef
 		$(if $(BASE_IMAGE), --build-arg BASE_IMAGE=$(BASE_IMAGE)) \
 		$(if $(BASE_TAG), --build-arg BASE_TAG=$(BASE_TAG)) \
 		$(if $(BUILD_WHEEL_ARGS), --build-arg BUILD_WHEEL_ARGS="$(BUILD_WHEEL_ARGS)") \
+		$(if $(BUILD_WHEEL_SCRIPT), --build-arg BUILD_WHEEL_SCRIPT="$(BUILD_WHEEL_SCRIPT)") \
 		$(if $(TORCH_INSTALL_TYPE), --build-arg TORCH_INSTALL_TYPE="$(TORCH_INSTALL_TYPE)") \
 		$(if $(CUDA_VERSION), --build-arg CUDA_VER="$(CUDA_VERSION)") \
 		$(if $(CUDNN_VERSION), --build-arg CUDNN_VER="$(CUDNN_VERSION)") \
diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy
index 77e12ee51003..5dae931b6ac1 100644
--- a/jenkins/Build.groovy
+++ b/jenkins/Build.groovy
@@ -460,6 +460,7 @@ def runLLMBuild(pipeline, buildFlags, tarName, is_linux_x86_64)
     sh "mkdir -p TensorRT-LLM/benchmarks/cpp"
     sh "cp ${LLM_ROOT}/cpp/build/benchmarks/bertBenchmark TensorRT-LLM/benchmarks/cpp"
     sh "cp ${LLM_ROOT}/cpp/build/benchmarks/gptManagerBenchmark TensorRT-LLM/benchmarks/cpp"
+    sh "cp ${LLM_ROOT}/cpp/build/benchmarks/disaggServerBenchmark TensorRT-LLM/benchmarks/cpp"
     sh "cp ${LLM_ROOT}/cpp/build/tensorrt_llm/libtensorrt_llm.so TensorRT-LLM/benchmarks/cpp"
     sh "cp ${LLM_ROOT}/cpp/build/tensorrt_llm/plugins/libnvinfer_plugin_tensorrt_llm.so TensorRT-LLM/benchmarks/cpp"
 
diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy
index 88ab2650374a..5a76c5f72d61 100644
--- a/jenkins/BuildDockerImage.groovy
+++ b/jenkins/BuildDockerImage.groovy
@@ -189,6 +189,23 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
 }
 
 
+def prepareWheelFromBuildStage(makefileStage, arch) {
+    if (!makefileStage || !arch) {
+        echo "Error: makefileStage and arch are required parameters"
+        return ""
+    }
+
+    if (makefileStage != "release") {
+        echo "prepareWheelFromBuildStage: ${makefileStage} is not release"
+        return ""
+    }
+
+    def wheelScript = 'scripts/get_wheel_from_package.py'
+    def wheelArgs = "--arch ${arch} --upload_path " + env.uploadPath
+    // def wheelArgs = "--arch ${arch} --upload_path " + "sw-tensorrt-generic/llm-artifacts/LLM/PipelineMonitor/L0_MergeRequest_PR/89"
+    return " BUILD_WHEEL_SCRIPT=${wheelScript} BUILD_WHEEL_ARGS='${wheelArgs}'"
+}
+
 def buildImage(config, imageKeyToTag)
 {
     def target = config.target
@@ -274,6 +291,7 @@ def buildImage(config, imageKeyToTag)
             }
         }
 
+        args += prepareWheelFromBuildStage(makefileStage, arch)
         // Avoid the frequency of OOM issue when building the wheel
         if (target == "trtllm") {
             if (arch == "x86_64") {
@@ -420,8 +438,8 @@ def launchBuildJobs(pipeline, globalVars, imageKeyToTag) {
                     } catch (InterruptedException e) {
                         throw e
                     } catch (Exception e) {
-                        echo "Build ${key} failed."
                         catchError(buildResult: 'FAILURE', stageResult: 'FAILURE') {
+                            echo "Build ${key} failed."
                             throw e
                         }
                     }
diff --git a/scripts/get_wheel_from_package.py b/scripts/get_wheel_from_package.py
new file mode 100644
index 000000000000..8d6af477849b
--- /dev/null
+++ b/scripts/get_wheel_from_package.py
@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import os
+import shutil
+import subprocess
+import time
+from argparse import ArgumentParser
+from pathlib import Path
+
+
+def get_project_dir():
+    return Path(__file__).parent.resolve().parent
+
+
+def add_arguments(parser: ArgumentParser):
+    parser.add_argument("--arch",
+                        "-a",
+                        required=True,
+                        help="Architecture to build for")
+    parser.add_argument("--upload_path",
+                        "-u",
+                        required=True,
+                        help="Upload path")
+
+
+def get_wheel_from_package(arch, upload_path):
+    if arch == "x86_64":
+        tarfile_name = "TensorRT-LLM.tar.gz"
+    else:
+        tarfile_name = "TensorRT-LLM-GH200.tar.gz"
+
+    tarfile_link = f"https://urm.nvidia.com/artifactory/{upload_path}/{tarfile_name}"
+    while True:
+        try:
+            subprocess.run(["wget", "-nv", tarfile_link], check=True)
+            print(f"Tarfile is available at {tarfile_link}")
+            break
+        except Exception:
+            print("Tarfile not ready yet, waiting 60 seconds...")
+            time.sleep(60)
+
+    llm_root = get_project_dir()
+    tmp_dir = llm_root / "tmp"
+    tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    subprocess.run(["tar", "-zxf", tarfile_name, "-C",
+                    str(tmp_dir)],
+                   check=True)
+
+    tmp_dir = tmp_dir / "TensorRT-LLM"
+
+    build_dir = llm_root / "build"
+    build_dir.mkdir(parents=True, exist_ok=True)
+
+    benchmarks_dir = llm_root / "build" / "benchmarks"
+    benchmarks_dir.mkdir(parents=True, exist_ok=True)
+
+    wheel_files = glob.glob(str(tmp_dir / "tensorrt_llm*.whl"))
+    for wheel_file in wheel_files:
+        shutil.move(wheel_file, str(build_dir))
+        print(f"Moved wheel file: {wheel_file} -> {build_dir}")
+
+    benchmark_files = [
+        "bertBenchmark", "gptManagerBenchmark", "disaggServerBenchmark"
+    ]
+
+    for benchmark in benchmark_files:
+        src_path = tmp_dir / "benchmarks" / "cpp" / benchmark
+        if src_path.exists():
+            dst_path = benchmarks_dir / benchmark
+            shutil.copy2(src_path, dst_path)
+            print(f"Copied benchmark file: {src_path} -> {dst_path}")
+        else:
+            print(f"Warning: Benchmark file not found: {src_path}")
+
+    shutil.rmtree(tmp_dir)
+
+    if os.path.exists(tarfile_name):
+        os.remove(tarfile_name)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    add_arguments(parser)
+    args = parser.parse_args()
+    get_wheel_from_package(**vars(args))

From f828316332ef61ca76f48174ff5a2ca72f54d856 Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Thu, 5 Jun 2025 20:58:53 -0700
Subject: [PATCH 2/8] Fix format

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 docker/Dockerfile.multi           | 15 +++++++--------
 scripts/get_wheel_from_package.py |  2 +-
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi
index 042ae035b95f..6f22d8cea1ca 100644
--- a/docker/Dockerfile.multi
+++ b/docker/Dockerfile.multi
@@ -139,9 +139,7 @@ RUN mkdir -p /root/.cache/pip
 
 WORKDIR /app/tensorrt_llm
 
-ARG SRC_DIR=/src/tensorrt_llm
-
-COPY --from=wheel ${SRC_DIR}/build/tensorrt_llm*.whl .
+COPY --from=wheel /src/tensorrt_llm/build/tensorrt_llm*.whl .
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install tensorrt_llm*.whl && \
@@ -160,13 +158,14 @@ RUN ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensor
     ldconfig && \
     ! ( ldd -v bin/executorWorker | grep tensorrt_llm | grep -q "not found" )
 
-COPY benchmarks benchmarks
+ARG SRC_DIR=/src/tensorrt_llm
+COPY --from=wheel ${SRC_DIR}/benchmarks benchmarks
 ARG CPP_BUILD_DIR=${SRC_DIR}/cpp/build
 COPY --from=wheel \
-    ${SRC_DIR}/build/benchmarks/bertBenchmark \
-    ${SRC_DIR}/build/benchmarks/gptManagerBenchmark \
-    ${SRC_DIR}/build/benchmarks/disaggServerBenchmark \
-    benchmarks/cpp/
+     ${CPP_BUILD_DIR}/benchmarks/bertBenchmark \
+     ${CPP_BUILD_DIR}/benchmarks/gptManagerBenchmark \
+     ${CPP_BUILD_DIR}/benchmarks/disaggServerBenchmark \
+     benchmarks/cpp/
 
 COPY examples examples
 RUN chmod -R a+w examples && \
diff --git a/scripts/get_wheel_from_package.py b/scripts/get_wheel_from_package.py
index 8d6af477849b..33c05be9d51c 100644
--- a/scripts/get_wheel_from_package.py
+++ b/scripts/get_wheel_from_package.py
@@ -67,7 +67,7 @@ def get_wheel_from_package(arch, upload_path):
     build_dir = llm_root / "build"
     build_dir.mkdir(parents=True, exist_ok=True)
 
-    benchmarks_dir = llm_root / "build" / "benchmarks"
+    benchmarks_dir = llm_root / "cpp" / "build" / "benchmarks"
     benchmarks_dir.mkdir(parents=True, exist_ok=True)
 
     wheel_files = glob.glob(str(tmp_dir / "tensorrt_llm*.whl"))

From c6a919c07e8b10dcde0b1b06d452bc4d9ce33e1a Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Fri, 6 Jun 2025 02:44:38 -0700
Subject: [PATCH 3/8] Fix format

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 docker/Dockerfile.multi | 2 --
 1 file changed, 2 deletions(-)

diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi
index 6f22d8cea1ca..0d156c7a764b 100644
--- a/docker/Dockerfile.multi
+++ b/docker/Dockerfile.multi
@@ -138,9 +138,7 @@ FROM ${DEVEL_IMAGE} AS release
 RUN mkdir -p /root/.cache/pip
 
 WORKDIR /app/tensorrt_llm
-
 COPY --from=wheel /src/tensorrt_llm/build/tensorrt_llm*.whl .
-
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install tensorrt_llm*.whl && \
     rm tensorrt_llm*.whl && \

From ff983503ddcd86687e1f19cad4ad668277885fd7 Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Fri, 6 Jun 2025 02:51:49 -0700
Subject: [PATCH 4/8] Fix manual can't run error

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 jenkins/BuildDockerImage.groovy | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy
index 5a76c5f72d61..8f42d9d0714f 100644
--- a/jenkins/BuildDockerImage.groovy
+++ b/jenkins/BuildDockerImage.groovy
@@ -27,6 +27,7 @@ LLM_SHORT_COMMIT = env.gitlabCommit ? env.gitlabCommit.substring(0, 7) : "undefi
 LLM_DEFAULT_TAG = env.defaultTag ?: "${LLM_SHORT_COMMIT}-${LLM_BRANCH_TAG}-${BUILD_NUMBER}"
 
 RUN_SANITY_CHECK = params.runSanityCheck ?: false
+TRIGGER_TYPE = env.triggerType ?: "manual"
 
 BUILD_JOBS = "32"
 BUILD_JOBS_RELEASE_X86_64 = "32"
@@ -190,6 +191,11 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
 
 
 def prepareWheelFromBuildStage(makefileStage, arch) {
+    if (TRIGGER_TYPE == "manual") {
+        echo "Trigger type is manual, skip preparing wheel from build stage"
+        return ""
+    }
+
     if (!makefileStage || !arch) {
         echo "Error: makefileStage and arch are required parameters"
         return ""
@@ -226,11 +232,15 @@ def buildImage(config, imageKeyToTag)
     def dependentImageWithTag = "${IMAGE_NAME}/${dependent.dockerfileStage}:${dependentTag}"
     def customImageWithTag = "${IMAGE_NAME}/${dockerfileStage}:${customTag}"
 
-    if (target == "ngc-release" && params.triggerType == "post-merge") {
-        echo "Use NGC artifacts for post merge build"
-        dependentImageWithTag = "${NGC_IMAGE_NAME}:${dependentTag}"
-        imageWithTag = "${NGC_IMAGE_NAME}:${tag}"
-        customImageWithTag = "${NGC_IMAGE_NAME}:${customTag}"
+    if (target == "ngc-release") {
+        if (TRIGGER_TYPE == "post-merge") {
+            echo "Use NGC artifacts for post merge build"
+            dependentImageWithTag = "${NGC_IMAGE_NAME}:${dependentTag}"
+            imageWithTag = "${NGC_IMAGE_NAME}:${tag}"
+            customImageWithTag = "${NGC_IMAGE_NAME}:${customTag}"
+        }
+        imageKeyToTag["NGC Devel Image ${config.arch}"] = dependentImageWithTag
+        imageKeyToTag["NGC Release Image ${config.arch}"] = imageWithTag
     }
 
     args += " GITHUB_MIRROR=https://urm.nvidia.com/artifactory/github-go-remote"

From 2955c1fae65fe5420b3cd1232b642c60570e4648 Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Mon, 9 Jun 2025 00:41:19 -0700
Subject: [PATCH 5/8] Remove test code

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 jenkins/BuildDockerImage.groovy | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy
index 8f42d9d0714f..3efe93569e15 100644
--- a/jenkins/BuildDockerImage.groovy
+++ b/jenkins/BuildDockerImage.groovy
@@ -208,7 +208,6 @@ def prepareWheelFromBuildStage(makefileStage, arch) {
 
     def wheelScript = 'scripts/get_wheel_from_package.py'
     def wheelArgs = "--arch ${arch} --upload_path " + env.uploadPath
-    // def wheelArgs = "--arch ${arch} --upload_path " + "sw-tensorrt-generic/llm-artifacts/LLM/PipelineMonitor/L0_MergeRequest_PR/89"
     return " BUILD_WHEEL_SCRIPT=${wheelScript} BUILD_WHEEL_ARGS='${wheelArgs}'"
 }
 

From 91c39ca0485769e760e079a79df802aa69c0a11f Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Mon, 23 Jun 2025 09:31:53 -0700
Subject: [PATCH 6/8] Fix rename

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 jenkins/BuildDockerImage.groovy | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy
index 3efe93569e15..fc8a9eb494d7 100644
--- a/jenkins/BuildDockerImage.groovy
+++ b/jenkins/BuildDockerImage.groovy
@@ -190,19 +190,19 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
 }
 
 
-def prepareWheelFromBuildStage(makefileStage, arch) {
+def prepareWheelFromBuildStage(dockerfileStage, arch) {
     if (TRIGGER_TYPE == "manual") {
         echo "Trigger type is manual, skip preparing wheel from build stage"
         return ""
     }
 
-    if (!makefileStage || !arch) {
-        echo "Error: makefileStage and arch are required parameters"
+    if (!dockerfileStage || !arch) {
+        echo "Error: dockerfileStage and arch are required parameters"
         return ""
     }
 
-    if (makefileStage != "release") {
-        echo "prepareWheelFromBuildStage: ${makefileStage} is not release"
+    if (dockerfileStage != "release") {
+        echo "prepareWheelFromBuildStage: ${dockerfileStage} is not release"
         return ""
     }
 
@@ -300,7 +300,7 @@ def buildImage(config, imageKeyToTag)
             }
         }
 
-        args += prepareWheelFromBuildStage(makefileStage, arch)
+        args += prepareWheelFromBuildStage(dockerfileStage, arch)
         // Avoid the frequency of OOM issue when building the wheel
         if (target == "trtllm") {
             if (arch == "x86_64") {

From 56389f2fd21c5cc87c3e5cd1ddcecea2a7b62950 Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Mon, 7 Jul 2025 00:43:53 -0700
Subject: [PATCH 7/8] Add timeout for wait build stage

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 jenkins/BuildDockerImage.groovy   |  4 +++-
 scripts/get_wheel_from_package.py | 17 ++++++++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy
index fc8a9eb494d7..4fed814c6fbc 100644
--- a/jenkins/BuildDockerImage.groovy
+++ b/jenkins/BuildDockerImage.groovy
@@ -29,6 +29,8 @@ LLM_DEFAULT_TAG = env.defaultTag ?: "${LLM_SHORT_COMMIT}-${LLM_BRANCH_TAG}-${BUI
 RUN_SANITY_CHECK = params.runSanityCheck ?: false
 TRIGGER_TYPE = env.triggerType ?: "manual"
 
+WAIT_TIME_FOR_BUILD_STAGE = 60  // minutes
+
 BUILD_JOBS = "32"
 BUILD_JOBS_RELEASE_X86_64 = "32"
 BUILD_JOBS_RELEASE_SBSA = "32"
@@ -207,7 +209,7 @@ def prepareWheelFromBuildStage(dockerfileStage, arch) {
     }
 
     def wheelScript = 'scripts/get_wheel_from_package.py'
-    def wheelArgs = "--arch ${arch} --upload_path " + env.uploadPath
+    def wheelArgs = "--arch ${arch} --timeout ${WAIT_TIME_FOR_BUILD_STAGE} --upload_path " + env.uploadPath
     return " BUILD_WHEEL_SCRIPT=${wheelScript} BUILD_WHEEL_ARGS='${wheelArgs}'"
 }
 
diff --git a/scripts/get_wheel_from_package.py b/scripts/get_wheel_from_package.py
index 33c05be9d51c..fff243d92f36 100644
--- a/scripts/get_wheel_from_package.py
+++ b/scripts/get_wheel_from_package.py
@@ -36,22 +36,33 @@ def add_arguments(parser: ArgumentParser):
                         "-u",
                         required=True,
                         help="Upload path")
+    parser.add_argument("--timeout",
+                        "-t",
+                        type=int,
+                        default=60,
+                        help="Timeout in minutes")
 
 
-def get_wheel_from_package(arch, upload_path):
+def get_wheel_from_package(arch, upload_path, timeout):
     if arch == "x86_64":
         tarfile_name = "TensorRT-LLM.tar.gz"
     else:
         tarfile_name = "TensorRT-LLM-GH200.tar.gz"
 
     tarfile_link = f"https://urm.nvidia.com/artifactory/{upload_path}/{tarfile_name}"
-    while True:
+    for attempt in range(timeout):
         try:
             subprocess.run(["wget", "-nv", tarfile_link], check=True)
             print(f"Tarfile is available at {tarfile_link}")
             break
         except Exception:
-            print("Tarfile not ready yet, waiting 60 seconds...")
+            if attempt == timeout - 1:
+                raise TimeoutError(
+                    f"Failed to download file after {timeout} attempts: {tarfile_link}"
+                )
+            print(
+                f"Tarfile not ready yet, waiting 60 seconds... (attempt {attempt + 1}/{timeout})"
+            )
             time.sleep(60)
 
     llm_root = get_project_dir()

From 1c814414266bc34d0ae545f14a5164e0df28475e Mon Sep 17 00:00:00 2001
From: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
Date: Sun, 27 Jul 2025 07:57:56 -0700
Subject: [PATCH 8/8] Fix review

Signed-off-by: ZhanruiSunCh <184402041+ZhanruiSunCh@users.noreply.github.com>
---
 jenkins/BuildDockerImage.groovy   |  6 +++---
 scripts/get_wheel_from_package.py | 10 +++++-----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy
index 4fed814c6fbc..b09a91352502 100644
--- a/jenkins/BuildDockerImage.groovy
+++ b/jenkins/BuildDockerImage.groovy
@@ -193,8 +193,8 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
 
 
 def prepareWheelFromBuildStage(dockerfileStage, arch) {
-    if (TRIGGER_TYPE == "manual") {
-        echo "Trigger type is manual, skip preparing wheel from build stage"
+    if (TRIGGER_TYPE != "post-merge") {
+        echo "Trigger type is not post-merge, skip preparing wheel from build stage"
         return ""
     }
 
@@ -209,7 +209,7 @@ def prepareWheelFromBuildStage(dockerfileStage, arch) {
     }
 
     def wheelScript = 'scripts/get_wheel_from_package.py'
-    def wheelArgs = "--arch ${arch} --timeout ${WAIT_TIME_FOR_BUILD_STAGE} --upload_path " + env.uploadPath
+    def wheelArgs = "--arch ${arch} --timeout ${WAIT_TIME_FOR_BUILD_STAGE} --artifact_path " + env.uploadPath
     return " BUILD_WHEEL_SCRIPT=${wheelScript} BUILD_WHEEL_ARGS='${wheelArgs}'"
 }
 
diff --git a/scripts/get_wheel_from_package.py b/scripts/get_wheel_from_package.py
index fff243d92f36..cb604482c27b 100644
--- a/scripts/get_wheel_from_package.py
+++ b/scripts/get_wheel_from_package.py
@@ -31,11 +31,11 @@ def add_arguments(parser: ArgumentParser):
     parser.add_argument("--arch",
                         "-a",
                         required=True,
-                        help="Architecture to build for")
-    parser.add_argument("--upload_path",
+                        help="Architecture of the built package")
+    parser.add_argument("--artifact_path",
                         "-u",
                         required=True,
-                        help="Upload path")
+                        help="the path of the built package")
     parser.add_argument("--timeout",
                         "-t",
                         type=int,
@@ -43,13 +43,13 @@ def add_arguments(parser: ArgumentParser):
                         help="Timeout in minutes")
 
 
-def get_wheel_from_package(arch, upload_path, timeout):
+def get_wheel_from_package(arch, artifact_path, timeout):
     if arch == "x86_64":
         tarfile_name = "TensorRT-LLM.tar.gz"
     else:
         tarfile_name = "TensorRT-LLM-GH200.tar.gz"
 
-    tarfile_link = f"https://urm.nvidia.com/artifactory/{upload_path}/{tarfile_name}"
+    tarfile_link = f"https://urm.nvidia.com/artifactory/{artifact_path}/{tarfile_name}"
     for attempt in range(timeout):
         try:
             subprocess.run(["wget", "-nv", tarfile_link], check=True)