@@ -10,15 +10,24 @@ set -x
10
10
set -o pipefail
11
11
12
12
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
+
15
20
if [[ $gpu_count -gt 0 ]]; then
16
21
echo " GPU found."
17
22
else
18
23
echo " Need at least 1 GPU to run benchmarking."
19
24
exit 1
20
25
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
22
31
echo " GPU type is $gpu_type "
23
32
}
24
33
@@ -90,9 +99,15 @@ kill_gpu_processes() {
90
99
91
100
92
101
# 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
96
111
97
112
# remove vllm config file
98
113
rm -rf ~ /.config/vllm
0 commit comments