TaskManager

项目背景

跑深度学习模型的时候需要用到GPU, 这时候通常会有面临两种情况:

1. 组里有自己的服务器, 但是服务器是组里的同学公用的, 这导致服务器的GPU资源比较紧张, 跑实验的时候需要排队才可以抢到卡
2. 在GPU集群上跑实验的时候需要用slurm提交任务, 但是在集群上跑有种种不方便。
   1. 集群上提交一次任务之后失败了只能重新提交, 但重新提交后又要重新排队, 但slurm允许一个任务最多运行14天, 这导致时间利用率非常低
   2. 同一个任务只能跑一个程序, 但是集群上是A100显卡, 单卡80G显存单个任务根本用不完, 这导致显存资源利用率低

本项目主要用于解决上述的问题, 对于问题1, 我们只需要让程序可以自动监控GPU然后抢占显存即可, 对于问题2, 只需要提纲提交/杀死任务的功能, 然后提交任务时提交一个运行这个程序的脚本, 即可实现对分配显卡的手动控制

功能简介

这个项目的核心功能是对GPU进行控制和管理, 提供的功能包括:

1. 自动抢占显存
2. 对抢占的显存进行分配
3. 提交任务/杀死任务

快速开始

0. 安装mysql
   • mysql只是在管理任务时才会使用, 如果只想管理显存的话可以不安装mysql, 但是3里面会报错, 但是不影响显存管理功能
1. 创建conda环境

conda create -n taskmanager python=3.10
conda activate taskmanager
pip install -r requirements.txt
pip install -e .
pip install numpy pymysql

2. 修改test/test.py里面关于mysql的配置信息
3. 运行

python test/test.py --type=server

4. 启动一个新的终端, 并运行

python test/test.py --type=client

此时便可以管理本地的GPU了

因为核心通信功能是使用ZMQ实现的, 它原生支持基于tcp的连接, 因此它也可以管理多服务器上的显卡