[CI/Build] Fix flaky failures in CPU CI image building

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -15,12 +15,15 @@ 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=${DISK_USAGE_THRESHOLD:-70}
-BUILDKIT_CACHE_MAX=${BUILDKIT_CACHE_MAX:-80GB}
+DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-80}
+BUILDKIT_CACHE_MAX=${BUILDKIT_CACHE_MAX:-200GB}
 
 # 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.
 prune_if_disk_pressure() {
     local docker_root disk_usage
     docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true)
@@ -31,7 +34,7 @@ prune_if_disk_pressure() {
     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" || true
+        docker builder prune -f --keep-storage="$BUILDKIT_CACHE_MAX" --filter type!=exec.cachemount || true
     else
         echo "Disk usage ${disk_usage:-unknown}% within ${DISK_USAGE_THRESHOLD}% threshold; skipping prune"
     fi
@@ -50,9 +53,35 @@ trap cleanup EXIT
 # Free space up front so a nearly-full host doesn't fail the build.
 prune_if_disk_pressure
 
+# Opportunistically warm the local base image cache; `docker build` below
+# omits `--pull` so it already prefers a cached image over the network. A
+# single attempt only -- the retry loop below covers a miss here too.
+echo "--- :docker: Pre-fetching base image"
+docker pull ubuntu:22.04 || true
+
 # building the docker image
 echo "--- :docker: Building Docker image"
-docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
+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_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 \
+            -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
+        rm -f "$build_log"
+        exit 1
+    fi
+    wait_s="${BUILD_RETRY_WAITS[$((attempt - 1))]}"
+    echo "--- :docker: Transient network error during build (attempt $attempt/$BUILD_MAX_ATTEMPTS), retrying in ${wait_s}s"
+    sleep "$wait_s"
+    attempt=$((attempt + 1))
+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" \

csrc/cpu/sgl-kernels/mla_cache.cpp

@@ -59,9 +59,9 @@ void concat_and_cache_mla_kernel_impl(
 
 // kv_c_normed : [num_tokens, kv_lora_rank]
 // k_pe        : [num_tokens, qk_rope_head_dim] or [num_tokens, 1,
-// qk_rope_head_dim] kv_cache    : [num_blocks, block_size, kv_lora_rank +
-// qk_rope_head_dim] slot_mapping: [num_tokens] int32/int64, absolute physical
-// row index
+//               qk_rope_head_dim]
+// kv_cache    : [num_blocks, block_size, kv_lora_rank + qk_rope_head_dim]
+// slot_mapping: [num_tokens] int32/int64, absolute physical row index
 //               (block_id * block_size + block_offset); negative entries are
 //               skipped (padded tokens), matching `concat_and_cache_mla`.
 void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed,

docker/Dockerfile.cpu

@@ -165,6 +165,9 @@ ARG USE_SCCACHE
 ARG SCCACHE_ENDPOINT
 ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
 ARG SCCACHE_REGION_NAME=us-west-2
+# When set, sccache skips the S3 backend and uses a local disk cache instead
+# (for builds without S3 credentials, e.g. a developer's local `docker build`).
+ARG SCCACHE_LOCAL_ONLY=0
 
 ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
 ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
@@ -186,10 +189,16 @@ ENV CARGO_BUILD_JOBS=4
 # inputs. The child stage below relinks the crates that embed the exact version.
 RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
     --mount=type=cache,target=/root/.cargo/git,sharing=locked \
+    --mount=type=cache,target=/root/.cache/sccache-local,sharing=locked \
     --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
     if [ "$USE_SCCACHE" = "1" ]; then \
         export RUSTC_WRAPPER=sccache; \
-        if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \
+        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; \
         sccache --show-stats; \
     fi && \
     SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+docker.cache" bash build_rust.sh && \
@@ -272,13 +281,20 @@ FROM vllm-src AS vllm-build
 
 ARG USE_SCCACHE
 ARG SCCACHE_ENDPOINT
+ARG SCCACHE_LOCAL_ONLY=0
 
 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=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
     if [ "$USE_SCCACHE" = "1" ]; then \
-        if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \
+        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; \
         sccache --show-stats; \
     fi && \
     VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 && \
@@ -296,22 +312,30 @@ 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 [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \
+            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; \

tests/models/language/pooling/test_max_tokens_per_doc.py

@@ -14,8 +14,6 @@
 from tests.utils import VLLM_PATH, RemoteOpenAIServer
 from vllm.entrypoints.pooling.scoring.protocol import RerankResponse
 
-os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
-
 TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template")
 ExpectedPromptTokens = int | tuple[int, ...]
 
@@ -128,7 +126,10 @@ def assert_prompt_tokens(actual: int, expected: ExpectedPromptTokens) -> None:
 @pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model)
 def server(request):
     config: TestConfig = request.param
-    with RemoteOpenAIServer(config.model, config.args) as remote_server:
+    env_dict = {"VLLM_LOGGING_LEVEL": "WARNING"}
+    with RemoteOpenAIServer(
+        config.model, config.args, env_dict=env_dict
+    ) as remote_server:
         yield config, remote_server