[TRTLLMINF-336][infra] Activate BOLT optimization of the SBSA release image wheel

jenkins/BuildDockerImage.groovy

@@ -58,6 +58,13 @@ BOLT_OVERLAY_ENABLED = (params.boltOverlayEnabled ?: env.boltOverlayEnabled ?: "
 // carries profiles). Left false until the enable PR wires it true for the
 // release/nightly path; premerge/new-branch stays lenient (retag plain build).
 BOLT_PROFILES_REQUIRED = (params.boltProfilesRequired ?: env.boltProfilesRequired ?: "false").toString() == "true"
+// Separate from the two above: those govern the profile BUNDLE baked in as a
+// thin layer (Dockerfile.bolt), which documents how to reproduce a BOLTed build
+// but does NOT optimize the binaries the image actually installs. This one
+// governs the INSTALLED wheel -- when true, the wheel unpacked from the build
+// tarball is BOLT-optimized in the image build, before the release stage pip
+// installs it. Kept independent so it can be rolled back on its own.
+BOLT_OPTIMIZE_WHEEL = (params.boltOptimizeWheel ?: env.boltOptimizeWheel ?: "false").toString() == "true"
 // The bundle this pipeline pinned, hoisted out of globalVars in launchBuildJobs
 // because overlayBoltBundle runs well below the scope globalVars is passed into.
 // Empty means unpinned, i.e. take whatever `latest` is.
@@ -281,14 +288,48 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false)
 
 
 def prepareWheelFromBuildStage(dockerfileStage, arch) {
-    if (!ENABLE_USE_WHEEL_FROM_BUILD_STAGE) {
-        echo "useWheelFromBuildStage is false, skip preparing wheel from build stage"
-        return ""
-    }
+    // Whether THIS image has to ship a BOLT-optimized wheel. Answered before the
+    // gates below, and deliberately not subject to them.
+    //
+    // Optimizing the installed wheel is only possible on this path: the wheel is
+    // BOLTed by get_wheel_from_package.py as it is unpacked from the build
+    // tarball, so an image that compiles its own wheel in-container has nothing
+    // to optimize. That makes the two gates below load-bearing for BOLT, and
+    // both are unreliable for reasons that have nothing to do with BOLT:
+    //
+    //   useWheelFromBuildStage -- a kill switch added for nvbug 5433581 in Aug
+    //       2025 and never reverted. It was also read without being declared, so
+    //       it was not merely false but incapable of being true.
+    //   triggerType -- read from env and not a declared parameter of this job,
+    //       so whether it survives the launch is a property of job registration
+    //       rather than of this repo. When it does not, TRIGGER_TYPE is "manual"
+    //       and this returns early no matter what the caller asked for.
+    //
+    // boltOptimizeWheel, by contrast, IS declared, so it reliably arrives. Key
+    // off it and let a BOLT request carry itself past both gates. The bypass is
+    // deliberately narrow -- one arch, one dockerfile stage, only when BOLT was
+    // asked for -- so the nvbug's blast radius stays at the SBSA release image
+    // instead of being reopened for every image this job builds.
+    def boltWheelRequired = BOLT_OPTIMIZE_WHEEL && arch == "sbsa" && dockerfileStage == "release"
+
+    if (!boltWheelRequired) {
+        if (!ENABLE_USE_WHEEL_FROM_BUILD_STAGE) {
+            echo "useWheelFromBuildStage is false, skip preparing wheel from build stage"
+            return ""
+        }
 
-    if (!(TRIGGER_TYPE in ["post-merge", "nightly-release"])) {
-        echo "Trigger type does not use the build stage wheel"
-        return ""
+        if (!(TRIGGER_TYPE in ["post-merge", "nightly-release"])) {
+            echo "Trigger type does not use the build stage wheel"
+            return ""
+        }
+    } else if (!ENABLE_USE_WHEEL_FROM_BUILD_STAGE ||
+               !(TRIGGER_TYPE in ["post-merge", "nightly-release"])) {
+        // Say so rather than doing it quietly: this is the one place the image
+        // build departs from what its parameters literally asked for.
+        echo "[BOLT] boltOptimizeWheel is set for the ${arch} release image, so the " +
+             "build-stage wheel path runs even though useWheelFromBuildStage=" +
+             "${ENABLE_USE_WHEEL_FROM_BUILD_STAGE} and triggerType=${TRIGGER_TYPE} would " +
+             "otherwise skip it. Optimizing the installed wheel is not possible any other way."
     }
 
     if (!dockerfileStage || !arch) {
@@ -302,10 +343,49 @@ def prepareWheelFromBuildStage(dockerfileStage, arch) {
     }
 
     def wheelScript = 'scripts/get_wheel_from_package.py'
-    def wheelArgs = "--arch ${arch} --timeout ${WAIT_TIME_FOR_BUILD_STAGE} --artifact_path " + env.uploadPath
+    // UPLOAD_PATH, not env.uploadPath: they agree whenever the parent passed the
+    // parameter, but an unset env leaves the raw reference interpolating to
+    // "null" and the download then polls .../null/<tarball> until it times out.
+    def wheelArgs = "--arch ${arch} --timeout ${WAIT_TIME_FOR_BUILD_STAGE} --artifact_path ${UPLOAD_PATH}"
+
+    // Only aarch64 has a promoted profile bundle, so only the SBSA image has
+    // anything to apply; x86 installs the wheel as built.
+    if (BOLT_OPTIMIZE_WHEEL && arch == "sbsa") {
+        // The branch whose promoted bundle to apply, resolved the same way the
+        // image overlay resolves it: an explicit override, else this build's own
+        // branch, else main. Deliberately NOT this run's own profiles -- those
+        // are not published until BoltProfileGen finishes, hours after this
+        // build starts, and waiting on them would serialize every release
+        // behind a multi-hour GPU job. get_wheel_from_package.py tries these in
+        // order and fails if none has a bundle.
+        def branches = [params.boltProfileBranch, LLM_BRANCH, "main"]
+            .collect { it?.toString()?.trim() }
+            .findAll { it }
+            .unique()
+        // Pinned, the candidate list collapses to the branch the pin came from:
+        // the ref names one immutable bundle under one promote directory, so
+        // falling through to another branch would optimize the image's wheel
+        // with different profiles than the release wheel and the tested build.
+        if (BOLT_PINNED_REF && BOLT_PINNED_BRANCH) {
+            branches = [BOLT_PINNED_BRANCH]
+            wheelArgs += " --bolt-profile-ref ${BOLT_PINNED_REF}"
+            echo "Release image for ${arch} is pinned to BOLT bundle ${BOLT_PINNED_REF} on ${BOLT_PINNED_BRANCH}"
+        }
+        echo "Release image for ${arch} will BOLT-optimize its wheel using profiles from: ${branches.join(', ')}"
+        wheelArgs += " --bolt-branch ${branches.join(',')}"
+    }
     return " BUILD_WHEEL_SCRIPT=${wheelScript} BUILD_WHEEL_ARGS='${wheelArgs}'"
 }
 
+// Whether a docker build that failed WITH the downloaded-wheel args may be
+// retried without them. The retry rebuilds the wheel from source in-container,
+// which is a fine recovery for an ordinary build but silently defeats the point
+// when that build was also responsible for optimizing the wheel -- the retry
+// would produce an unoptimized release image that looks identical.
+def mayRetryWithoutBuildStageWheel(arch) {
+    return !(BOLT_OPTIMIZE_WHEEL && arch == "sbsa")
+}
+
 // Produce each CANONICAL image from its raw `-noprofiles` build by
 // overlaying the merged LLVM BOLT profile bundle as a thin layer (docker/
 // Dockerfile.bolt via the docker/Makefile `bolt_overlay` target). Canonical is
@@ -614,6 +694,11 @@ def buildImage(config, imageKeyToTag, versionOverride)
                 if (buildWheelArgs.trim().isEmpty()) {
                     throw ex
                 }
+                if (!mayRetryWithoutBuildStageWheel(arch)) {
+                    echo "Build failed with wheel arguments and the BOLT-optimized wheel is required for ${arch}; " +
+                         "NOT retrying from source (that would publish an unoptimized release image)"
+                    throw ex
+                }
                 echo "Build failed with wheel arguments, retrying without them"
                 buildWheelArgs = ""
                 trtllm_utils.llmExecStepWithRetry(this, script: """
@@ -888,6 +973,16 @@ pipeline {
             defaultValue: false,
             description: "When boltOverlayEnabled is true, treat a missing/empty BOLT bundle as a FATAL error instead of retagging the plain build as canonical. Enable for the release/nightly path to guarantee canonical images carry profiles."
         )
+        booleanParam(
+            name: "useWheelFromBuildStage",
+            defaultValue: false,
+            description: "Install the wheel from the build-stage tarball instead of compiling one inside the image. Read since Aug 2025 but never DECLARED, so params.useWheelFromBuildStage was always null and prepareWheelFromBuildStage() returned early on every run -- see nvbug 5433581, whose temporary kill switch was never reverted. Declared here so the flag is at least capable of being set; it stays false by default, and nothing turns it on. boltOptimizeWheel does NOT depend on it: a BOLT request carries itself past this gate for the SBSA release image only."
+        )
+        booleanParam(
+            name: "boltOptimizeWheel",
+            defaultValue: false,
+            description: "BOLT-optimize the wheel the SBSA release image installs, applying the branch's latest promoted profile bundle during the image build, and fail if no bundle can be applied. Independent of boltOverlayEnabled, which only bakes the bundle in as a layer and leaves the installed binaries unoptimized. Uses the last promoted bundle rather than this run's, so the image build never waits on BoltProfileGen. Ignored on x86_64, which has no promoted bundle."
+        )
         string(
             name: "boltProfileBranch",
             defaultValue: "",

jenkins/L0_MergeRequest.groovy

@@ -2617,6 +2617,17 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars)
                             // main, so a ref with no promoted bundle still gets profiles.
                             'boltOverlayEnabled': true,
                             'boltProfilesRequired': true,
+                            // The overlay above only bakes in the profile bundle; it
+                            // leaves the installed wheel unoptimized. This BOLTs the
+                            // wheel the SBSA release image installs, applying the
+                            // branch's last promoted bundle during the image build.
+                            // Deliberately not this run's profiles: those are not
+                            // published until BoltProfileGen finishes, hours after
+                            // this build starts, so depending on them would serialize
+                            // every release behind a multi-hour GPU job. Inert on
+                            // x86_64 (no promoted bundle) and whenever the wheel is
+                            // built from source rather than downloaded.
+                            'boltOptimizeWheel': true,
                         ]
                         if (runMode == "nightly_release") {
                             additionalParameters += [
@@ -2676,9 +2687,12 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars)
                             'uploadPath': UPLOAD_PATH,
                             // Must match Build-Docker-Images above: this path pushes the
                             // same tags, so the scanned+registered image has to be the
-                            // BOLTed canonical one rather than a plain build.
+                            // BOLTed canonical one rather than a plain build, built on
+                            // the BOLT-optimized wheel rather than the first tarball to
+                            // appear.
                             'boltOverlayEnabled': true,
                             'boltProfilesRequired': true,
+                            'boltOptimizeWheel': true,
                         ]
                         if (runMode == "nightly_release") {
                             additionalParameters += [

jenkins/L0_Test.groovy

@@ -5712,7 +5712,8 @@ def checkKitmakerWheelDryRun(pipeline, kitmakerDryRunMetadata)
 // a missing bundle as a skip; it has to, because the same switch covers x86_64,
 // where nothing is promoted yet. This path is aarch64-only and the switch is
 // main-only, so there is always a bundle to find.
-def applyLatestBoltToWheel(pipeline, String wheel, String cpu_arch)
+def applyLatestBoltToWheel(pipeline, String wheel, String cpu_arch, String boltProfileRef = "",
+                           String boltProfileBranch = "")
 {
     def llvmArch = (cpu_arch == AARCH64_TRIPLE) ? "ARM64" : "X64"
     // apply_latest.sh resolves exactly one branch, so try the build's own branch
@@ -5723,6 +5724,13 @@ def applyLatestBoltToWheel(pipeline, String wheel, String cpu_arch)
         .collect { it?.toString()?.trim() }
         .findAll { it }
         .unique()
+    // A pin now supplies its own branch, so nothing is left to guess: the ref
+    // names one object under that branch's promote directory. Walking candidates
+    // would only add ways to fetch something other than what was pinned.
+    if (boltProfileRef && boltProfileBranch) {
+        branches = [boltProfileBranch]
+        echo "[bolt-wheel] pinned to BOLT bundle ${boltProfileRef} on ${boltProfileBranch}"
+    }
 
     stage("BOLT release wheel") {
         sh """
@@ -5748,15 +5756,24 @@ def applyLatestBoltToWheel(pipeline, String wheel, String cpu_arch)
         for (b in branches) {
             rc = sh(returnStatus: true, script: """
                 export PATH="\$PWD/.bolt-llvm/bin:\$PATH"
+                export BOLT_PROFILE_REF='${boltProfileRef}'
+                # Quoted: the branch comes from job env and the wheel name from a
+                # directory listing, so an unquoted expansion would let a shell
+                # metacharacter in either run before apply_latest.sh starts.
                 bash tensorrt_llm/scripts/bolt/internal/apply_latest.sh \
-                     ${b} ${cpu_arch} ${wheel} ${wheel}.bolted
+                     '${b}' '${cpu_arch}' '${wheel}' '${wheel}.bolted'
             """)
             if (rc != 3) {
                 appliedFrom = b
                 break
             }
             echo "[bolt-wheel] no promoted bundle for ${b}/${cpu_arch}; trying next candidate branch"
         }
+        if (rc == 3 && boltProfileRef) {
+            error("[bolt-wheel] pinned BOLT bundle ${boltProfileRef} (${cpu_arch}) not found under any of " +
+                  "${branches.join(', ')}. This pipeline pinned a bundle the release wheel cannot fetch; " +
+                  "refusing to publish a wheel optimized with anything else.")
+        }
         if (rc == 3) {
             error("[bolt-wheel] no promoted BOLT bundle for any of ${branches.join(', ')} (${cpu_arch}); " +
                   "refusing to upload an unoptimized release wheel (promote a bundle via BoltProfileGen, " +
@@ -5766,7 +5783,8 @@ def applyLatestBoltToWheel(pipeline, String wheel, String cpu_arch)
             error("[bolt-wheel] apply_latest.sh failed (rc=${rc}) for ${appliedFrom}/${cpu_arch}")
         }
         sh "mv -f ${wheel}.bolted ${wheel}"
-        echo "[bolt-wheel] ${wheel} is now BOLTed (profiles from ${appliedFrom}/${cpu_arch})"
+        echo "[bolt-wheel] ${wheel} is now BOLTed (profiles from ${appliedFrom}/${cpu_arch}" +
+             (boltProfileRef ? ", bundle ${boltProfileRef})" : ")")
     }
 }
 
@@ -5779,7 +5797,9 @@ def runLLMBuild(
     cpver="cp312",
     plat_name="",
     is_dlfw=false,
-    boltConsume=false)
+    boltConsume=false,
+    boltProfileRef="",
+    boltProfileBranch="")
 {
     sh "pwd && ls -alh"
     sh "env | sort"
@@ -5842,7 +5862,8 @@ def runLLMBuild(
     // inside an already-released image to prove that image can still build from
     // source; optimizing it would prove nothing and only add a failure mode.
     if (boltConsume && cpu_arch == AARCH64_TRIPLE && !wheel_path) {
-        applyLatestBoltToWheel(pipeline, "tensorrt_llm/build/${wheelName}", cpu_arch)
+        applyLatestBoltToWheel(pipeline, "tensorrt_llm/build/${wheelName}", cpu_arch,
+                               boltProfileRef, boltProfileBranch)
     }
 
     def rootWheelUploadPath = "${cpu_arch}/${wheel_path}"
@@ -7018,11 +7039,17 @@ def launchTestJobs(pipeline, testFilter, globalVars)
             // Same switch the build helpers use for the tarball, so the released
             // wheel and the released tarball are never optimized differently.
             def boltConsume = globalVars[BOLT_CONSUME_BUILD]?.toString() == "true"
+            // The bundle this pipeline pinned, so the released wheel is optimized
+            // with the same profiles the tested build was. Empty means unpinned:
+            // apply_latest.sh then takes whatever `latest` is, i.e. today's
+            // behaviour.
+            def boltProfileRef = globalVars[BOLT_PROFILE_REF]?.toString() ?: ""
+            def boltProfileBranch = globalVars[BOLT_PROFILE_BRANCH]?.toString() ?: ""
 
             buildRunner("[${toStageName(values[1], key)}] Build") {
                 wheelPath = runLLMBuild(
                     pipeline, cpu_arch, values[3], "", versionOverride, cpver,
-                    values[7], isDlfw, boltConsume)
+                    values[7], isDlfw, boltConsume, boltProfileRef, boltProfileBranch)
             }
 
             // TODO: Re-enable the sanity check after updating GPU testers' driver version.

scripts/bolt/internal/apply_latest.sh

@@ -56,6 +56,15 @@ fi
 DEST="$(mktemp -d)"
 trap 'rm -rf "$DEST"' EXIT
 
+# 0) Put llvm-bolt on PATH. No-op when the caller already staged it (the Jenkins
+#    build pods do), which keeps this free for them and lets callers that cannot
+#    stage it themselves -- notably the image build, where this runs inside a
+#    docker layer -- just call apply_latest.sh and get a working toolchain.
+if ! . "$HERE/stage_llvm_bolt.sh"; then
+    echo "[apply_latest] FATAL: could not stage llvm-bolt" >&2
+    exit 2
+fi
+
 # 1) Pull the branch `latest` bundle. A missing bundle is fatal here (see header):
 #    consumption was requested but there is nothing promoted to consume.
 if ! bash "$HERE/artifactory.sh" pull-latest "$BRANCH" "$TRIPLE" "$DEST"; then