服务器等待显卡可用炼丹

服务器脚本——等待显卡可用炼丹

subprocess模块

subprocess模块概述

subprocess模块是Python标准库的一部分，它为创建和管理子进程提供了强大的接口。在Python 2.4中引入以取代旧的模块（如os.system和os.spawn*等），并提供了更复杂的进程管理功能，比如超时控制、I/O管道通信等

subprocess.check_output

1
result = subprocess.check_output(['nvidia-smi', '--query-gpu=utilization.gpu', '--format=csv,noheader,nounits'])
  1. subprocess.check_output:
    • 这是 subprocess 模块中的一个函数，用于执行一个 shell 命令并捕获其输出。
    • 它会等待命令执行完毕，并返回命令的标准输出作为字节串
  2. 命令参数:
    • nvidia-smi: 这是一个用于查询 NVIDIA GPU 信息的工具
    • --query-gpu=utilization.gpu: 这个选项告诉 nvidia-smi 只查询 GPU 的利用率。这里的 utilization.gpu 指的是 GPU 的利用率，即 GPU 的活跃程度
    • --format=csv,noheader,nounits: 这个选项设置了输出格式。csv 表示输出为逗号分隔的值格式，noheader 表示输出中不包含表头，nounits 表示输出中不包含单位信息
  3. result 变量:
    • result 是一个字节串，包含从 nvidia-smi 命令获取的 GPU 使用率信息

subprocess.check_output

1
subprocess.call(["python", "train.py"])
  • subprocess.call(): 这个函数用来运行一个命令并等待它完成。它返回命令的退出码（如果命令正常结束，通常为0；如果有错误，会是非零值）。
  • ["python", "train.py"]: 这是一个列表，其中每个元素代表命令行中的一个参数。在这个例子中，“python”是你想要运行的程序，而“train.py”是你希望传递给该程序的文件名。

当你执行 subprocess.call(["python", "train.py"]) 时，你的Python脚本会启动一个新的Python解释器进程来运行 train.py 脚本。当前的Python脚本会暂停，等待 train.py 完成执行。一旦 train.py 执行完毕，控制权会返回到调用 subprocess.call 的地方，然后继续执行后续代码

检查服务器是否可用，可用即训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import os
import subprocess
import time


def check_gpu_memory():
    """
    使用 nvidia-smi 获取 GPU 显存使用情况
    返回 True 表示至少有一块 GPU 的剩余显存大于 23GB，False 表示没有
    同时返回可用 GPU 的 ID
    """
    try:
        # 使用 nvidia-smi 查询 GPU 显存信息
        result = subprocess.check_output(
            ['nvidia-smi', '--query-gpu=index,memory.used,memory.total', '--format=csv,noheader,nounits'])
        # 将结果转换为字符串并按行分割
        memory_info = result.decode('utf-8').strip().split('\n')
        print("GPU Memory Info:", memory_info)
        # 设置显存阈值（单位：GB）
        memory_threshold_gb = 23

        # 初始化最大剩余显存和对应的 GPU ID
        max_free_memory_gb = 0
        gpu_id_with_max_free_memory = None

        # 检查 GPU 显存使用情况
        for info in memory_info:
            index, used, total = info.split(',')
            free = int(total) - int(used)
            free_gb = free / 1024  # 将 MB 转换为 GB

            if free_gb >= memory_threshold_gb and free_gb > max_free_memory_gb:
                max_free_memory_gb = free_gb
                gpu_id_with_max_free_memory = index

        # 如果找到了剩余显存大于 23GB 的 GPU
        if max_free_memory_gb > 0:
            return True, gpu_id_with_max_free_memory
        else:
            return False, None
    except Exception as e:
        print(f"Error checking GPU memory usage: {e}")
        return False, None  # 出错时，默认认为没有合适的 GPU


def train_network(gpu_id):
    print(f"Starting training on GPU {gpu_id}...")
    # 这里可以替换为你自己的训练脚本或命令
    # 使用 CUDA_VISIBLE_DEVICES 设置要使用的 GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
    subprocess.call(["python", "train.py"])


def main():
    max_wait_time_seconds = 10 * 60 * 60  # 10 hours in seconds
    start_time = time.time()  # 获取当前时间作为开始时间

    while True:
        can_train, gpu_id = check_gpu_memory()
        if can_train:
            train_network(gpu_id)
            break  # 训练完成后退出循环
        else:
            print("No GPU with sufficient memory available. Waiting...")
            time.sleep(60)  # 每隔一分钟检查一次

        # 检查是否超过了最大等待时间
        elapsed_time = time.time() - start_time
        if elapsed_time > max_wait_time_seconds:
            print("Maximum wait time reached. Exiting.")
            break


if __name__ == "__main__":
    main()