awk筛nvidia-smi的PID

最新推荐文章于 2024-09-19 10:00:38 发布

原创最新推荐文章于 2024-09-19 10:00:38 发布 · 1.3k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#nvidia-smi #awk #shell #kill #xargs

环境专栏收录该内容

97 篇文章

订阅专栏

本文介绍如何通过nvidia-smi工具获取正在运行的GPU程序的PID，并使用这些PID来终止占用GPU资源的所有程序。此外，还提供了一个脚本示例，用于循环检查并杀死GPU上的进程。

部署运行你感兴趣的模型镜像

想根据 nvidia-smi 的输出 kill 掉所有 gpu 程序。其输出形如：

Wed Jan 12 11:10:11 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.44       Driver Version: 440.44       CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 108...  Off  | 00000000:01:00.0 Off |                  N/A |
|  0%   50C    P8    10W / 250W |   7630MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      2763      C   python                                      1905MiB |
|    0      2776      C   python                                      1905MiB |
|    0      2779      C   python                                      1905MiB |
|    0      2782      C   python                                      1905MiB |
+-----------------------------------------------------------------------------+

或者：

Wed Jan 12 11:51:54 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.80.02    Driver Version: 450.80.02    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-DGXS...  On   | 00000000:07:00.0  On |                    0 |
| N/A   39C    P0    50W / 300W |  14448MiB / 32505MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      7989      C   python                           3611MiB |
|    0   N/A  N/A      7992      C   python                           3611MiB |
|    0   N/A  N/A      7995      C   python                           3611MiB |
|    0   N/A  N/A      8000      C   python                           3611MiB |
+-----------------------------------------------------------------------------+

sieve PID

想获得其中的 PID 一列，可以用 awk：

nvidia-smi | awk '$3 ~/[0-9]+/{if((NR>15)) {print $3}}'

其中：

第一个 $3 似乎是指定针对第 3 列操作（上面第二种输出则换成 $5），其后 ~/<正则>/ 启用正则表达式匹配，参考 [3]
后面就是匹配成功时的操作。NR>15 指 15 行之后，因为从上面可以看出 PID 至少从第 16 行开始（第二种输出是 19，换成 NR>18）；print $3 就是打印第 3 列，即 PID。

如果当时无用 gpu 的是实验在跑，就会是空。

kill

筛完之后，用 xargs 遍历传参，逐一 kill（可能要加 sudo）：

nvidia-smi | awk '$3 ~/[0-9]+/ {if((NR>15)) {print $3}}' | xargs kill -9

如果当初跑程序的时候用了循环和 & 挂后台，kill 一次可能停不完，可以挂个死循环等一下，kill 完为止：

#!/bin/bash

while true; do
    pid_list=$(nvidia-smi | awk '$3 ~/[0-9]+/{if((NR>15)) {print $3}}')
    # echo ${pid_list[@]}
    if [ ${#pid_list[@]} -eq 0 ]; then
        break
    else
        for pid in ${pid_list[@]}; do
            # echo $pid
            kill -9 $pid
        done
    fi
    # 等个 20 秒，看有没有新程序启动
    sleep 20s
done