| 1 | #!/bin/bash |
| 2 | set -euox pipefail |
| 3 | export VLLM_CPU_CI_ENV=0 |
| 4 | export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM |
| 5 | |
| 6 | MODE=${1:-all} |
| 7 | |
| 8 | run_scenario() { |
| 9 | local label="$1" result_file="$2" |
| 10 | shift 2 |
| 11 | echo "--- $label" |
| 12 | vllm serve meta-llama/Llama-3.2-3B-Instruct "$@" --max-model-len=4096 & |
| 13 | local server_pid=$! |
| 14 | timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 |
| 15 | vllm bench serve \ |
| 16 | --backend vllm \ |
| 17 | --dataset-name random \ |
| 18 | --model meta-llama/Llama-3.2-3B-Instruct \ |
| 19 | --num-prompts 20 \ |
| 20 | --result-dir ./test_results \ |
| 21 | --result-filename "$result_file" \ |
| 22 | --save-result \ |
| 23 | --endpoint /v1/completions |
| 24 | kill -s SIGTERM "$server_pid"; wait "$server_pid" || true |
| 25 | if [ "$(jq '.failed' "./test_results/$result_file")" -ne 0 ]; then |
| 26 | echo "Some requests were failed in $label!" |
| 27 | exit 1 |
| 28 | fi |
| 29 | } |
| 30 | |
| 31 | case "$MODE" in |
| 32 | tp_pp) run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2 ;; |
| 33 | dp_tp) run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2 ;; |
| 34 | all) |
| 35 | run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2 |
| 36 | run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2 |
| 37 | ;; |
| 38 | *) echo "ERROR: unknown mode '$MODE' (expected: tp_pp | dp_tp | all)" >&2; exit 1 ;; |
| 39 | esac |