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

队列Queue实现多进程Process间通信"/>

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 = 0file_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_idwhile 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卡的利用。