Skip to content

Commit e2d1e1e

Browse files
huydhnnishith-fujitsu
authored andcommitted
[ROCm] Tweak the benchmark script to run on ROCm (vllm-project#14252)
1 parent 9f8ee65 commit e2d1e1e

File tree

1 file changed

+21
-6
lines changed

1 file changed

+21
-6
lines changed

.buildkite/nightly-benchmarks/scripts/run-performance-benchmarks.sh

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,24 @@ set -x
1010
set -o pipefail
1111

1212
check_gpus() {
13-
# check the number of GPUs and GPU type.
14-
declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
13+
if command -v nvidia-smi; then
14+
# check the number of GPUs and GPU type.
15+
declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
16+
elif command -v amd-smi; then
17+
declare -g gpu_count=$(amd-smi list | grep 'GPU' | wc -l)
18+
fi
19+
1520
if [[ $gpu_count -gt 0 ]]; then
1621
echo "GPU found."
1722
else
1823
echo "Need at least 1 GPU to run benchmarking."
1924
exit 1
2025
fi
21-
declare -g gpu_type=$(nvidia-smi --query-gpu=name --format=csv,noheader | awk '{print $2}')
26+
if command -v nvidia-smi; then
27+
declare -g gpu_type=$(nvidia-smi --query-gpu=name --format=csv,noheader | awk '{print $2}')
28+
elif command -v amd-smi; then
29+
declare -g gpu_type=$(amd-smi static -g 0 -a | grep 'MARKET_NAME' | awk '{print $2}')
30+
fi
2231
echo "GPU type is $gpu_type"
2332
}
2433

@@ -90,9 +99,15 @@ kill_gpu_processes() {
9099

91100

92101
# wait until GPU memory usage smaller than 1GB
93-
while [ "$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -n 1)" -ge 1000 ]; do
94-
sleep 1
95-
done
102+
if command -v nvidia-smi; then
103+
while [ "$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -n 1)" -ge 1000 ]; do
104+
sleep 1
105+
done
106+
elif command -v amd-smi; then
107+
while [ "$(amd-smi metric -g 0 | grep 'USED_VRAM' | awk '{print $2}')" -ge 1000 ]; do
108+
sleep 1
109+
done
110+
fi
96111

97112
# remove vllm config file
98113
rm -rf ~/.config/vllm

0 commit comments

Comments
 (0)