asdTools-自动搜索可用GPU

Github
Blog

摘要:自动获取可用 GPU 的工具介绍。

场景

  • 有十个实验要跑,但不想每次跑的时候都确认哪个GPU可用。
  • 又有十个实验要跑,但不想定时上号开实验,想让他有GPU时自动跑。

使用示例

  • 开源地址:asdTools/GPUMonitor.py at main
  • 在指定GPU的地方,调用本模块,并启动程序,实验将会等待直至本模块返回可用
from asdTools.Classes.AI.GPUMonitor import GPUMonitor

def SampleInChinese():
    """
    假设我有8张卡,我想获取4张卡,每张卡至少16G显存,并且我需要字符串格式的结果。
    使用如下代码,将持续搜索可用GPU,每次搜索间隔不超过300秒。
    搜索不到时,等待random(1, 300)秒后,再次搜索。
    如果某次搜索发现 GPU 0, 1, 4, 5, 6, 7 空闲,则会返回"0,1,4,5",并结束搜索。
    """
    gpus_id = [0, 1, 2, 3, 4, 5, 6, 7]
    # 方式1
    monitor = GPUMonitor(capacity_thres=16, output_type="str", time_cooldown=300)
    gpus_available = monitor(gpus_id, gpu_need=4)
    # 方式2
    gpus_available = GPUMonitor()(gpus_id, gpu_need=4, capacity_thres=16, output_type="str", time_cooldown=300)

    """
    我想知道哪些实验正在等待GPU,则使用extra_info,为每次输出增加前缀。
    我需要list类型的结果,则output_type="list"
    2023-05-14 20:14:06: test_No available gpus, sleep 3.99s.
    2023-05-14 20:14:10: test_No available gpus, sleep 2.51s.
    2023-05-14 20:14:13: test_Find available gpus: [0, 1, 4, 5].
    """
    gpus_available = monitor(gpus_id, gpu_need=4, extra_info="test", output_type="list")

if __name__ == "__main__":
    SampleInChinese()

版权声明:
作者:MWHLS
链接:https://mwhls.top/4564.html
来源:无镣之涯
文章版权归作者所有,未经允许请勿转载。

THE END
分享
二维码
打赏
< <上一篇
下一篇>>