[python脚本]批量杀死某个用户跑在GPU上的所有进程

1. 代码

import argparse
import os
import time

parser = argparse.ArgumentParser()
parser.add_argument('--user', type=str, default='zhangr', help='killed process who run on the server')
args = parser.parse_args()


def get_exec_out(sxcute_str):
    out_list = os.popen(sxcute_str).readlines()
    return out_list


if __name__ == '__main__':
    execute_str = 'nvidia-smi'
    user = args.user.strip()
    dir_keyword = 'cwd ->'
    out_list = get_exec_out(execute_str)
    # print(out_list)
    for oo in out_list:
        if oo.find('python') != -1:
            # split()函数默认可以按空格分割，并且把结果中的空字符串删除掉，留下有用信息
            proc_list = oo.split()
            pid = proc_list[2].strip()
            proc_execute = 'ls -l /proc/' + pid
            # print(proc_execute)
            proc_details = get_exec_out(proc_execute)
            for detail in proc_details:
                if detail.find(user) != -1 and detail.find(dir_keyword) != -1:
                    # print(detail)
                    kill_str = 'kill -9 ' + pid
                    print(kill_str)
                    os.system(kill_str)

2. 运行命令

python kill_me.py --user zhangr