[CI] Shard CPU jobs above the 24h P90 threshold (#54751)
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com>
Kevin H. Luu committed
Sep 2, 2026 at 02:16 UTC
76ba32160a501e5e8aadb5e1820c51c765714220
3 files changed
+61
-40
.buildkite/hardware_tests/cpu.yaml
+13
-9
@@ -62,7 +62,7 @@ steps:
62
# bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
63
# bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
64
65
-- label: CPU-Language Generation and Pooling Model Tests
65
+- label: CPU-Language Generation and Pooling Model Tests Shard %N
66
key: cpu-language-generation-and-pooling-model-tests
67
depends_on: []
68
device: intel_cpu
@@ -75,10 +75,9 @@ steps:
75
- tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
76
commands:
77
- |
78
- bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m "
79
- pytest -x -v -s tests/models/language/generation -m cpu_model
80
- pytest -x -v -s tests/models/language/pooling -m cpu_model
81
- pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py"
78
+ bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 25m "
79
+ pytest -x -v -s tests/models/language/generation tests/models/language/pooling tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
80
+ parallelism: 3
81
82
- label: CPU-Quantization Model Tests
83
key: cpu-quantization-model-tests
@@ -162,11 +161,16 @@ steps:
161
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
162
VLLM_CI_ENV=0 pytest -x -v -s tests/models/multimodal/generation/test_qwen2_5_vl.py"
163
165
-- label: "Arm CPU Test"
164
+- label: "Arm CPU Test Shard %N"
165
key: arm-cpu-test
167
- depends_on: []
166
+ depends_on:
167
+ - cpu-arm64-image-build
168
soft_fail: false
169
device: arm_cpu
170
no_plugin: true
171
- commands:
172
- - bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
171
+ commands:
172
+ - >-
173
+ bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
174
+ "$$BUILDKITE_PARALLEL_JOB"
175
+ "$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64-cpu"
176
+ parallelism: 3
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+46
-31
@@ -1,8 +1,17 @@
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 -ex
3
+# Run the Arm CPU test suites against the image built by cpu-arm64-image-build.
4
+set -euxo pipefail
5
+
6
+SHARD_ID=${1:?Usage: run-cpu-test-arm.sh SHARD_ID IMAGE}
7
+IMAGE=${2:?Usage: run-cpu-test-arm.sh SHARD_ID IMAGE}
8
+if [[ ! "$SHARD_ID" =~ ^[0-2]$ ]]; then
9
+ echo "SHARD_ID must be 0, 1, or 2" >&2
10
+ exit 2
11
+fi
12
+
13
+JOB_SUFFIX=${BUILDKITE_JOB_ID:-local-$SHARD_ID}
14
+CONTAINER_NAME="cpu-test-${JOB_SUFFIX//[^a-zA-Z0-9_.-]/-}"
15
16
# allow to bind to different cores
17
CORE_RANGE=${CORE_RANGE:-0-31}
@@ -13,26 +22,25 @@ export CMAKE_BUILD_PARALLEL_LEVEL=32
22
# Setup cleanup
23
remove_docker_container() {
24
set -e;
16
- docker rm -f cpu-test || true;
25
+ docker rm -f "$CONTAINER_NAME" || true;
26
}
27
trap remove_docker_container EXIT
28
remove_docker_container
29
21
-# Try building the docker image
22
-docker build --tag cpu-test --target vllm-test -f docker/Dockerfile.cpu .
30
+docker pull "$IMAGE"
31
32
# Run the image
25
-docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test cpu-test
26
-
27
-function cpu_tests() {
28
- set -e
33
+docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name "$CONTAINER_NAME" "$IMAGE"
34
30
- docker exec cpu-test bash -c "
35
+print_packages() {
36
+ docker exec "$CONTAINER_NAME" bash -c "
37
set -e
38
pip list"
39
+}
40
34
- # Run kernel tests
35
- docker exec cpu-test bash -c "
41
+kernel_tests() {
42
+ set -e
43
+ docker exec "$CONTAINER_NAME" bash -c "
44
set -e
45
pytest -x -v -s tests/kernels/test_onednn.py
46
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
@@ -43,32 +51,37 @@ function cpu_tests() {
51
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
52
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
53
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
54
+}
55
47
- # skip tests requiring model downloads if HF_TOKEN is not set
48
- # due to rate-limits
49
- if [ -z "$HF_TOKEN" ]; then
56
+model_tests() {
57
+ set -e
58
+ if [ -z "${HF_TOKEN:-}" ]; then
59
echo "Warning: HF_TOKEN is not set. Skipping tests that require model downloads."
60
return
61
fi
62
54
- # offline inference
55
- docker exec cpu-test bash -c "
63
+ docker exec "$CONTAINER_NAME" bash -c "
64
set -e
65
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m"
66
59
- # Test encoder-decoder and encoder-only models
60
- docker exec cpu-test bash -c "
67
+ docker exec "$CONTAINER_NAME" bash -c "
68
set -e
69
pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model
70
pytest -x -v -s 'tests/models/language/pooling/test_embedding.py::test_models[sentence-transformers/all-MiniLM-L12-v2]'"
71
65
- # Run quantized model tests
66
- docker exec cpu-test bash -c "
72
+ docker exec "$CONTAINER_NAME" bash -c "
73
set -e
74
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
75
+}
76
+
77
+serving_tests() {
78
+ set -e
79
+ if [ -z "${HF_TOKEN:-}" ]; then
80
+ echo "Warning: HF_TOKEN is not set. Skipping tests that require model downloads."
81
+ return
82
+ fi
83
70
- # basic online serving
71
- docker exec cpu-test bash -c '
84
+ docker exec "$CONTAINER_NAME" bash -c '
85
set -e
86
VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3-0.6B --max-model-len 2048 &
87
server_pid=$!
@@ -81,8 +94,7 @@ function cpu_tests() {
94
--endpoint /v1/completions
95
kill -s SIGTERM $server_pid &'
96
84
- # smoke test for Gated DeltaNet
85
- docker exec cpu-test bash -c '
97
+ docker exec "$CONTAINER_NAME" bash -c '
98
set -e
99
VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3.5-0.8B --max-model-len 2048 &
100
server_pid=$!
@@ -94,10 +106,13 @@ function cpu_tests() {
106
--num-prompts 20 \
107
--endpoint /v1/completions
108
kill -s SIGTERM $server_pid &'
97
-
109
}
110
100
-# All of CPU tests are expected to be finished less than 40 mins.
101
-export -f cpu_tests
102
-timeout 2h bash -c cpu_tests
103
-
111
+print_packages
112
+export CONTAINER_NAME
113
+export -f kernel_tests model_tests serving_tests
114
+case "$SHARD_ID" in
115
+ 0) timeout 30m bash -c kernel_tests ;;
116
+ 1) timeout 30m bash -c model_tests ;;
117
+ 2) timeout 30m bash -c serving_tests ;;
118
+esac
tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
+2
@@ -10,6 +10,8 @@ from tests.models.utils import check_logprobs_close
10
from vllm import LLM, SamplingParams
11
from vllm.platforms import current_platform
12
13
+pytestmark = pytest.mark.cpu_model
14
+
15
if not current_platform.is_cpu():
16
pytest.skip("skipping CPU-only tests", allow_module_level=True)
17