[CI/Build][CPU] Fix flaky rust downloads, broken prune flag, and triton-cpu cache coupling (#56247)

Signed-off-by: jiang1.li <jiang1.li@intel.com>
This commit is contained in:
Li, Jiang
2026-09-10 22:27:26 +08:00
committed by GitHub
parent 2e0ee66cab
commit 46bae7e98c
3 changed files with 111 additions and 68 deletions
+32 -16
View File
@@ -13,17 +13,16 @@ TIMEOUT_VAL=$1
TEST_COMMAND=$2
# Disk hygiene knobs. Reclaim space only once the Docker root filesystem crosses
# DISK_USAGE_THRESHOLD percent, and cap the shared BuildKit cache at
# BUILDKIT_CACHE_MAX so subsequent builds keep reusing the hottest layers.
# DISK_USAGE_THRESHOLD percent, and only prune images/cache unused for at least
# CACHE_MAX_AGE so recently-built layers survive for reuse.
DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-80}
BUILDKIT_CACHE_MAX=${BUILDKIT_CACHE_MAX:-200GB}
CACHE_MAX_AGE=${CACHE_MAX_AGE:-24h}
# Reclaim disk only when the host is under pressure. We trim (not purge) the
# shared BuildKit cache so cross-job/cross-agent reuse stays intact, and only
# touch dangling images; other agents' uniquely tagged images are left alone.
# Cache mounts (exec.cachemount: uv/cargo/apt) are excluded so they survive
# pruning -- they're what let installs reuse packages instead of hitting the
# network, and are otherwise reclaimable like any other build cache record.
# Reclaim disk only when the host is under pressure, aging out anything unused
# for less than CACHE_MAX_AGE so hot layers survive -- same `--filter
# until=<N>h` pattern the TPU CI scripts already rely on. `docker buildx
# prune --max-used-space` is a no-op on this host's `docker` driver (BuildKit
# embedded in dockerd never enforces the size cap), so we don't use it.
prune_if_disk_pressure() {
local docker_root disk_usage
docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true)
@@ -33,8 +32,7 @@ prune_if_disk_pressure() {
disk_usage=$(df "$docker_root" 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%')
if [ "${disk_usage:-0}" -gt "$DISK_USAGE_THRESHOLD" ]; then
echo "--- :broom: Disk usage ${disk_usage}% exceeds ${DISK_USAGE_THRESHOLD}%, reclaiming space"
docker image prune -f || true
docker builder prune -f --keep-storage="$BUILDKIT_CACHE_MAX" --filter type!=exec.cachemount || true
docker system prune --force --all --filter "until=${CACHE_MAX_AGE}" || true
else
echo "Disk usage ${disk_usage:-unknown}% within ${DISK_USAGE_THRESHOLD}% threshold; skipping prune"
fi
@@ -61,18 +59,18 @@ docker pull ubuntu:25.04 || true
# building the docker image
echo "--- :docker: Building Docker image"
BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:'
BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:'
BUILD_MAX_ATTEMPTS=4 # 1 initial + 3 retries
BUILD_RETRY_WAITS=(10 20 40) # seconds to wait before retry 1/2/3
build_log="$(mktemp)"
attempt=1
while true; do
if docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test \
--build-arg USE_SCCACHE=1 --build-arg SCCACHE_LOCAL_ONLY=1 \
--build-arg USE_SCCACHE=1 --build-arg SCCACHE_LOCAL_ONLY=1 --build-arg max_jobs=16 \
-f docker/Dockerfile.cpu . 2>&1 | tee "$build_log"; then
break
fi
if [ "$attempt" -ge "$BUILD_MAX_ATTEMPTS" ] || ! grep -qE "$BUILD_RETRY_PATTERN" "$build_log"; then
if [ "$attempt" -ge "$BUILD_MAX_ATTEMPTS" ] || ! grep -qiE "$BUILD_RETRY_PATTERN" "$build_log"; then
rm -f "$build_log"
exit 1
fi
@@ -83,6 +81,24 @@ while true; do
done
rm -f "$build_log"
# Run the image, setting --shm-size=4g for tensor parallel.
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
# Run the image, setting --shm-size=4g for tensor parallel. Default to
# HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network;
# retry once online if the cache is missing something.
OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)'
run_test() {
local hf_offline=$1
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
}
test_log="$(mktemp)"
if run_test 1 2>&1 | tee "$test_log"; then
rm -f "$test_log"
elif grep -qE "$OFFLINE_RETRY_PATTERN" "$test_log"; then
rm -f "$test_log"
echo "--- :warning: HF_HUB_OFFLINE caused a cache miss, retrying with online fallback"
run_test 0
else
rm -f "$test_log"
exit 1
fi
+60 -49
View File
@@ -105,6 +105,60 @@ ENV UV_HTTP_TIMEOUT=500
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
ENV TARGETARCH=${TARGETARCH}
######################### TRITON-CPU BUILD IMAGE #########################
# Branches off base-common directly (before requirements/cpu.txt is installed
# below) so a torch/etc. version bump doesn't invalidate this stage's cache --
# it builds triton-cpu in an isolated `uv build` environment and never touches
# these packages.
FROM base-common AS vllm-triton-cpu-build
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
# Re-declared here because this stage is `FROM base-common` (not `vllm-build`),
# so it does not inherit the ARG/ENV defined there. Without it, the guard
# below would see an empty value and build triton-cpu on non-x86 targets
# (e.g. arm64).
ARG VLLM_CPU_X86=0
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
ARG SCCACHE_LOCAL_ONLY=0
WORKDIR /vllm-workspace
RUN mkdir dist
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/sccache-local,sharing=locked \
--mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
--mount=type=cache,target=/root/.triton \
--mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
set -ex; \
if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
cd triton-cpu; \
git checkout "270e696d"; \
if [ "$USE_SCCACHE" = "1" ]; then \
if [ "$SCCACHE_LOCAL_ONLY" = "1" ]; then \
unset SCCACHE_BUCKET SCCACHE_REGION; \
export SCCACHE_DIR=/root/.cache/sccache-local SCCACHE_CACHE_SIZE=20G; \
elif [ -n "${SCCACHE_ENDPOINT}" ]; then \
export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; \
fi; \
export TRITON_APPEND_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache"; \
sccache --show-stats; \
fi; \
uv build --wheel --out-dir=../dist; \
if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi; \
cd ..; \
rm -rf triton-cpu; \
fi
######################### COMMON BASE IMAGE + PYTHON DEPS #########################
FROM base-common AS base-common-pydeps
# PyTorch provides its own indexes for standard and nightly builds.
# PYTORCH_NIGHTLY=1 builds the CPU image against torch nightly (parity with the
# GPU torch-nightly lane); default 0 keeps the pinned release CPU build.
@@ -133,15 +187,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/cpu.txt --torch-backend cpu; \
fi
ENV TARGETARCH=${TARGETARCH}
######################### x86_64 BASE IMAGE #########################
FROM base-common AS base-amd64
FROM base-common-pydeps AS base-amd64
ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4:/opt/venv/lib/libiomp5.so"
######################### arm64 BASE IMAGE #########################
FROM base-common AS base-arm64
FROM base-common-pydeps AS base-arm64
ENV LD_PRELOAD="/usr/lib/aarch64-linux-gnu/libtcmalloc_minimal.so.4"
@@ -210,6 +262,10 @@ RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+docker.cache" bash build_rust.sh && \
if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
# build_rust.sh installed rustup here via rustup.rs; put it on PATH so the
# child stage below finds it on disk instead of re-downloading it.
ENV PATH="/root/.cargo/bin:${PATH}"
# Relink with the exact Git-derived package version.
FROM rust-build-cache AS rust-build
@@ -308,51 +364,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 && \
if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
######################### TRITON-CPU BUILD IMAGE #########################
FROM base-arch AS vllm-triton-cpu-build
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
# Re-declared here because this stage is `FROM base-arch` (not `vllm-build`),
# so it does not inherit the ARG/ENV defined there. Without it, the guard
# below would see an empty value and build triton-cpu on non-x86 targets
# (e.g. arm64).
ARG VLLM_CPU_X86=0
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
ARG SCCACHE_LOCAL_ONLY=0
WORKDIR /vllm-workspace
RUN mkdir dist
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/sccache-local,sharing=locked \
--mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
--mount=type=cache,target=/root/.triton \
--mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
set -ex; \
if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
cd triton-cpu; \
git checkout "270e696d"; \
if [ "$USE_SCCACHE" = "1" ]; then \
if [ "$SCCACHE_LOCAL_ONLY" = "1" ]; then \
unset SCCACHE_BUCKET SCCACHE_REGION; \
export SCCACHE_DIR=/root/.cache/sccache-local SCCACHE_CACHE_SIZE=20G; \
elif [ -n "${SCCACHE_ENDPOINT}" ]; then \
export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; \
fi; \
export TRITON_APPEND_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache"; \
sccache --show-stats; \
fi; \
uv build --wheel --out-dir=../dist; \
if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi; \
cd ..; \
rm -rf triton-cpu; \
fi
######################### TEST DEPS #########################
FROM base-arch AS vllm-test-deps
+19 -3
View File
@@ -12,7 +12,8 @@ from typing import Any
import psutil
import torch
from vllm.config import VllmConfig
from vllm import envs
from vllm.config import CompilationMode, VllmConfig
from vllm.logger import init_logger
from vllm.platforms import CpuArchEnum, current_platform
from vllm.profiler.wrapper import TorchProfilerWrapper
@@ -197,8 +198,23 @@ class CPUWorker(Worker):
logger.warning("sleep mode is not supported on CPU, ignore it.")
pass
def _should_warm_up_model(self) -> bool:
# VLLM_CPU_CI_ENV always skips warmup to save CI time.
if envs.VLLM_CPU_CI_ENV:
return False
# With eager (CompilationMode.NONE) execution and an explicit KV
# cache size, warmup serves no purpose: there's no compiled graph to
# prewarm, and the auto KV-cache-size calc below -- which relies on
# a warmup forward pass having already bumped RSS to a realistic
# steady-state value -- is bypassed whenever the size is explicit.
return not (
self.compilation_config.mode == CompilationMode.NONE
and self.cache_config.kv_cache_memory_bytes is not None
)
def determine_available_memory(self) -> int:
self.model_runner.warming_up_model()
if self._should_warm_up_model():
self.model_runner.warming_up_model()
allowed_cpu_list = get_allowed_cpu_list()
cpu_core = allowed_cpu_list[0]
@@ -259,7 +275,7 @@ class CPUWorker(Worker):
def compile_or_warm_up_model(self) -> CompilationTimes:
# Note: the model has been compiled in determine_available_memory(),
# Only compile here for models without kv cache
if len(self.model_runner.kv_caches) == 0:
if len(self.model_runner.kv_caches) == 0 and self._should_warm_up_model():
self.model_runner.warming_up_model()
# Reset the seed to ensure that the random state is not affected by
# the model initialization and profiling.