前言

同样的,这篇博客也源自于我在做组内2030项目所产生的结果。当时,5个硕士生需要进行类似的微调工作,偶尔还会有博士生使用服务器上的GPU,但服务器上仅有8块GPU。

因此,如何快速抢占到 \(n\) 块GPU,从而高效完成手里的工作,便是一个很重要的问题啦~^ _ ^

问题

我首先在网上看了下现有的抢GPU的脚本,但发现简单的脚本要么只能抢1块GPU,要么是一个复杂项目操作起来较麻烦。

于是便萌生了自己写个Python脚本,这样以后凡是涉及到需要抢GPU的场景,我都可以通过运行该脚本抢占到 \(n\) 块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)

抢占GPU的脚本的更多相关文章

  1. ubuntu12.04通过Ganglia利用NVML模块进行GPU监控

    1.安装Ganglia,这里安装的是3.1*版本,因为监控GPU的模块只支持3.1*版本系列的 apt-get install ganglia* 2.下载并安装PyNVML和NVML模块,下载地址ht ...

  2. 谁在使用GPU?

    nvidia-smi命令可以查看GPU使用情况,但是只能看到占用每个GPU的进程ID.根据进程ID可以得到进程详情,进程详情中包括用户ID,根据用户ID可以获取用户名称,从而知道哪个用户在使用GPU. ...

  3. keepalived 配置文件参数详解

    global_defs 全局配置vrrpd 1. vrrp_script添加一个周期性执行的脚本.脚本的退出状态码会被调用它的所有的VRRP Instance记录. 2. vrrp_sync_grou ...

  4. Ogre 整体框架入门

    ogre 是面向对象的3d图形引擎. root 是引擎的一个界面类,包含很多快捷的调用其他类的接口. 在ogre中,广泛的使用了单件模式,同时最大的保证了你不需要自己管理资源,除了是你自己new的对象 ...

  5. Keepalived+nginx高可用

    这种方法会把Keepalived进程结束掉,在教育机构学习到的方法,我个人对这种方法不认可. 参考: https://www.cnblogs.com/gshelldon/p/14504236.html ...

  6. 从 PyTorch DDP 到 Accelerate 到 Trainer,轻松掌握分布式训练

    概述 本教程假定你已经对于 PyToch 训练一个简单模型有一定的基础理解.本教程将展示使用 3 种封装层级不同的方法调用 DDP (DistributedDataParallel) 进程,在多个 G ...

  7. [转]各种移动GPU压缩纹理的使用方法

    介绍了各种移动设备所使用的GPU,以及各个GPU所支持的压缩纹理的格式和使用方法.1. 移动GPU大全 目前移动市场的GPU主要有四大厂商系列:1)Imagination Technologies的P ...

  8. Caffe学习笔记2--Ubuntu 14.04 64bit 安装Caffe(GPU版本)

    0.检查配置 1. VMWare上运行的Ubuntu,并不能支持真实的GPU(除了特定版本的VMWare和特定的GPU,要求条件严格,所以我在VMWare上搭建好了Caffe环境后,又重新在Windo ...

  9. GPU深度发掘(一)::GPGPU数学基础教程

    作者:Dominik Göddeke                 译者:华文广 Contents 介绍 准备条件 硬件设备要求 软件设备要求 两者选择 初始化OpenGL GLUT OpenGL ...

  10. Linux后台执行脚本文件,nohup

    看运维人员执行nohup命令后,把程序放在后台执行,很高大上,就研究了一下,这个命令. nohup命令及其输出文件 nohup命令:如果你正在运行一个进程,而且你觉得在退出帐户时该进程还不会结束,那么 ...

随机推荐

  1. MySQL 数据库中的数据类型

    整数类型 标准 SQL 中支持 INTEGER 和 SMALLINT 这两种类型,MySQL 数据库除了支持这两种类型以外,还扩展支持了 TINYINT.MEDIUMINT 和 BIGINT 整数类型 ...

  2. 多智能体协同控制研究中光学动作捕捉与UWB定位技术比较

    人类在进行任何工作时,总是强调团队合作,teamwork.随着控制科学.计算机科学等多学科的交叉发展与融合,在智能体控制领域,对于单个机器人,无人机,无人车的控制已经不能满足现在领域的技术需求,从而和 ...

  3. Go--下载安装

    下载包地址:https://go.dev/dl/ linux: 下载后上传实例解压 tar -xvf go1.19.8.linux-amd64.tar.gz -C /usr/local/ 创建工作目录 ...

  4. 【有奖体验】AIGC小说创作大赛开启!通义千问X函数计算部署AI助手

    一个 AI 助手到底能做什么? 可以书写小说 可以解析编写代码 可以鼓舞心灵 提供职业建议 还有更多能力需要您自己去探索.接下来我们将花费 5 分钟,基于函数计算X通义千问部署一个 AI 助手,帮你撰 ...

  5. 函数计算 HTTP 触发器支持异步,解放双手搭建 Web 服务

    作者| 阿里云Serverless技术专家 澈尔 当前阿里云函数计算支持两种类型的函数:事件函数和 HTTP 函数.其中 HTTP 函数结合 HTTP 触发器,能够支持用户直接通过 HTTP 请求利用 ...

  6. 这才是你应该了解的Redis数据结构!

    深入了解Redis数据结构 Redis,作为一种高性能的内存数据库,支持多种数据结构,从简单的字符串到复杂的哈希表.在这篇博文中,我们将深入探讨Redis的一些主要数据结构,并通过详细的例子展示它们的 ...

  7. P5731

    https://www.luogu.com.cn/problem/P5731 这道题被标为红题,真实难度应该介于红题和橙题之间,问题在于我高估了它的难度,以为至少有橙题的难度,一般不打表的我毫不犹豫选 ...

  8. 如何使用chatgpt编写代码

    功能列举 回答编程问题 我想让你充当 Stackoverflow 的帖子.我将提出与编程有关的问题,你将回答答案是什么.我希望你只回答给定的答案,在没有足够的细节时写出解释.当我需要用英语告诉你一些事 ...

  9. Chrome 控制台 换行编写js调试代码

    转载请注明出处: 在 chrome 浏览器的console 控制台编写 js 调试或验证代码时,每输一行换行时,就会执行当前行的函数,再重新换行输入时,就会将之前的代码忽略,这种方式就会导致 chro ...

  10. python pip手动安装二进制包

    python中使用pip安装扩展包的时候,有时候会遇到如下类似报错: Running setup.py install for mysqlclient ... error ...(中间报错信息省略) ...