@@ -10,15 +10,24 @@ set -x
1010set -o pipefail
1111
1212check_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