[TRTLLM-9514][infra] Support Build Docker Images and Test CI in one pipeline

jenkins/BuildDockerImage.groovy

@@ -1,4 +1,4 @@
-@Library(['bloom-jenkins-shared-lib@main', 'trtllm-jenkins-shared-lib@main']) _
+@Library(['bloom-jenkins-shared-lib@main', 'trtllm-jenkins-shared-lib@user/zhanruis/0105_test_mirror']) _
 
 import java.lang.Exception
 import groovy.transform.Field
@@ -31,6 +31,12 @@ TRIGGER_TYPE = env.triggerType ?: "manual"
 
 ENABLE_USE_WHEEL_FROM_BUILD_STAGE = params.useWheelFromBuildStage ?: false
 
+GITHUB_CREDENTIALS_ID = env.GithubCredencialId ?: 'github-cred-trtllm-ci'
+
+// "normal": will build and run all stages
+// "build_for_ci": will only build the images for CI
+MODE = params.mode ?: "normal"
+
 WAIT_TIME_FOR_BUILD_STAGE = 60  // minutes
 
 BUILD_JOBS = "32"
@@ -47,19 +53,22 @@ def CACHED_CHANGED_FILE_LIST = "cached_changed_file_list"
 def ACTION_INFO = "action_info"
 @Field
 def IMAGE_KEY_TO_TAG = "image_key_to_tag"
+@Field
+def GITHUB_SOURCE_REPO_AND_BRANCH = "github_source_repo_and_branch"
 def globalVars = [
     (GITHUB_PR_API_URL): null,
     (CACHED_CHANGED_FILE_LIST): null,
     (ACTION_INFO): null,
     (IMAGE_KEY_TO_TAG): [:],
+    (GITHUB_SOURCE_REPO_AND_BRANCH): null,
 ]
 
 @Field
 def imageKeyToTag = [:]
 
 def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
 {
-    def targetCould = "kubernetes-cpu"
+    def targetCloud = "kubernetes-cpu"
     def containerConfig = ""
     def selectors = """
                 nodeSelector:
@@ -178,7 +187,7 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
     def buildID = env.BUILD_ID
     def nodeLabel = trtllm_utils.appendRandomPostfix("${nodeLabelPrefix}---tensorrt-${jobName}-${buildID}")
     def podConfig = [
-        cloud: targetCould,
+        cloud: targetCloud,
         namespace: "sw-tensorrt",
         label: nodeLabel,
         yaml: """
@@ -356,6 +365,9 @@ def buildImage(config, imageKeyToTag)
                 STAGE=${dockerfileStage} \
                 BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs}
                 """, sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200)
+                if (MODE == "build_for_ci") {
+                    imageKeyToTag[config.stageName] = imageWithTag
+                }
             } catch (InterruptedException ex) {
                 throw ex
             } catch (Exception ex) {
@@ -373,6 +385,9 @@ def buildImage(config, imageKeyToTag)
                 STAGE=${dockerfileStage} \
                 BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs}
                 """, sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200)
+                if (MODE == "build_for_ci") {
+                    imageKeyToTag[config.stageName] = imageWithTag
+                }
             }
             if (target == "ngc-release") {
                 imageKeyToTag["NGC Release Image ${config.arch}"] = imageWithTag
@@ -477,13 +492,36 @@ def launchBuildJobs(pipeline, globalVars, imageKeyToTag) {
             dockerfileStage: "release",
         ],
     ]
+    if (MODE == "build_for_ci") {
+        buildConfigs = buildConfigs.findAll { key, config ->
+            key.contains("Build CI Image")
+        }
+        // Add NGC devel build configs for CI
+        buildConfigs += [
+            "Build NGC devel (x86_64)": [
+                target: "ngc-devel",
+                action: release_action,
+                args: "DOCKER_BUILD_OPTS='--load --platform linux/amd64'",
+                dockerfileStage: "devel",
+            ],
+            "Build NGC devel (SBSA)": [
+                target: "ngc-devel",
+                action: release_action,
+                args: "DOCKER_BUILD_OPTS='--load --platform linux/arm64'",
+                arch: "arm64",
+                dockerfileStage: "devel",
+            ],
+        ]
+        echo "Build configs only for CI"
+    }
     // Override all fields in build config with default values
     buildConfigs.each { key, config ->
         defaultBuildConfig.each { defaultKey, defaultValue ->
             if (!(defaultKey in config)) {
                 config[defaultKey] = defaultValue
             }
         }
+        config.stageName = key
         config.podConfig = createKubernetesPodConfig("build", config.arch, config.build_wheel)
     }
     echo "Build configs:"
@@ -518,6 +556,130 @@ def launchBuildJobs(pipeline, globalVars, imageKeyToTag) {
 }
 
 
+def updateCIImageTag(globalVars) {
+    echo "Update CI Image Tag"
+    // Update jenkins/current_image_tags.properties with newly built image tags and push to PR branch
+
+    def imageTagKeys = [
+        "LLM_DOCKER_IMAGE",
+        "LLM_SBSA_DOCKER_IMAGE",
+        "LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE",
+        "LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE"
+    ]
+
+    def newImageTags = [
+        "LLM_DOCKER_IMAGE" : imageKeyToTag["Build CI Image (x86_64 tritondevel)"],
+        "LLM_SBSA_DOCKER_IMAGE" : imageKeyToTag["Build CI Image (SBSA tritondevel)"],
+        "LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE" : imageKeyToTag["Build CI Image (RockyLinux8 Python310)"],
+        "LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE" : imageKeyToTag["Build CI Image (RockyLinux8 Python312)"],
+    ]
+
+    def emptyKeys = newImageTags.findAll { k, v -> v == null || v.trim().isEmpty() }.keySet()
+    if (!emptyKeys.isEmpty()) {
+        error "Not found image tags for CI: ${emptyKeys.join(', ')}"
+    }
+
+    def filePath = "jenkins/current_image_tags.properties"
+
+    withCredentials([usernamePassword(credentialsId: GITHUB_CREDENTIALS_ID, usernameVariable: 'GITHUB_USERNAME', passwordVariable: 'GITHUB_PASSWORD')]) {
+        // 1. Validate and parse source repo and branch
+        def srcRepoAndBranch = globalVars[GITHUB_SOURCE_REPO_AND_BRANCH]
+        if (!srcRepoAndBranch || !srcRepoAndBranch.contains(":")) {
+            echo "WARNING: No GitHub source repo and branch found. Skipping update."
+            return
+        }
+
+        def parts = srcRepoAndBranch.split(":", 2)
+        if (parts.size() != 2) {
+            error "Invalid GITHUB_SOURCE_REPO_AND_BRANCH format: '${srcRepoAndBranch}'. Expected 'owner/repo:branch'"
+        }
+        def repoPart = parts[0]  // e.g., "ZhanruiSunCh/TensorRT-LLM"
+        def branchName = parts[1]  // e.g., "user/zhanruis/feature_branch"
+        echo "Target fork repo: ${repoPart}, branch: ${branchName}"
+
+        // 2. Setup workspace with upstream repo
+        def workDir = "update_ci_image_tag_workspace"
+
+        // Extract repo path from LLM_REPO (e.g., "https://github.com/NVIDIA/TensorRT-LLM" -> "NVIDIA/TensorRT-LLM")
+        def upstreamRepoPath = 'NVIDIA/TensorRT-LLM'
+        def upstreamRepoUrl = "https://${GITHUB_PASSWORD}@github.com/${upstreamRepoPath}.git"
+        def forkRepoUrl = "https://${GITHUB_PASSWORD}@github.com/${repoPart}.git"
+
+        echo "Setting up workspace with upstream repo: ${upstreamRepoPath}"
+        echo "Fork repo: ${repoPart}"
+
+        // Clean up and prepare workspace
+        sh "rm -rf ${workDir}"
+        sh "mkdir -p ${workDir}"
+
+        // Disable git-lfs globally to avoid lock verification
+        sh "git lfs uninstall || true"
+
+        // Clone upstream repository without LFS
+        sh """
+        export GIT_LFS_SKIP_SMUDGE=1
+        cd ${workDir}
+        git clone --depth 20 ${upstreamRepoUrl} repo
+        """
+
+        // Disable LFS in the cloned repo
+        sh "cd ${workDir}/repo && git lfs uninstall --local || true"
+        sh "cd ${workDir}/repo && git config --local lfs.locksverify false"
+
+        // Add contributor's fork as remote
+        sh "cd ${workDir}/repo && git remote add contributor ${forkRepoUrl}"
+
+        // Fetch PR branch from contributor's fork
+        sh "cd ${workDir}/repo && git fetch contributor ${branchName}"
+
+        // Checkout the PR branch
+        sh "cd ${workDir}/repo && git checkout -b pr-branch contributor/${branchName}"
+
+        // Configure Git user
+        sh "cd ${workDir}/repo && git config user.name 'tensorrt-cicd'"
+        sh "cd ${workDir}/repo && git config user.email '[email protected]'"
+
+        // 3. Read current file and update content
+        echo "Reading and updating ${filePath}"
+        def currentContent = readFile("${workDir}/repo/${filePath}")
+        def lines = currentContent.split("\n") as List
+        def updatedLines = lines.collect { line ->
+            def matchedKey = imageTagKeys.find { key -> line.startsWith(key + "=") }
+            matchedKey ? "${matchedKey}=${newImageTags[matchedKey]}" : line
+        }
+        def updatedContent = updatedLines.join("\n") + "\n"
+
+        // Write updated content
+        writeFile file: "${workDir}/repo/${filePath}", text: updatedContent
+
+        // 4. Commit and push back to contributor's fork
+        echo "Committing and pushing changes"
+
+        // Ensure LFS is still disabled (prevent lock verification)
+        sh "cd ${workDir}/repo && git lfs uninstall --local || true"
+        sh "cd ${workDir}/repo && git config --local lfs.locksverify false"
+
+        // Stage changes
+        sh "cd ${workDir}/repo && git add ${filePath}"
+
+        // Commit with sign-off
+        sh "cd ${workDir}/repo && git commit -s -m '[auto] Update CI image tags with newly built images'"
+
+        // Push to contributor's fork branch (maintainer permission allows this)
+        sh """
+        export GIT_LFS_SKIP_SMUDGE=1
+        cd ${workDir}/repo
+        git push contributor HEAD:${branchName}
+        """
+
+        echo "✅ Successfully updated ${filePath} and pushed to ${repoPart}/${branchName}"
+
+        // Cleanup
+        sh "rm -rf ${workDir}"
+    }
+}
+
+
 def getCommonParameters()
 {
     return [
@@ -583,7 +745,24 @@ pipeline {
                 }
             }
         }
+        stage("Update CI Image Tag") {
+            when {
+                expression {
+                    MODE == "build_for_ci"
+                }
+            }
+            steps {
+                script {
+                    updateCIImageTag(globalVars)
+                }
+            }
+        }
         stage("Upload Artifact") {
+            when {
+                expression {
+                    MODE != "build_for_ci"
+                }
+            }
             steps {
                 script {
                     String imageKeyToTagJson = writeJSON returnText: true, json: imageKeyToTag
@@ -597,7 +776,7 @@ pipeline {
         stage("Wait For Build Job Complete") {
             when {
                 expression {
-                    RUN_SANITY_CHECK
+                    RUN_SANITY_CHECK && MODE != "build_for_ci"
                 }
             }
             steps {
@@ -658,7 +837,7 @@ pipeline {
         stage("Sanity Check For NGC Image") {
             when {
                 expression {
-                    RUN_SANITY_CHECK
+                    RUN_SANITY_CHECK && MODE != "build_for_ci"
                 }
             }
             steps {
@@ -694,7 +873,7 @@ pipeline {
         stage("Register NGC Image For Security Check") {
             when {
                 expression {
-                    return params.nspect_id && params.action == "push"
+                    return params.nspect_id && params.action == "push" && MODE != "build_for_ci"
                 }
             }
             steps {