@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hcl 147 lines 4.43 KB
Raw
1 # docker-bake-rocm.hcl - vLLM ROCm Docker build configuration
2 #
3 # This file lives in the vLLM repo at docker/docker-bake-rocm.hcl
4 # Equivalent of docker-bake.hcl for ROCm builds.
5 #
6 # Usage:
7 # docker buildx bake -f docker/docker-bake-rocm.hcl # Build test (default)
8 # docker buildx bake -f docker/docker-bake-rocm.hcl final-rocm # Build final image
9 # docker buildx bake -f docker/docker-bake-rocm.hcl --print # Show resolved config
10 #
11 # CI usage (with the vLLM-owned CI overlay):
12 # docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci
13
14 variable "MAX_JOBS" {
15 # Empty string lets the Dockerfile fall back to $(nproc) via
16 # MAX_JOBS="${MAX_JOBS:-$(nproc)}" in each RUN step, which uses all
17 # available cores on whatever machine the build runs on.
18 # Override with --set '*.args.max_jobs=8' for local builds on small machines.
19 default = ""
20 }
21
22 variable "PYTORCH_ROCM_ARCH" {
23 default = "gfx90a;gfx942;gfx950"
24 }
25
26 variable "COMMIT" {
27 default = ""
28 }
29
30 # Content hash of ci_base-affecting files. Computed by ci-bake-rocm.sh and
31 # embedded as a label so future builds can compare without rebuilding.
32 variable "CI_BASE_CONTENT_HASH" {
33 default = ""
34 }
35
36 variable "CI_BASE_DOCKERFILE" {
37 default = "docker/Dockerfile.rocm"
38 }
39
40 # REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/)
41 # REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source)
42 variable "REMOTE_VLLM" {
43 default = "0"
44 }
45
46 variable "VLLM_BRANCH" {
47 default = "main"
48 }
49
50 # CI_BASE_IMAGE: pre-built ci_base image for per-PR test builds.
51 # Defaults to the local "ci_base" stage for standalone/local builds.
52 # CI overrides this to "rocm/vllm-dev:ci_base" via environment variable.
53 variable "CI_BASE_IMAGE" {
54 default = "rocm/vllm-dev:ci_base"
55 }
56
57 # Upstream dependency commit pins. Plain local bake builds use the Dockerfile
58 # ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
59 # overrides) and writes a small HCL override before invoking CI targets.
60 variable "NIXL_BRANCH" {
61 default = ""
62 }
63
64 variable "UCX_BRANCH" {
65 default = ""
66 }
67
68 variable "ROCSHMEM_BRANCH" {
69 default = ""
70 }
71
72 variable "DEEPEP_BRANCH" {
73 default = ""
74 }
75
76 group "default" {
77 targets = ["test-rocm"]
78 }
79
80 target "_common-rocm" {
81 dockerfile = CI_BASE_DOCKERFILE
82 context = "."
83 args = {
84 max_jobs = MAX_JOBS
85 ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH
86 REMOTE_VLLM = REMOTE_VLLM
87 VLLM_BRANCH = VLLM_BRANCH
88 CI_BASE_IMAGE = CI_BASE_IMAGE
89 }
90 }
91
92 target "_labels" {
93 labels = {
94 "org.opencontainers.image.source" = "https://github.com/vllm-project/vllm"
95 "org.opencontainers.image.vendor" = "vLLM"
96 "org.opencontainers.image.title" = "vLLM ROCm"
97 "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs (ROCm)"
98 "org.opencontainers.image.licenses" = "Apache-2.0"
99 "org.opencontainers.image.revision" = COMMIT
100 }
101 annotations = [
102 "manifest:org.opencontainers.image.revision=${COMMIT}",
103 ]
104 }
105
106 target "test-rocm" {
107 inherits = ["_common-rocm", "_labels"]
108 target = "test"
109 tags = ["rocm/vllm:test"]
110 output = ["type=docker"]
111 }
112
113 # CI base image target - builds only the ci_base stage (NIXL, DeepEP,
114 # torchcodec, requirements, etc.). Used by the weekly scheduled build and
115 # the auto-rebuild trigger when requirements change in a PR.
116 target "ci-base-rocm" {
117 inherits = ["_common-rocm", "_labels"]
118 target = "ci_base"
119 labels = {
120 "vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH
121 }
122 tags = ["rocm/vllm-dev:ci_base"]
123 output = ["type=docker"]
124 }
125
126 # Wheel export target - extracts the built vLLM wheel + test workspace
127 # to local disk. Used by CI to upload the wheel as a Buildkite artifact
128 # so test jobs can assemble images locally from ci_base + wheel instead
129 # of pulling the full large image from Docker Hub.
130 #
131 # Usage:
132 # docker buildx bake -f docker/docker-bake-rocm.hcl export-wheel-rocm
133 # # Creates ./wheel-export/*.whl, ./wheel-export/requirements/, etc.
134 #
135 # After a full bake build, BuildKit cache makes this nearly instant.
136 target "export-wheel-rocm" {
137 inherits = ["_common-rocm"]
138 target = "export_vllm"
139 output = ["type=local,dest=./wheel-export"]
140 }
141
142 target "final-rocm" {
143 inherits = ["_common-rocm", "_labels"]
144 target = "final"
145 tags = ["rocm/vllm:latest"]
146 output = ["type=docker"]
147 }