GUnicorn+CUDA preload模式fork子进程CUDA重初始化报错求解

问题背景 #

基于PyTorch、Gunicorn、Flask搭建调用CUDA的推理服务时，为降低多worker带来的额外资源开销，启用Gunicorn的--preload选项尝试让多个worker进程共享加载的模型，但触发CUDA运行时异常。

最小复现代码 #

服务代码run_server.py内容如下：

from flask import Flask, request
import torch

app = Flask('dummy')

model = torch.rand(500)
model = model.to('cuda:0')


@app.route('/', methods=['POST'])
def f():
    data = request.get_json()
    x = torch.rand((data['number'], 500))
    x = x.to('cuda:0')
    res = x * model
    return {
        "result": res.sum().item()
    }

报错现象 #

使用如下命令可正常启动服务进程：

CUDA_VISIBLE_DEVICES=1 gunicorn -w 3 -b $HOST_IP:8080 --preload run_server:app

发起首次POST请求时worker进程抛出错误，请求命令：

curl -X POST -d '{"number": 1}' http://$HOST_IP:8080/

错误栈信息：

[2022-06-28 09:42:00,378] ERROR in app: Exception on / [POST]
Traceback (most recent call last):
  File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 2447, in wsgi_app
    response = self.full_dispatch_request()
  File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1952, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1821, in handle_user_exception
    reraise(exc_type, exc_value, tb)
  File "/home/user/.local/lib/python3.6/site-packages/flask/_compat.py", line 39, in reraise
    raise value
  File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1950, in full_dispatch_request
    rv = self.dispatch_request()
  File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1936, in dispatch_request
    return self.view_functions[rule.endpoint](**req.view_args)
  File "/home/user/project/run_server.py", line 14, in f
    x = x.to('cuda:0')
  File "/home/user/.local/lib/python3.6/site-packages/torch/cuda/__init__.py", line 195, in _lazy_init
    "Cannot re-initialize CUDA in forked subprocess. " + msg)
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

根因分析 #

• Gunicorn启用--preload选项时，会先在主进程加载全部应用代码（包括将模型迁移到CUDA的逻辑、完成CUDA上下文初始化），再通过fork系统调用创建worker进程
• CUDA上下文不支持在fork出的子进程中复用，子进程第一次执行CUDA张量操作时会尝试重新初始化CUDA，直接触发运行时错误
• 手动配置torch.multiprocessing.set_start_method('spawn')或multiprocessing.set_start_method('spawn')无效，因为Gunicorn的--preload模式固定使用fork创建子进程，不受Python多进程启动方法配置影响
• 若关闭--preload选项，每个worker会独立加载模型到显存，产生多份模型副本，显存、内存开销过高，不符合资源优化预期

可行解决方案 #

所有方案均不需要在每个worker中重复执行完整的模型加载逻辑，可根据资源情况和并发需求选择：

方案1：延迟CUDA初始化（改动最小，CPU内存零冗余） #

核心逻辑是避免主进程触发任何CUDA操作，利用fork的写时复制机制让所有worker共享CPU侧的模型内存，worker启动后再独立初始化CUDA：

1. 修改服务代码，主进程加载模型时仅保留CPU上的副本，不执行.to('cuda')操作，保证主进程全程不初始化CUDA上下文
2. 新增Gunicorn配置文件，通过post_fork钩子在每个worker进程fork完成后，再执行CUDA初始化、将模型迁移到GPU

修改后的run_server.py：

from flask import Flask, request
import torch

app = Flask('dummy')
# 主进程仅在CPU加载模型，不触发CUDA初始化
model = torch.rand(500)

@app.route('/', methods=['POST'])
def f():
    data = request.get_json()
    x = torch.rand((data['number'], 500)).to('cuda:0')
    res = x * model
    return {"result": res.sum().item()}

新增Gunicorn配置文件gunicorn_conf.py：

def post_fork(server, worker):
    # worker进程fork完成后初始化CUDA，将模型迁移到GPU
    import run_server
    run_server.model = run_server.model.to('cuda:0')

启动命令：

CUDA_VISIBLE_DEVICES=1 gunicorn -w 3 -b $HOST_IP:8080 --preload -c gunicorn_conf.py run_server:app

该方案下CPU侧的模型内存会被所有worker共享，无额外CPU内存开销；显存侧每个worker持有独立的CUDA上下文，会保留一份模型副本，适合显存充足、希望最小化代码改动的场景。

方案2：单worker多线程模式（全资源单副本，配置最简单） #

如果需要显存、内存都只保留一份模型副本，可改用Gunicorn的线程worker模式：

• 仅启动1个worker进程，进程内启动多个线程处理请求
• 模型只在该进程内加载一次到CUDA，完全没有多副本开销
• PyTorch执行CUDA推理时会主动释放GIL，多线程不会造成明显的性能阻塞，线程数可根据并发需求设置为4-16，足以应对绝大多数推理场景

启动命令：

CUDA_VISIBLE_DEVICES=1 gunicorn -w 1 --threads 8 -b $HOST_IP:8080 run_server:app

该模式不需要开启--preload，也不会触发fork相关的CUDA错误，配置成本最低。

方案3：独立推理进程架构（全资源单副本，支持高并发） #

如果并发要求高、单进程多线程无法满足需求，可将Web服务逻辑和推理逻辑解耦：

• 单独启动1个常驻推理进程，该进程唯一持有CUDA上下文和模型副本
• Gunicorn启动多worker（不需要preload，worker本身不加载模型），仅负责HTTP请求解析、参数校验等轻量逻辑
• worker和推理进程之间通过本地Unix套接字、共享内存等高效IPC方式传输输入数据和推理结果

该架构下模型始终只有一份副本，不会触发fork相关的CUDA错误，同时支持多worker并行处理请求，稳定性更高，适合生产环境高并发场景，仅需要额外开发简单的IPC通信逻辑。

内容的提问来源于stack exchange，提问作者Green 绿色