谁在用我的显卡?自动抢占服务器空闲的GPU

查看当前显卡进程使用者,以及在显卡空闲时自动抢占运行的脚本

实验室的显卡总是被别人用,写了两个脚本,放在这里备用

第一个脚本,根据pid查找进程的启动用户、运行目录、启动命令、已运行时间等信息,把那些不认识的人不知道跑的什么玩意还长时间占着显卡的进程kill掉

#!/bin/bash

# 检查是否提供了 PID
if [ -z "$1" ]; then
  echo "Usage: $0 <pid>"
  exit 1
fi

PID=$1

# 检查进程是否存在
if [ ! -d /proc/$PID ]; then
  echo "Process with PID $PID does not exist."
  exit 1
fi

# 获取工作目录
WORK_DIR=$(readlink -f /proc/$PID/cwd)
# 获取命令行
CMD=$(cat /proc/$PID/cmdline | tr '\0' ' ')
# 获取用户
USER=$(ps -o user= -p $PID)
# 获取运行时长
UPTIME=$(ps -o etime= -p $PID)

# 输出信息
echo "PID: $PID"
echo "User: $USER"
echo "Working Directory: $WORK_DIR"
echo "Command: $CMD"
echo "Uptime: $UPTIME"

第二个脚本,设置一个显存阈值,当显存充足的时候运行我们的脚本,显存不足的时候阻塞等待

#!/bin/bash

# 设置显存占用的阈值 (单位为MiB)
THRESHOLD=$1

# 获取显卡显存占用情况的函数
get_gpu_memory_usage() {
    nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '{sum+=$1} END {print sum}'
}
counter=0
# 检查显卡显存占用是否低于阈值
while true; do
    USED_MEMORY=$(get_gpu_memory_usage)
    
    if (( USED_MEMORY < THRESHOLD )); then
        echo "显存 $USED_MEMORY 占用低于阈值 $THRESHOLD MiB,可以继续执行脚本。"
        break
    else
        if (( counter % 100 == 0 )); then
            echo "显存占用 $USED_MEMORY 高于 $THRESHOLD MiB,等待... $(date +"%m-%d %H:%M:%S")"
        fi
        sleep 2  # 等待2秒后再次检查
    fi
    ((counter++))
done

# 在此处添加显卡显存占用低于阈值后的其他操作
echo "显卡可用,执行后续操作... $(date +"%m-%d %H:%M:%S")"

使用方法:插入我们要执行的命令之前bash wait.sh 20000,这样只有在显存占用低于20000MB的时候才会向下执行我们的代码,否则一直等待,每隔2秒检查一次显卡占用