赞
踩
我在利用大模型进行推理的时候,一个机器上面有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卡的利用。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。