A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[TPU][Test] Add script to run benchmark on TPU for buildkite (#19039)
Signed-off-by: Qiliang Cui <derrhein@gmail.com>
QiliangCui committed
Jun 6, 2025 at 20:10 UTC
66c508b137acbd2d82590cbbf4d144e8c51271d1
4 files changed
+234
.buildkite/scripts/tpu/cleanup_docker.sh
new
+24
@@ -0,0 +1,24 @@
1
+#!/bin/bash
2
+
3
+set -euo pipefail
4
+
5
+docker_root=$(docker info -f '{{.DockerRootDir}}')
6
+if [ -z "$docker_root" ]; then
7
+ echo "Failed to determine Docker root directory."
8
+ exit 1
9
+fi
10
+echo "Docker root directory: $docker_root"
11
+# Check disk usage of the filesystem where Docker's root directory is located
12
+disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
13
+# Define the threshold
14
+threshold=70
15
+if [ "$disk_usage" -gt "$threshold" ]; then
16
+ echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
17
+ # Remove dangling images (those that are not tagged and not used by any container)
18
+ docker image prune -f
19
+ # Remove unused volumes / force the system prune for old images as well.
20
+ docker volume prune -f && docker system prune --force --filter "until=72h" --all
21
+ echo "Docker images and volumes cleanup completed."
22
+else
23
+ echo "Disk usage is below $threshold%. No cleanup needed."
24
+fi
.buildkite/scripts/tpu/config_v6e_1.env
new
+14
@@ -0,0 +1,14 @@
1
+# Environment config
2
+TEST_NAME=llama8b
3
+CONTAINER_NAME=vllm-tpu
4
+
5
+# vllm config
6
+MODEL=meta-llama/Llama-3.1-8B-Instruct
7
+MAX_NUM_SEQS=512
8
+MAX_NUM_BATCHED_TOKENS=512
9
+TENSOR_PARALLEL_SIZE=1
10
+MAX_MODEL_LEN=2048
11
+DOWNLOAD_DIR=/mnt/disks/persist
12
+EXPECTED_THROUGHPUT=8.0
13
+INPUT_LEN=1800
14
+OUTPUT_LEN=128
.buildkite/scripts/tpu/docker_run_bm.sh
new
+102
@@ -0,0 +1,102 @@
1
+#!/bin/bash
2
+
3
+if [ ! -f "$1" ]; then
4
+ echo "Error: The env file '$1' does not exist."
5
+ exit 1 # Exit the script with a non-zero status to indicate an error
6
+fi
7
+
8
+ENV_FILE=$1
9
+
10
+# For testing on local vm, use `set -a` to export all variables
11
+source /etc/environment
12
+source $ENV_FILE
13
+
14
+remove_docker_container() {
15
+ docker rm -f tpu-test || true;
16
+ docker rm -f vllm-tpu || true;
17
+ docker rm -f $CONTAINER_NAME || true;
18
+}
19
+
20
+trap remove_docker_container EXIT
21
+
22
+# Remove the container that might not be cleaned up in the previous run.
23
+remove_docker_container
24
+
25
+# Build docker image.
26
+# TODO: build the image outside the script and share the image with other
27
+# tpu test if building time is too long.
28
+DOCKER_BUILDKIT=1 docker build \
29
+ --build-arg max_jobs=16 \
30
+ --build-arg USE_SCCACHE=1 \
31
+ --build-arg GIT_REPO_CHECK=0 \
32
+ --tag vllm/vllm-tpu-bm \
33
+ --progress plain -f docker/Dockerfile.tpu .
34
+
35
+LOG_ROOT=$(mktemp -d)
36
+# If mktemp fails, set -e will cause the script to exit.
37
+echo "Results will be stored in: $LOG_ROOT"
38
+
39
+if [ -z "$HF_TOKEN" ]; then
40
+ echo "Error: HF_TOKEN is not set or is empty."
41
+ exit 1
42
+fi
43
+
44
+# Make sure mounted disk or dir exists
45
+if [ ! -d "$DOWNLOAD_DIR" ]; then
46
+ echo "Error: Folder $DOWNLOAD_DIR does not exist. This is useually a mounted drive. If no mounted drive, just create a folder."
47
+ exit 1
48
+fi
49
+
50
+echo "Run model $MODEL"
51
+echo
52
+
53
+echo "starting docker...$CONTAINER_NAME"
54
+echo
55
+docker run \
56
+ -v $DOWNLOAD_DIR:$DOWNLOAD_DIR \
57
+ --env-file $ENV_FILE \
58
+ -e HF_TOKEN="$HF_TOKEN" \
59
+ -e TARGET_COMMIT=$BUILDKITE_COMMIT \
60
+ -e MODEL=$MODEL \
61
+ -e WORKSPACE=/workspace \
62
+ --name $CONTAINER_NAME \
63
+ -d \
64
+ --privileged \
65
+ --network host \
66
+ -v /dev/shm:/dev/shm \
67
+ vllm/vllm-tpu-bm tail -f /dev/null
68
+
69
+echo "run script..."
70
+echo
71
+docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/hardware_ci/run_bm.sh"
72
+
73
+echo "copy result back..."
74
+VLLM_LOG="$LOG_ROOT/$TEST_NAME"_vllm_log.txt
75
+BM_LOG="$LOG_ROOT/$TEST_NAME"_bm_log.txt
76
+docker cp "$CONTAINER_NAME:/workspace/vllm_log.txt" "$VLLM_LOG"
77
+docker cp "$CONTAINER_NAME:/workspace/bm_log.txt" "$BM_LOG"
78
+
79
+throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g')
80
+echo "throughput for $TEST_NAME at $BUILDKITE_COMMIT: $throughput"
81
+
82
+if [ "$BUILDKITE" = "true" ]; then
83
+ echo "Running inside Buildkite"
84
+ buildkite-agent artifact upload "$VLLM_LOG"
85
+ buildkite-agent artifact upload "$BM_LOG"
86
+else
87
+ echo "Not running inside Buildkite"
88
+fi
89
+
90
+#
91
+# compare the throughput with EXPECTED_THROUGHPUT
92
+# and assert meeting the expectation
93
+#
94
+if [[ -z "$throughput" || ! "$throughput" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
95
+ echo "Failed to get the throughput"
96
+ exit 1
97
+fi
98
+
99
+if (( $(echo "$throughput < $EXPECTED_THROUGHPUT" | bc -l) )); then
100
+ echo "Error: throughput($throughput) is less than expected($EXPECTED_THROUGHPUT)"
101
+ exit 1
102
+fi
.buildkite/scripts/tpu/run_bm.sh
new
+94
@@ -0,0 +1,94 @@
1
+#!/bin/bash
2
+
3
+set -euo pipefail
4
+
5
+VLLM_LOG="$WORKSPACE/vllm_log.txt"
6
+BM_LOG="$WORKSPACE/bm_log.txt"
7
+
8
+if [ -n "$TARGET_COMMIT" ]; then
9
+ head_hash=$(git rev-parse HEAD)
10
+ if [ "$TARGET_COMMIT" != "$head_hash" ]; then
11
+ echo "Error: target commit $TARGET_COMMIT does not match HEAD: $head_hash"
12
+ exit 1
13
+ fi
14
+fi
15
+
16
+echo "model: $MODEL"
17
+echo
18
+
19
+#
20
+# create a log folder
21
+#
22
+mkdir "$WORKSPACE/log"
23
+
24
+# TODO: Move to image building.
25
+pip install pandas
26
+pip install datasets
27
+
28
+#
29
+# create sonnet_4x
30
+#
31
+echo "Create sonnet_4x.txt"
32
+echo "" > benchmarks/sonnet_4x.txt
33
+for _ in {1..4}
34
+ do
35
+ cat benchmarks/sonnet.txt >> benchmarks/sonnet_4x.txt
36
+done
37
+
38
+#
39
+# start vllm service in backend
40
+#
41
+echo "lanching vllm..."
42
+echo "logging to $VLLM_LOG"
43
+echo
44
+
45
+VLLM_USE_V1=1 vllm serve $MODEL \
46
+ --seed 42 \
47
+ --disable-log-requests \
48
+ --max-num-seqs $MAX_NUM_SEQS \
49
+ --max-num-batched-tokens $MAX_NUM_BATCHED_TOKENS \
50
+ --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
51
+ --no-enable-prefix-caching \
52
+ --download_dir $DOWNLOAD_DIR \
53
+ --max-model-len $MAX_MODEL_LEN > "$VLLM_LOG" 2>&1 &
54
+
55
+
56
+echo "wait for 20 minutes.."
57
+echo
58
+# sleep 1200
59
+# wait for 10 minutes...
60
+for i in {1..120}; do
61
+ # TODO: detect other type of errors.
62
+ if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then
63
+ echo "Detected RuntimeError, exiting."
64
+ exit 1
65
+ elif grep -Fq "Application startup complete" "$VLLM_LOG"; then
66
+ echo "Application started"
67
+ break
68
+ else
69
+ echo "wait for 10 seconds..."
70
+ sleep 10
71
+ fi
72
+done
73
+
74
+#
75
+# run test
76
+#
77
+echo "run benchmark test..."
78
+echo "logging to $BM_LOG"
79
+echo
80
+python benchmarks/benchmark_serving.py \
81
+ --backend vllm \
82
+ --model $MODEL \
83
+ --dataset-name sonnet \
84
+ --dataset-path benchmarks/sonnet_4x.txt \
85
+ --sonnet-input-len $INPUT_LEN \
86
+ --sonnet-output-len $OUTPUT_LEN \
87
+ --ignore-eos > "$BM_LOG"
88
+
89
+echo "completed..."
90
+echo
91
+
92
+throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g')
93
+echo "throughput: $throughput"
94
+echo