[CI] Build CPU image against torch nightly for TORCH_NIGHTLY runs (#48750)
Signed-off-by: Andrey Talman <atalman@users.noreply.github.com> Signed-off-by: Andrey Talman <atalman@fb.com> Co-authored-by: Andrey Talman <atalman@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Andrey Talman committed
Aug 31, 2026 at 20:25 UTC
8600db5dff18054f7a4314f6f8bba4259e3e2a98
4 files changed
+104
-13
.buildkite/image_build/image_build_arm64.sh
+12
-1
@@ -9,7 +9,17 @@ fi
9
REGISTRY=$1
10
REPO=$2
11
BUILDKITE_COMMIT=$3
12
-IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64"
12
+
13
+# When TORCH_NIGHTLY=1, build against torch nightly and tag it with the
14
+# -torch-nightly-arm64 suffix the DGX Spark test steps pull on the nightly lane
15
+# (ci-infra get_image(arm64=True) appends -torch-nightly before -arm64).
16
+PYTORCH_NIGHTLY_ARGS=()
17
+if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then
18
+ IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-torch-nightly-arm64"
19
+ PYTORCH_NIGHTLY_ARGS=(--build-arg PYTORCH_NIGHTLY=1)
20
+else
21
+ IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64"
22
+fi
23
24
# authenticate with AWS ECR
25
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
@@ -30,6 +40,7 @@ else
40
--build-arg torch_cuda_arch_list="9.0 10.0 11.0 12.0" \
41
--build-arg USE_SCCACHE=1 \
42
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
43
+ "${PYTORCH_NIGHTLY_ARGS[@]}" \
44
--tag "$IMAGE" \
45
--target test \
46
--progress plain .
.buildkite/image_build/image_build_cpu.sh
+25
-8
@@ -9,7 +9,16 @@ fi
9
REGISTRY=$1
10
REPO=$2
11
BUILDKITE_COMMIT=$3
12
-IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu"
12
+
13
+# When TORCH_NIGHTLY=1, build the CPU image against torch nightly and tag it
14
+# with the -torch-nightly-cpu suffix the test steps pull on the nightly lane.
15
+PYTORCH_NIGHTLY_ARGS=()
16
+if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then
17
+ IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-torch-nightly-cpu"
18
+ PYTORCH_NIGHTLY_ARGS=(--build-arg PYTORCH_NIGHTLY=1)
19
+else
20
+ IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu"
21
+fi
22
23
# replace invalid characters in Docker image tags and truncate to 128 chars
24
clean_docker_tag() {
@@ -20,32 +29,39 @@ clean_docker_tag() {
29
# resolve and set: CACHE_TO, CACHE_FROM, CACHE_FROM_BASE_BRANCH, CACHE_FROM_MAIN
30
# Reuses the same ECR cache repos as the CUDA image build, with an
31
# "x86_cpu" suffix on every tag so CPU and CUDA cache blobs stay isolated.
32
+# The nightly lane appends "-nightly": it resolves a different torch, so
33
+# sharing a cache ref with the regular lane would have the two evict each
34
+# other from the same ref on every run.
35
prepare_cache_tags() {
36
TEST_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache"
37
MAIN_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache"
38
+ CACHE_SUFFIX="x86_cpu"
39
+ if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then
40
+ CACHE_SUFFIX="x86_cpu-nightly"
41
+ fi
42
43
if [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]]; then
44
if [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
29
- cache="${MAIN_CACHE_ECR}:latest-x86_cpu"
45
+ cache="${MAIN_CACHE_ECR}:latest-${CACHE_SUFFIX}"
46
else
47
clean_branch=$(clean_docker_tag "${BUILDKITE_BRANCH:-unknown}")
32
- cache="${TEST_CACHE_ECR}:${clean_branch}-x86_cpu"
48
+ cache="${TEST_CACHE_ECR}:${clean_branch}-${CACHE_SUFFIX}"
49
fi
50
CACHE_TO="$cache"
51
CACHE_FROM="$cache"
52
CACHE_FROM_BASE_BRANCH="$cache"
53
else
38
- CACHE_TO="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}-x86_cpu"
39
- CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}-x86_cpu"
54
+ CACHE_TO="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}-${CACHE_SUFFIX}"
55
+ CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}-${CACHE_SUFFIX}"
56
if [[ "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" == "main" ]]; then
41
- CACHE_FROM_BASE_BRANCH="${MAIN_CACHE_ECR}:latest-x86_cpu"
57
+ CACHE_FROM_BASE_BRANCH="${MAIN_CACHE_ECR}:latest-${CACHE_SUFFIX}"
58
else
59
clean_base=$(clean_docker_tag "${BUILDKITE_PULL_REQUEST_BASE_BRANCH}")
44
- CACHE_FROM_BASE_BRANCH="${TEST_CACHE_ECR}:${clean_base}-x86_cpu"
60
+ CACHE_FROM_BASE_BRANCH="${TEST_CACHE_ECR}:${clean_base}-${CACHE_SUFFIX}"
61
fi
62
fi
63
48
- CACHE_FROM_MAIN="${MAIN_CACHE_ECR}:latest-x86_cpu"
64
+ CACHE_FROM_MAIN="${MAIN_CACHE_ECR}:latest-${CACHE_SUFFIX}"
65
}
66
67
# authenticate with AWS ECR (public, for the image; private, for the cache repos)
@@ -90,6 +106,7 @@ else
106
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
107
--build-arg VLLM_CPU_X86=true \
108
--build-arg USE_SCCACHE=1 \
109
+ "${PYTORCH_NIGHTLY_ARGS[@]}" \
110
--tag "$IMAGE" \
111
--target vllm-test \
112
"${CACHE_FROM_ARGS[@]}" \
docker/Dockerfile.cpu
+65
-4
@@ -100,13 +100,33 @@ ENV UV_HTTP_TIMEOUT=500
100
ENV UV_INDEX_STRATEGY="unsafe-best-match"
101
ENV UV_LINK_MODE="copy"
102
103
+# PyTorch provides its own indexes for standard and nightly builds.
104
+# PYTORCH_NIGHTLY=1 builds the CPU image against torch nightly (parity with the
105
+# GPU torch-nightly lane); default 0 keeps the pinned release CPU build.
106
+ARG PYTORCH_NIGHTLY=0
107
+ARG PYTORCH_CPU_INDEX_BASE_URL=https://download.pytorch.org/whl
108
+
109
# Copy requirements files for installation
110
COPY requirements/common.txt requirements/common.txt
111
COPY requirements/cpu.txt requirements/cpu.txt
112
+COPY use_existing_torch.py use_existing_torch.py
113
114
+# For nightly: install torch nightly (CPU) first, strip torch pins from the
115
+# requirements via use_existing_torch.py so they can't pull it back to a release
116
+# build, then install the rest from the nightly CPU index (no --torch-backend,
117
+# which would force the stable channel). Release path is unchanged.
118
RUN --mount=type=cache,target=/root/.cache/uv \
119
uv pip install --upgrade pip && \
109
- uv pip install -r requirements/cpu.txt --torch-backend cpu
120
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
121
+ echo "Installing torch nightly (CPU)..." \
122
+ && uv pip install torch torchaudio torchvision torchcodec --pre \
123
+ --index-url ${PYTORCH_CPU_INDEX_BASE_URL}/nightly/cpu \
124
+ && python3 use_existing_torch.py --prefix \
125
+ && uv pip install -r requirements/cpu.txt \
126
+ --extra-index-url ${PYTORCH_CPU_INDEX_BASE_URL}/nightly/cpu; \
127
+ else \
128
+ uv pip install -r requirements/cpu.txt --torch-backend cpu; \
129
+ fi
130
131
ENV TARGETARCH=${TARGETARCH}
132
@@ -191,6 +211,10 @@ RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
211
# build it never uses.
212
FROM base-arch AS vllm-src
213
214
+# Re-declare torch-nightly args (ARGs do not cross FROM)
215
+ARG PYTORCH_NIGHTLY=0
216
+ARG PYTORCH_CPU_INDEX_BASE_URL=https://download.pytorch.org/whl
217
+
218
ARG GIT_REPO_CHECK=0
219
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
220
ARG VLLM_CPU_X86=0
@@ -215,7 +239,13 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \
239
COPY requirements/build/cpu.txt requirements/build/cpu.txt
240
241
RUN --mount=type=cache,target=/root/.cache/uv \
218
- uv pip install -r requirements/build/cpu.txt --torch-backend cpu
242
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
243
+ python3 /workspace/use_existing_torch.py --prefix \
244
+ && uv pip install -r requirements/build/cpu.txt \
245
+ --extra-index-url ${PYTORCH_CPU_INDEX_BASE_URL}/nightly/cpu; \
246
+ else \
247
+ uv pip install -r requirements/build/cpu.txt --torch-backend cpu; \
248
+ fi
249
250
COPY . .
251
@@ -226,6 +256,18 @@ COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
256
257
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
258
259
+# For torch nightly, strip the torch/vision/audio/codec pins from the source
260
+# tree's requirements *after* `COPY . .` (which restored the pinned files) so
261
+# the wheel built below does not hard-pin the release torch (e.g.
262
+# torch==2.11.0+cpu). setup.py reads requirements/cpu.txt into install_requires;
263
+# leaving the pin makes `uv pip install dist/*.whl` unsatisfiable against the
264
+# nightly index. The nightly torch installed in the base image satisfies the
265
+# now-unpinned dependency. Done at the end of vllm-src so vllm-build and
266
+# vllm-dev, which both derive from it, inherit the unpinned tree.
267
+RUN if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
268
+ python3 /workspace/use_existing_torch.py --prefix; \
269
+ fi
270
+
271
######################### BUILD IMAGE #########################
272
FROM vllm-src AS vllm-build
273
@@ -283,6 +325,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
325
######################### TEST DEPS #########################
326
FROM base-arch AS vllm-test-deps
327
328
+# Re-declare torch-nightly args (ARGs do not cross FROM)
329
+ARG PYTORCH_NIGHTLY=0
330
+ARG PYTORCH_CPU_INDEX_BASE_URL=https://download.pytorch.org/whl
331
+
332
WORKDIR /vllm-workspace
333
334
# Test requirements are compiled from requirements/test/cuda.in into
@@ -297,11 +343,21 @@ RUN case "$(uname -m)" in \
343
esac
344
345
RUN --mount=type=cache,target=/root/.cache/uv \
300
- uv pip install -r requirements/test/cpu.txt --torch-backend cpu
346
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
347
+ python3 /workspace/use_existing_torch.py --prefix \
348
+ && uv pip install -r requirements/test/cpu.txt \
349
+ --extra-index-url ${PYTORCH_CPU_INDEX_BASE_URL}/nightly/cpu; \
350
+ else \
351
+ uv pip install -r requirements/test/cpu.txt --torch-backend cpu; \
352
+ fi
353
354
######################### DEV IMAGE #########################
355
FROM vllm-src AS vllm-dev
356
357
+# Re-declare torch-nightly args (ARGs do not cross FROM)
358
+ARG PYTORCH_NIGHTLY=0
359
+ARG PYTORCH_CPU_INDEX_BASE_URL=https://download.pytorch.org/whl
360
+
361
WORKDIR /vllm-workspace
362
363
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
@@ -323,7 +379,12 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement
379
380
RUN --mount=type=cache,target=/root/.cache/uv \
381
uv pip install -r requirements/lint.txt && \
326
- uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
382
+ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
383
+ uv pip install -r requirements/test/cpu.txt \
384
+ --extra-index-url ${PYTORCH_CPU_INDEX_BASE_URL}/nightly/cpu; \
385
+ else \
386
+ uv pip install -r requirements/test/cpu.txt --torch-backend cpu; \
387
+ fi && \
388
pre-commit install --hook-type pre-commit --hook-type commit-msg
389
390
# Installed after requirements/test/cpu.txt (which pins triton==3.6.0) so the
use_existing_torch.py
+2
@@ -11,10 +11,12 @@ TORCH_LIB_PREFIXES = (
11
"torch=",
12
"torchvision=",
13
"torchaudio=",
14
+ "torchcodec=",
15
# pyproject.toml
16
'"torch =',
17
'"torchvision =',
18
'"torchaudio =',
19
+ '"torchcodec =',
20
)
21
22