[TRTLLM-8236][infra] fix platform tag for public wheel

jenkins/L0_Test.groovy

@@ -99,6 +99,12 @@ TESTER_MEMORY = "96Gi"
 CCACHE_DIR="/mnt/sw-tensorrt-pvc/scratch.trt_ccache/llm_ccache"
 MODEL_CACHE_DIR="/scratch.trt_llm_data/llm-models"
 
+KITMAKER_CREDENTIALS_ID = env.kitmakerCredentialsId ? env.kitmakerCredentialsId : "svc_tensorrt_kitmaker_api_token"
+KITMAKER_DRY_RUN_PIC_EMAIL = env.kitmakerDryRunPicEmail ? env.kitmakerDryRunPicEmail : "kefu@nvidia.com"
+KITMAKER_PUBLISH_TO = env.kitmakerPublishTo ? env.kitmakerPublishTo : "both_devzone_pypi"
+RELEASE_SCRIPT_REPO = env.releaseScriptRepo ? env.releaseScriptRepo.trim() : ""
+RELEASE_SCRIPT_COMMIT = env.releaseScriptCommit ? env.releaseScriptCommit.trim() : ""
+
 // GPU types that require open driver
 REQUIRED_OPEN_DRIVER_TYPES = ["b100-ts2", "rtx-5080", "rtx-5090", "rtx-pro-6000", "rtx-pro-6000d"]
 
@@ -3451,7 +3457,94 @@ def checkPipInstall(pipeline, wheel_path, version_local)
 }
 
 
-def runLLMBuild(pipeline, cpu_arch, reinstall_dependencies=false, wheel_path="", version_local="", cpver="cp312")
+def pythonVersionFromCpver(cpver)
+{
+    if (cpver == "cp310") {
+        return "3.10"
+    }
+    if (cpver == "cp312") {
+        return "3.12"
+    }
+    error "Unsupported Python ABI for Kitmaker dry run: ${cpver}"
+}
+
+
+def runKitmakerWheelDryRun(pipeline, wheel_path, python_bin, publish_to)
+{
+    def wheelUrl = "https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/${wheel_path}"
+    def releaseScriptsDir = "release-scripts"
+
+    echo "Running Kitmaker wheel dry run for ${wheelUrl} with publish target ${publish_to}"
+    def kitmakerDryRunMetadata = null
+    stage("Kitmaker Publish Dry Run") {
+        catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') {
+            sh "rm -rf ${releaseScriptsDir}"
+            trtllm_utils.checkoutSource(RELEASE_SCRIPT_REPO, RELEASE_SCRIPT_COMMIT, releaseScriptsDir, false, true)
+            trtllm_utils.llmExecStepWithRetry(
+                pipeline,
+                script: "${python_bin} -m pip install -r ${releaseScriptsDir}/requirements.txt")
+            withCredentials([string(credentialsId: KITMAKER_CREDENTIALS_ID, variable: 'KITMAKER_API_TOKEN')]) {
+                retry(3) {
+                    echo "Publishing Kitmaker wheel dry run for ${wheelUrl}"
+                    def resultData = sh(script: """${python_bin} ${releaseScriptsDir}/kitmaker_wheel.py publish \
+                    --pic-email ${KITMAKER_DRY_RUN_PIC_EMAIL} \
+                    --publish-to ${publish_to} \
+                    --size large \
+                    --wheel-urls ${wheelUrl} \
+                    --no-upload
+                """, returnStdout: true).trim()
+                    echo "${resultData}"
+                    def resultJson = readJSON text: resultData
+                    kitmakerDryRunMetadata = [
+                        releaseUuid: resultJson["release_uuid"],
+                        releaseScriptsDir: releaseScriptsDir,
+                        pythonBin: python_bin,
+                    ]
+                }
+            }
+        }
+    }
+    return kitmakerDryRunMetadata
+}
+
+
+def isKitmakerWheelDryRunEnabled()
+{
+    return RELEASE_SCRIPT_REPO && RELEASE_SCRIPT_COMMIT
+}
+
+
+def checkKitmakerWheelDryRun(pipeline, kitmakerDryRunMetadata)
+{
+    if (!kitmakerDryRunMetadata?.releaseUuid) {
+        echo "Skipping Kitmaker wheel dry run check because publish did not return a release UUID"
+        return
+    }
+    stage("Kitmaker Check Dry Run") {
+        catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') {
+            retry(3) {
+                echo "Checking Kitmaker wheel dry run ${kitmakerDryRunMetadata.releaseUuid}"
+                withCredentials([string(credentialsId: KITMAKER_CREDENTIALS_ID, variable: 'KITMAKER_API_TOKEN')]) {
+                    sh """${kitmakerDryRunMetadata.pythonBin} ${kitmakerDryRunMetadata.releaseScriptsDir}/kitmaker_wheel.py check \
+                        ${kitmakerDryRunMetadata.releaseUuid} \
+                        --wait \
+                        --ignore-missing-logs-error
+                    """
+                }
+            }
+        }
+    }
+}
+
+
+def runLLMBuild(
+    pipeline,
+    cpu_arch,
+    reinstall_dependencies=false,
+    wheel_path="",
+    version_local="",
+    cpver="cp312",
+    plat_name="")
 {
     sh "pwd && ls -alh"
     sh "env | sort"
@@ -3481,6 +3574,7 @@ def runLLMBuild(pipeline, cpu_arch, reinstall_dependencies=false, wheel_path="",
     if (cpu_arch == AARCH64_TRIPLE) {
         buildArgs += " -a '90-real;100-real;103-real;120-real'"
     }
+    def platNameArg = plat_name ? " --plat-name ${plat_name}" : ""
 
     if (version_local) {
         sh """
@@ -3495,7 +3589,7 @@ def runLLMBuild(pipeline, cpu_arch, reinstall_dependencies=false, wheel_path="",
     }
 
     withCredentials([usernamePassword(credentialsId: "urm-artifactory-creds", usernameVariable: 'CONAN_LOGIN_USERNAME', passwordVariable: 'CONAN_PASSWORD')]) {
-        trtllm_utils.llmExecStepWithRetry(pipeline, script: "#!/bin/bash \n" + "cd tensorrt_llm/ && python3 scripts/build_wheel.py --use_ccache -G Ninja -j ${BUILD_JOBS} -D 'WARNING_IS_ERROR=ON' ${buildArgs}")
+        trtllm_utils.llmExecStepWithRetry(pipeline, script: "#!/bin/bash \n" + "cd tensorrt_llm/ && python3 scripts/build_wheel.py --use_ccache -G Ninja -j ${BUILD_JOBS} -D 'WARNING_IS_ERROR=ON' ${buildArgs}${platNameArg}")
     }
     if (env.alternativeTRT) {
         sh "bash -c 'pip3 show tensorrt || true'"
@@ -3504,6 +3598,16 @@ def runLLMBuild(pipeline, cpu_arch, reinstall_dependencies=false, wheel_path="",
     def wheelName = sh(returnStdout: true, script: 'cd tensorrt_llm/build && ls -1 *.whl').trim()
     echo "uploading ${wheelName} to ${cpu_arch}/${wheel_path}"
     trtllm_utils.uploadArtifacts("tensorrt_llm/build/${wheelName}",  "${UPLOAD_PATH}/${cpu_arch}/${wheel_path}")
+    def uploadedWheelPath = "${cpu_arch}/${wheel_path}${wheelName}"
+    def kitmakerDryRunMetadata = null
+    if (version_local) {
+        echo "Skipping Kitmaker wheel dry run for local version '+${version_local}'"
+    } else if (!isKitmakerWheelDryRunEnabled()) {
+        echo "Skipping Kitmaker wheel dry run because releaseScriptRepo or releaseScriptCommit is not set"
+    } else {
+        def kitmakerPython = "tensorrt_llm/.venv-${pythonVersionFromCpver(cpver)}/bin/python3"
+        kitmakerDryRunMetadata = runKitmakerWheelDryRun(pipeline, uploadedWheelPath, kitmakerPython, KITMAKER_PUBLISH_TO)
+    }
 
     if (reinstall_dependencies) {
         // Test installation in the new environment
@@ -3532,6 +3636,7 @@ def runLLMBuild(pipeline, cpu_arch, reinstall_dependencies=false, wheel_path="",
             trtllm_utils.uploadArtifacts("${attrDir}/${f}", "${UPLOAD_PATH}/${cpu_arch}/attribution/${wheel_path}${wheelBase}/")
         }
     }
+    checkKitmakerWheelDryRun(pipeline, kitmakerDryRunMetadata)
 
     return wheelName
 }
@@ -4276,7 +4381,7 @@ def launchTestJobs(pipeline, testFilter)
     }]]}
 
     // Python version and OS for sanity check
-    // Slots: [buildImage, gpuType, cpuArch, reinstallDependencies, isDlfw, pipInstallImage, extraPytorchInstall]
+    // Slots: [buildImage, gpuType, cpuArch, reinstallDependencies, isDlfw, pipInstallImage, extraPytorchInstall, platName]
     x86SanityCheckConfigs = [
         "PY312-DLFW": [
             LLM_DOCKER_IMAGE,
@@ -4286,6 +4391,7 @@ def launchTestJobs(pipeline, testFilter)
             true,
             DLFW_IMAGE,
             false,
+            'manylinux_2_39_x86_64',
         ],
         "PY310-UB2204": [
             LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE,
@@ -4295,6 +4401,7 @@ def launchTestJobs(pipeline, testFilter)
             false,
             UBUNTU_22_04_IMAGE,
             true, // Extra install PyTorch CUDA 13.0 package to align with the CUDA version used for building TensorRT LLM wheels.
+            'manylinux_2_28_x86_64',
         ],
         "PY312-UB2404": [
             LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE,
@@ -4304,6 +4411,7 @@ def launchTestJobs(pipeline, testFilter)
             false,
             UBUNTU_24_04_IMAGE,
             true, // Extra PyTorch CUDA 13.0 install
+            'manylinux_2_28_x86_64',
         ],
     ]
 
@@ -4316,6 +4424,7 @@ def launchTestJobs(pipeline, testFilter)
             false,
             UBUNTU_24_04_IMAGE,
             true, // Extra PyTorch CUDA 13.0 install
+            'manylinux_2_39_aarch64',
         ],
         "PY312-DLFW": [
             LLM_DOCKER_IMAGE,
@@ -4325,6 +4434,7 @@ def launchTestJobs(pipeline, testFilter)
             true,
             DLFW_IMAGE,
             false,
+            'manylinux_2_39_aarch64',
         ],
     ]
 
@@ -4380,7 +4490,7 @@ def launchTestJobs(pipeline, testFilter)
             }
 
             buildRunner("[${toStageName(values[1], key)}] Build") {
-                wheelName = runLLMBuild(pipeline, cpu_arch, values[3], "", versionLocal, cpver)
+                wheelName = runLLMBuild(pipeline, cpu_arch, values[3], "", versionLocal, cpver, values[7])
             }
 
             // TODO: Re-enable the sanity check after updating GPU testers' driver version.

scripts/build_wheel.py

@@ -22,14 +22,14 @@
 import sysconfig
 import tempfile
 import warnings
-from argparse import ArgumentParser
+from argparse import ArgumentParser, ArgumentTypeError
 from contextlib import contextmanager
 from functools import partial
 from multiprocessing import cpu_count
 from pathlib import Path
 from shutil import copy, copytree, rmtree
 from subprocess import DEVNULL, CalledProcessError, check_output, run
-from typing import Sequence
+from typing import Optional, Sequence
 
 try:
     from packaging.requirements import Requirement
@@ -510,7 +510,8 @@ def main(*,
          no_venv: bool = False,
          nvrtc_dynamic_linking: bool = False,
          mypyc: bool = False,
-         require_dynamic_attributions: bool = False):
+         require_dynamic_attributions: bool = False,
+         plat_name: Optional[str] = None):
 
     if clean:
         clean_wheel = True
@@ -1101,6 +1102,11 @@ def get_binding_lib(subdirectory, name):
             clear_folder(dist_dir)
 
         extra_wheel_build_args = os.getenv("EXTRA_WHEEL_BUILD_ARGS", "")
+        plat_name_arg = ""
+        if plat_name:
+            plat_name_arg = f'--config-setting="--build-option=--plat-name={plat_name}"'
+            extra_wheel_build_args = " ".join(
+                arg for arg in (extra_wheel_build_args, plat_name_arg) if arg)
 
         # Attempt to generate attributions using the dependency database
         # Skip if output already exists and the build system hasn't changed
@@ -1143,7 +1149,7 @@ def get_binding_lib(subdirectory, name):
             env["TRTLLM_ENABLE_MYPYC"] = "0"
 
         build_run(
-            f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check --no-isolation --wheel --outdir "{dist_dir}"',
+            f'\"{venv_python}\" -m build {project_dir} --skip-dependency-check {plat_name_arg} --no-isolation --wheel --outdir "{dist_dir}"',
             env=env)
 
     if install:
@@ -1305,6 +1311,21 @@ def add_arguments(parser: ArgumentParser):
                         action="store_true",
                         help="Fail the build if attribution generation fails")
 
+    def _plat_name_type(value):
+        import re
+        if not re.fullmatch(r'[a-zA-Z0-9_]+', value):
+            raise ArgumentTypeError(
+                f"Invalid plat name '{value}': only alphanumerics and underscores are allowed"
+            )
+        return value
+
+    parser.add_argument(
+        "--plat-name",
+        type=_plat_name_type,
+        help=
+        "Wheel platform tag passed to bdist_wheel --plat-name (e.g. linux_x86_64, manylinux_2_28_x86_64)"
+    )
+
 
 if __name__ == "__main__":
     parser = ArgumentParser()