| 1 | #!/bin/bash |
| 2 | |
| 3 | # This script build the CPU docker image and run the offline inference inside the container. |
| 4 | # It serves a sanity check for compilation and basic model usage. |
| 5 | set -euox pipefail |
| 6 | |
| 7 | # allow to bind to different cores |
| 8 | CORE_RANGE=${CORE_RANGE:-48-95} |
| 9 | NUMA_NODE=${NUMA_NODE:-1} |
| 10 | AGENT_SLOT=${AGENT_SLOT:-} |
| 11 | IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}" |
| 12 | TIMEOUT_VAL=$1 |
| 13 | TEST_COMMAND=$2 |
| 14 | |
| 15 | # Disk hygiene knobs. Reclaim space only once the Docker root filesystem crosses |
| 16 | # DISK_USAGE_THRESHOLD percent, and only prune images/cache unused for at least |
| 17 | # CACHE_MAX_AGE so recently-built layers survive for reuse. |
| 18 | DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-80} |
| 19 | CACHE_MAX_AGE=${CACHE_MAX_AGE:-24h} |
| 20 | |
| 21 | # Reclaim disk only when the host is under pressure, aging out anything unused |
| 22 | # for less than CACHE_MAX_AGE so hot layers survive -- same `--filter |
| 23 | # until=<N>h` pattern the TPU CI scripts already rely on. `docker buildx |
| 24 | # prune --max-used-space` is a no-op on this host's `docker` driver (BuildKit |
| 25 | # embedded in dockerd never enforces the size cap), so we don't use it. |
| 26 | prune_if_disk_pressure() { |
| 27 | local docker_root disk_usage |
| 28 | docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true) |
| 29 | if [ -z "$docker_root" ]; then |
| 30 | return 0 |
| 31 | fi |
| 32 | disk_usage=$(df "$docker_root" 2>/dev/null | tail -1 | awk '{print $5}' | tr -d '%') |
| 33 | if [ "${disk_usage:-0}" -gt "$DISK_USAGE_THRESHOLD" ]; then |
| 34 | echo "--- :broom: Disk usage ${disk_usage}% exceeds ${DISK_USAGE_THRESHOLD}%, reclaiming space" |
| 35 | docker system prune --force --all --filter "until=${CACHE_MAX_AGE}" || true |
| 36 | else |
| 37 | echo "Disk usage ${disk_usage:-unknown}% within ${DISK_USAGE_THRESHOLD}% threshold; skipping prune" |
| 38 | fi |
| 39 | } |
| 40 | |
| 41 | # Always drop this agent's image once the job ends (the default builder never |
| 42 | # uses it as a cache source, so removing it costs no rebuild speed), then |
| 43 | # reclaim space if needed. Guard every docker call with `|| true` so the trap |
| 44 | # never overrides the test's exit code. |
| 45 | cleanup() { |
| 46 | docker image rm -f "$IMAGE_NAME" || true |
| 47 | prune_if_disk_pressure |
| 48 | } |
| 49 | trap cleanup EXIT |
| 50 | |
| 51 | # Free space up front so a nearly-full host doesn't fail the build. |
| 52 | prune_if_disk_pressure |
| 53 | |
| 54 | # Opportunistically warm the local base image cache; `docker build` below |
| 55 | # omits `--pull` so it already prefers a cached image over the network. A |
| 56 | # single attempt only -- the retry loop below covers a miss here too. |
| 57 | echo "--- :docker: Pre-fetching base image" |
| 58 | docker pull ubuntu:25.04 || true |
| 59 | |
| 60 | # building the docker image |
| 61 | echo "--- :docker: Building Docker image" |
| 62 | 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:|The read operation timed out|BrokenPipeError:.*Broken pipe|HTTP/[0-9.]+ stream [0-9]+ was not closed cleanly|fetch-pack: unexpected disconnect|fatal: early EOF|invalid index-pack output' |
| 63 | BUILD_MAX_ATTEMPTS=4 # 1 initial + 3 retries |
| 64 | BUILD_RETRY_WAITS=(10 20 40) # seconds to wait before retry 1/2/3 |
| 65 | build_log="$(mktemp)" |
| 66 | attempt=1 |
| 67 | while true; do |
| 68 | if docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test \ |
| 69 | --build-arg USE_SCCACHE=1 --build-arg SCCACHE_LOCAL_ONLY=1 --build-arg max_jobs=16 \ |
| 70 | -f docker/Dockerfile.cpu . 2>&1 | tee "$build_log"; then |
| 71 | break |
| 72 | fi |
| 73 | if [ "$attempt" -ge "$BUILD_MAX_ATTEMPTS" ] || ! grep -qiE "$BUILD_RETRY_PATTERN" "$build_log"; then |
| 74 | rm -f "$build_log" |
| 75 | exit 1 |
| 76 | fi |
| 77 | wait_s="${BUILD_RETRY_WAITS[$((attempt - 1))]}" |
| 78 | echo "--- :docker: Transient network error during build (attempt $attempt/$BUILD_MAX_ATTEMPTS), retrying in ${wait_s}s" |
| 79 | sleep "$wait_s" |
| 80 | attempt=$((attempt + 1)) |
| 81 | done |
| 82 | rm -f "$build_log" |
| 83 | |
| 84 | # Run the image, setting --shm-size=4g for tensor parallel. Default to |
| 85 | # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network; |
| 86 | # retry once online if the cache is missing something. vllm wraps offline |
| 87 | # cache-miss errors in ways that don't preserve the raw huggingface_hub |
| 88 | # exception name: get_config() re-raises as a generic ValueError |
| 89 | # (transformers_utils/config.py), and default_loader.py raises its own |
| 90 | # RuntimeError when a snapshot dir has a config but no weight files. Match |
| 91 | # those messages too or the fallback never triggers for those cache misses. |
| 92 | # |
| 93 | # Also mount ~/.cache/vllm: multimodal test fixtures (e.g. VideoAsset, via |
| 94 | # vllm/assets/video.py) download into VLLM_ASSETS_CACHE (~/.cache/vllm/assets |
| 95 | # by default), a directory distinct from the HF hub cache above. Without a |
| 96 | # persistent mount there, those assets never survive past one container's |
| 97 | # lifetime, so every run re-fetches them from the network on the "offline" |
| 98 | # attempt and unconditionally falls back to the online retry. |
| 99 | OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with' |
| 100 | run_test() { |
| 101 | local hf_offline=$1 |
| 102 | docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --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" -e TERM=xterm-256color -e PY_COLORS=1 -e FORCE_COLOR=1 -e CLICOLOR_FORCE=1 --shm-size=4g "$IMAGE_NAME" \ |
| 103 | timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" |
| 104 | } |
| 105 | |
| 106 | test_log="$(mktemp)" |
| 107 | if run_test 1 2>&1 | tee "$test_log"; then |
| 108 | rm -f "$test_log" |
| 109 | elif grep -qE "$OFFLINE_RETRY_PATTERN" "$test_log"; then |
| 110 | rm -f "$test_log" |
| 111 | echo "--- :warning: HF_HUB_OFFLINE caused a cache miss, retrying with online fallback" |
| 112 | run_test 0 |
| 113 | else |
| 114 | rm -f "$test_log" |
| 115 | exit 1 |
| 116 | fi |