抢占GPU的脚本
时间:2023-08-26
本文章向大家介绍抢占GPU的脚本,主要内容包括其使用实例、应用技巧、基本知识点总结和需要注意事项,具有一定的参考价值,需要的朋友可以参考一下。
我主要利用Python多进程编程,通过占用GPU内存,从而达到占用GPU的目的。关于代码的解释见我的个人博客,以下主要介绍如何使用该脚本。
我的Python版本为3.11,执行命令如下
python grab_gpu.py --n 3 --otime 30 --spath ./train.sh
其中n表示需要占用的GPU个数,otime表示占用时间,spath表示一旦释放GPU后,我们需要执行的脚本。
运行结果如下:
完整代码如下
import os
import subprocess
import time
import argparse
from multiprocessing import Process, Value, Lock, Array
def get_gpu_mem(gpu_id):
gpu_query = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader'])
gpu_memory = [int(x) for x in gpu_query.decode('utf-8').split('\n')[:-1]]
return gpu_memory[gpu_id]
def get_free_gpus()->list:
gpu_query = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader'])
gpu_memory = [int(x) for x in gpu_query.decode('utf-8').split('\n')[:-1]]
free_gpus = [i for i, mem in enumerate(gpu_memory) if mem < 100]
return free_gpus
def occupy_gpu(gpu_id:int, n, occupy_num, ocpy_gpus, lock, a_dim=140000):
with lock:
if get_gpu_mem(gpu_id) < 100 and occupy_num.value < n:
import torch
a = torch.ones((a_dim,a_dim)).cuda(gpu_id)
ocpy_gpus[occupy_num.value]= gpu_id
occupy_num.value += 1
print(f"Occupying GPU {gpu_id}, Total Occupied: {occupy_num.value}")
while True:
time.sleep(10)
def occupy_all_gpus(n:int, occupy_num, ocpy_gpus, interval=10):
print("Launching process to occupy GPU ...")
lock = Lock()
processes = [] #List to store the processes
while occupy_num.value < n:
free_gpus = get_free_gpus()
will_occupy_num = min(n, max(0,len(free_gpus)))
for i in range(will_occupy_num):
if occupy_num.value < n:
p = Process(target=occupy_gpu, args=(free_gpus[i], n, occupy_num, ocpy_gpus, lock))
p.start()
processes.append(p)
time.sleep(interval) # enough time to occupy gpus and update nvidia-smi
return processes, ocpy_gpus
def run_my_program(n, desired_script, processes, ocpy_gpus, occupy_num):
for p in processes:
p.terminate()
ocpy_gpus_list = list(ocpy_gpus[:occupy_num.value])
cuda_visible_devices = ",".join(map(str, ocpy_gpus_list))
os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices
subprocess.run([desired_script, str(n)])
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Arguments for Occupy GPUs")
parser.add_argument(
"--n", type=int, default=2, help="Number of GPUs to occupy"
)
parser.add_argument(
"--otime", type=int, default=10, help="Time of occupying gpu"
)
parser.add_argument(
"--spath", type=str, default='./train.sh', help="the execute script path"
)
args = parser.parse_args()
n = args.n
occupy_time = args.otime
desired_script = args.spath
occupy_num = Value('i', 0) # Shared variable to count occupied GPUs
ocpy_gpus = Array('i', [-1 for _ in range(8)])# Shared array to store occupied gpu
processes,ocpy_gpus = occupy_all_gpus(n, occupy_num, ocpy_gpus)
time.sleep(occupy_time)
run_my_program(n, desired_script, processes, ocpy_gpus, occupy_num)
原文地址:https://www.cnblogs.com/littletreee/p/17658006.html
- hdu 2818 Building Block
- hdu 3074 Zjnu Stadium (带权并查集)
- 线程练习题---简易多人聊天室
- Servlet的尾(yi)巴---filter ( 过滤器 )的小应用
- ijg库的使用的几点注意
- C++获取鼠标位置及全局检测鼠标行为
- C/C++ http协议加载sessionID
- 两个月入门深度学习,全靠动手实践!一位前端小哥的经验分享
- 服务器端Go程序对长短链接的处理及运行参数的保存
- CEF使用的几个注意点
- 排序
- c/c++ 软件集成 安装和可卸载软件
- C/C++ 如何来自动优雅的涮别银家的贴子
- C/C++-----------http协议发送字段,文件,单个和多张图片
- JavaScript 教程
- JavaScript 编辑工具
- JavaScript 与HTML
- JavaScript 与Java
- JavaScript 数据结构
- JavaScript 基本数据类型
- JavaScript 特殊数据类型
- JavaScript 运算符
- JavaScript typeof 运算符
- JavaScript 表达式
- JavaScript 类型转换
- JavaScript 基本语法
- JavaScript 注释
- Javascript 基本处理流程
- Javascript 选择结构
- Javascript if 语句
- Javascript if 语句的嵌套
- Javascript switch 语句
- Javascript 循环结构
- Javascript 循环结构实例
- Javascript 跳转语句
- Javascript 控制语句总结
- Javascript 函数介绍
- Javascript 函数的定义
- Javascript 函数调用
- Javascript 几种特殊的函数
- JavaScript 内置函数简介
- Javascript eval() 函数
- Javascript isFinite() 函数
- Javascript isNaN() 函数
- parseInt() 与 parseFloat()
- escape() 与 unescape()
- Javascript 字符串介绍
- Javascript length属性
- javascript 字符串函数
- Javascript 日期对象简介
- Javascript 日期对象用途
- Date 对象属性和方法
- Javascript 数组是什么
- Javascript 创建数组
- Javascript 数组赋值与取值
- Javascript 数组属性和方法
- Angular rxjs Observable的异步行为
- Angular rxjs里自定义operator的使用
- tensorflow 生成指定大小的赋值0的张量 np.zeros 在TF中对应的语句 生成全0张量
- Angular rxjs fromEvent使用的一个例子
- 在StackBlitz上setup SAP Spartacus
- [898]python获取两个list交集|并集|差集
- [897]使用Maxwell实时同步mysql数据
- Magicodes.IE之导入导出筛选器
- 一文搞懂Flink生成StreamGraph
- 知识全聚集 .Net Core 技术突破 | 如何实现一个模块化方案一
- ESP8266和ROS调试一些问题汇总
- tianbot_mini机器人上手ROS/SLAM/Navigation究竟有多简单???
- 如何知道我们的E2E测试覆盖率?
- tf.add_n函数的用法
- tensorflow 生成指定大小的赋值0的张量 np.zeros 在TF中对应的语句 生成全0张量