@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 120 lines 3.91 KB
Raw
1 #!/bin/bash
2
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}
18 OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-31}
19
20 export CMAKE_BUILD_PARALLEL_LEVEL=32
21
22 # Setup cleanup
23 remove_docker_container() {
24 set -e;
25 docker rm -f "$CONTAINER_NAME" || true;
26 }
27 trap remove_docker_container EXIT
28 remove_docker_container
29
30 docker pull "$IMAGE"
31
32 # Run the image
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" -e TERM=xterm-256color -e PY_COLORS=1 -e FORCE_COLOR=1 -e CLICOLOR_FORCE=1 --shm-size=4g --name "$CONTAINER_NAME" "$IMAGE"
34
35 print_packages() {
36 docker exec "$CONTAINER_NAME" bash -c "
37 set -e
38 pip list"
39 }
40
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
47 pytest -x -v -s tests/kernels/core/test_cpu_activation.py
48 pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
49 pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
50 pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
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
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
63 docker exec "$CONTAINER_NAME" bash -c "
64 set -e
65 python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m"
66
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
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
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=$!
88 timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
89 vllm bench serve \
90 --backend vllm \
91 --dataset-name random \
92 --model Qwen/Qwen3-0.6B \
93 --num-prompts 20 \
94 --endpoint /v1/completions
95 kill -s SIGTERM $server_pid &'
96
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=$!
101 timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
102 vllm bench serve \
103 --backend vllm \
104 --dataset-name random \
105 --model Qwen/Qwen3.5-0.8B \
106 --num-prompts 20 \
107 --endpoint /v1/completions
108 kill -s SIGTERM $server_pid &'
109 }
110
111 print_packages
112 export CONTAINER_NAME
113 export -f kernel_tests model_tests serving_tests
114 case "$SHARD_ID" in
115 # The kernel shard contains several large parametrized suites and can exceed
116 # 30 minutes on otherwise healthy Arm runners.
117 0) timeout 60m bash -c kernel_tests ;;
118 1) timeout 30m bash -c model_tests ;;
119 2) timeout 30m bash -c serving_tests ;;
120 esac