Python利用队列Queue实现多进程Process间通信

背景:

我在利用大模型进行推理的时候,一个机器上面有8张GPU卡,如何充分的利用这几张卡呢?最开始想到的是利用Python多线程方案,可代码实现之后发现虽然我在环境变量里面设置可以使用多张GPU卡,可程序依然只使用了第一张卡,于是就想利用多进程来实现

代码:

from multiprocessing import Queue
import time
import multiprocessing as mp
#这里是定义一个函数,将一个目录下面的文件添加到队列里面,该函数由主进程调用
def put_data(queue, root_dir):    
    message_dir = os.path.join(root_dir, "message");
    file_count = 0
    file_names =[]
    for item in os.listdir(message_dir):
        sub_path = os.path.join(message_dir, item)
        if os.path.isdir(sub_path):
            for subItem in os.listdir(sub_path):
                sub_item_path = os.path.join(sub_path, subItem)
                if os.path.isfile(sub_item_path):
                    file_names.append(sub_item_path)
                    file_count = file_count + 1
    #这里根据文件的创建时间进行排序
    file_names = sorted(file_names, key=lambda x: os.path.getctime(x))
    #取文件创建时间最新的一百个文件
    file_names = file_names[-100:]
    #文件添加到队列里
    for file_name in file_names:
        queue.put(file_name)
    return file_count
#这里定义一个函数,这个函数是后面创建进程将会执行的方法
#gpu_id:该进程将利用哪一块GPU,
def worker(gpu_id, queue, process_num):
    print(f"此时正在执行的GPU num:{gpu_id} --- process_num:{process_num}")
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    #通过设置该环境变量来确保GPU对该进程可见
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
    while True:
        try:
            #从队列里面获取文件名称
            file_name = queue.get();
            #通过重命名文件来解决队列里面可能存在的重复文件名问题,
            #解决主进程在添加文件时,添加了重复的文件
            new_filename = file_name + ".bak"
            try:
                os.rename(file_name, new_filename)
                print(f"此时正在执行的GPU num:{gpu_id} --- process_num:{process_num},文件已成功重命名为:{new_filename}")
            except:
                print(f"此时正在执行的GPU num:{gpu_id} --- process_num:{process_num},文件重命名失败:{new_filename}")
        except Exception as outExceiption:
            print('Error:', outExceiption)

processes = []
queue = Queue()
for i in range(5, 9):
    for process_num in range(1):
        #创建进程,并传递对应的执行函数,和函数的参数
        process = mp.Process(target=worker, args=(str(i),queue, process_num))
        processes.append(process)
        #启动进程
        process.start()
print(f"process count : {len(processes)}")

root_dir = "/data/test"

while True:
    if queue.qsize() < 20:
        file_count = put_data(queue, root_dir)
        print(f"main process find file size : {file_count}")
        if file_count == 0:
            time.sleep(5)
    else:
        time.sleep(5)

结论:

利用multiprocessing的Process可以非常方便的实现多进程架构,另外使用它的Queue也可以非常便捷的进行进程间通信,这样每个进程指定特定的GPU卡,就实现了GPU卡的利用。

你可能感兴趣的:(Python,进程,队列,multiprocessing,GPU)