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()