infra: [TRTLLM-5873] Use build stage wheels to speed up docker release image build

[ZhanruiSunCh]

Summary by CodeRabbit

• New Features

  • Added a script to automate downloading and extracting wheel files and benchmark binaries from remote packages.
  • Included a new benchmark executable in the packaged artifacts.

• Chores

  • Made the wheel build script path configurable in the Docker build process.
  • Introduced new variables and improved argument handling in Docker and build scripts.
  • Enhanced build automation with updated environment variable usage and improved error handling.

[ZhanruiSunCh]

/bot run --extra-stage "Build-Docker-Images" --skip-test

[tensorrt-cicd]

PR_Github #7657 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7657 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7663 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7663 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7670 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7670 [ run ] completed with state FAILURE
/LLM/PipelineMonitor/L0_MergeRequest_PR pipeline #78 (Partly Tested) completed with status: 'FAILURE'

[tensorrt-cicd]

PR_Github #7672 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7674 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7672 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7681 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7674 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7695 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7681 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7695 [ run ] completed with state FAILURE
/LLM/PipelineMonitor/L0_MergeRequest_PR pipeline #82 (Partly Tested) completed with status: 'FAILURE'

[tensorrt-cicd]

PR_Github #7703 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7703 [ run ] completed with state FAILURE
/LLM/PipelineMonitor/L0_MergeRequest_PR pipeline #83 (Partly Tested) completed with status: 'FAILURE'

[tensorrt-cicd]

PR_Github #7710 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7714 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7710 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7716 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7714 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7717 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7716 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7722 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7717 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7722 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #7724 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7729 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #7724 [ run ] completed with state ABORTED

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (1)

scripts/get_wheel_from_package.py (1)

30-44: Consider adding architecture validation

The argument parsing should validate the architecture parameter to ensure it's supported.

 def add_arguments(parser: ArgumentParser):
     parser.add_argument("--arch",
                         "-a",
                         required=True,
+                        choices=["x86_64", "aarch64", "arm64"],
                         help="Architecture to build for")

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e42f5a9 and 0c9ebd6.

📒 Files selected for processing (5)

• docker/Dockerfile.multi (1 hunks)
• docker/Makefile (2 hunks)
• jenkins/Build.groovy (1 hunks)
• jenkins/BuildDockerImage.groovy (5 hunks)
• scripts/get_wheel_from_package.py (1 hunks)

🔇 Additional comments (10)

jenkins/Build.groovy (1)

447-447: LGTM: Consistent benchmark packaging

The addition of disaggServerBenchmark follows the established pattern for copying benchmark executables and maintains consistency with the other benchmark files being packaged.

docker/Makefile (2)

42-42: LGTM: Clean variable declaration

The BUILD_WHEEL_SCRIPT variable with empty default value provides flexibility for configuring the wheel build script path.

---

84-84: LGTM: Proper conditional build argument

The conditional passing of BUILD_WHEEL_SCRIPT as a build argument follows the established pattern and only adds the argument when the variable is set, avoiding unnecessary build arguments.

docker/Dockerfile.multi (2)

131-131: LGTM: Configurable build script with sensible default

The BUILD_WHEEL_SCRIPT argument with default value "scripts/build_wheel.py" maintains backward compatibility while enabling customization for the wheel build optimization.

---

133-133: LGTM: Proper variable substitution

The use of ${BUILD_WHEEL_SCRIPT} variable instead of hardcoded path enables the configurable wheel build script functionality.

jenkins/BuildDockerImage.groovy (5)

28-30: LGTM: Clear variable declarations

The TRIGGER_TYPE variable with sensible default and the WAIT_TIME_FOR_BUILD_STAGE constant are well-defined and improve code readability.

---

190-209: LGTM: Well-structured wheel preparation function

The prepareWheelFromBuildStage function has good validation logic and clear conditional returns. The parameter validation and early returns for non-applicable scenarios are appropriate.

---

232-232: LGTM: Consistent variable usage

The change from params.triggerType to TRIGGER_TYPE ensures consistency with the centralized variable definition.

---

297-297: LGTM: Clean integration of wheel preparation

The integration of prepareWheelFromBuildStage into the build arguments is clean and maintains the existing argument building pattern.

---

442-442: LGTM: Improved error handling structure

Moving the echo statement inside the catchError block improves the error handling flow and ensures proper logging during failures.

[ZhanruiSunCh]

/bot run

[tensorrt-cicd]

PR_Github #13112 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #13112 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #9809 completed with status: 'FAILURE'

[ZhanruiSunCh]

/bot run

[tensorrt-cicd]

PR_Github #13311 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #13311 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #9945 completed with status: 'SUCCESS'