gunicorn给worker分配GPU

在gunicorn的config文件(myconfig.py)中添加如下配置代码：

import os
import time

try:
    import pynvml 
    pynvml.nvmlInit()
    gpuDeviceCount = pynvml.nvmlDeviceGetCount()
except:
    gpuDeviceCount = 1
    
gpuDevicePool = []

def pre_fork(server, worker):
    try:
        gid = gpuDevicePool.pop(0)
    except:
        gid = (worker.age - 1) % gpuDeviceCount
    worker.gid = gid

def post_fork(server,worker):
    time.sleep(worker.age % server.cfg.workers)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(worker.gid)
    server.log.info(f'worker(age:{worker.age}, pid:{worker.pid}, cuda:{worker.gid})')
    
def child_exit(server, worker):
    gpuDevicePool.append(worker.gid)

启动服务：

gunicorn -c myconfig.py server:app -w 16