通过跳板机批量管理GPU服务器 - peter-xbs/CommonCodes GitHub Wiki

1. 通过跳板机监测GPU服务器使用情况

背景：200台GPU服务器，多人共用，但哪台GPU服务器空闲不太确定，如何快速找出闲置GPU服务器

#!/bin/bash

# 文件名包含目标服务器的IP地址
SERVER_LIST="gpu_servers.txt"
# 输出符合要求的服务器的IP地址文件
OUTPUT_FILE="idle_gpu_servers.txt"
# 日志文件
LOG_FILE="gpu_check.log"

# 检查服务器列表文件是否存在
if [ ! -f "$SERVER_LIST" ]; then
  echo "服务器列表文件 '$SERVER_LIST' 不存在，请确认文件路径。"
  exit 1
fi

# 清空输出文件和日志文件
> "$OUTPUT_FILE"
> "$LOG_FILE"

# 显存占用率阈值
MEMORY_USAGE_THRESHOLD=5

# 检查服务器GPU显存使用情况
check_server_gpu_usage() {
  local server_ip=$1

  echo "检查服务器 $server_ip 的 GPU 使用情况..." | tee -a "$LOG_FILE"

  # 通过 SSH 获取 GPU 使用情况，使用 `timeout` 限制执行时间为10秒
  gpu_info=$(ssh -o BatchMode=yes -o ConnectTimeout=5 "$server_ip" "timeout 10s nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader,nounits; exit" 2>/dev/null)

  # 检查 SSH 或者 nvidia-smi 命令是否成功执行
  if [ $? -ne 0 ]; then
    echo "无法连接到服务器 $server_ip 或者无法获取 GPU 状态信息。" | tee -a "$LOG_FILE"
    return
  fi

  # 检查每个 GPU 的显存占用率，只要有一个GPU符合条件，就记录服务器IP
  is_idle=true
  IFS=$'\n'
  for line in $gpu_info; do
    memory_used=$(echo "$line" | awk -F, '{print $1}')
    memory_total=$(echo "$line" | awk -F, '{print $2}')
    memory_usage=$((100 * memory_used / memory_total))

    if [ "$memory_usage" -gt "$MEMORY_USAGE_THRESHOLD" ](/peter-xbs/CommonCodes/wiki/-"$memory_usage"--gt-"$MEMORY_USAGE_THRESHOLD"-); then
      is_idle=false
      break 
    fi
  done
  unset IFS

  if [ "$is_idle" = true ]; then
    echo "推荐：服务器 $server_ip 的所有 GPU 显存占用率均低于 $MEMORY_USAGE_THRESHOLD%." | tee -a "$LOG_FILE"
    echo "$server_ip" >> "$OUTPUT_FILE"
  else
    echo "服务器 $server_ip 不满足推荐条件." | tee -a "$LOG_FILE"
  fi
}

# 并发执行每台服务器的检查
export -f check_server_gpu_usage
export MEMORY_USAGE_THRESHOLD
export OUTPUT_FILE
export LOG_FILE

# 使用 xargs 并发执行，最多同时处理 30 台服务器
xargs -a "$SERVER_LIST" -I {} -P 30 bash -c 'check_server_gpu_usage "$@"' _ {}

echo "符合条件的服务器 IP 已写入文件：$OUTPUT_FILE"
exit 0

2. 通过跳板机批量向GPU服务器复制环境

背景：跳板机准备了一个新的nemo2.tar的环境包，批量复制到每台GPU服务器解压使用

#!/bin/bash

# servers.txt 文件包含了所有服务器的 IP 地址或主机名
SERVERS_FILE="gpu_servers.txt"

# 要复制和解压的文件名
TAR_FILE="nemo2_slim.tar.gz"

# 目标路径
DEST_PATH="/opt/conda/envs"

# 读取服务器列表并使用 xargs 并行处理
cat $SERVERS_FILE | xargs -P 10 -I {} bash -c "{
    echo 'Processing {}...'
    rsync -azq --progress $TAR_FILE {}:${DEST_PATH}/ > /dev/null 2>&1 && \
    ssh -q {} 'cd $DEST_PATH && tar -zxvf $TAR_FILE && rm -f $TAR_FILE' >/dev/null 2>&1
    if [ $? -eq 0 ]; then
        echo '{}: Deployment successful.'
    else
        echo '{}: Error occurred during deployment.'
    fi
}"

3. 批量复制docker镜像并load到GPU服务器内存中

#!/bin/bash

# 服务器地址列表文件路径
SERVERS_FILE="/apps/sharedstorage/oceanstor-a800/sunxinbao/monitor/gpu_servers.txt"

# 镜像文件路径
IMAGE_PATH="/apps/sharedstorage/oceanstor-a800/sunxinbao/TRANSFER/vllm-0.4.2.tar"

# 读取服务器列表并存入数组
mapfile -t SERVERS < "$SERVERS_FILE"

# 定义一个函数来加载镜像
load_image() {
  local server=$1
  echo "Starting image load on server: $server"

  # 使用 SSH 执行 docker load 命令
  ssh "$server" "docker load < $IMAGE_PATH" && echo "Image loaded successfully on $server" || echo "Failed to load image on $server"
}

# 循环遍历每个服务器并在后台并行执行 load_image 函数
for server in "${SERVERS[@]}"; do
  load_image "$server" &
done

# 等待所有后台任务完成
wait

echo "All operations completed."