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卡的利用。