python - 定期检查 gpu 内存并在空闲时运行脚本

Question

我的系统中有 4 个 GPU (Nvidia)。我想定期检查特定 GPU 是否空闲(例如，空闲内存是否超过 10GB)，如果空闲，我想运行一个 python 脚本。

我想我可以使用 nvidia-smi 检查给定 gpu 有多少空闲内存。我有一个想法，但我无法完全完成剧本。有人可以帮助我吗？

到目前为止我写的是:

check.sh

id=$1
free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id)
echo $free_mem # this prints out: memory.free [MiB] 1954 MiB
while [ $free_mem -lt 10000 ]
    free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id)
    sleep 5

CUDA_VISIBLE_DEVICES=$id python run_python_file.py

我相信代码应该类似于上面的代码片段，但是我找不到详细信息。

Answer 1

使用grep -Eo [0-9]+匹配一个或多个连续数字:

id=$1
free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo [0-9]+)

while [ $free_mem -lt 10000 ]; do
    free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo [0-9]+)
    sleep 5
done

CUDA_VISIBLE_DEVICES=$id python run_python_file.py

grep 的替代方法是使用 POSIX regex : sed 's/[^[:digit:]]*\([[:digit:]]\+\).*/\1/'.