Docker部署GPU+Torch+Gunicorn+Flask应用遇CUDA报错，求可用示例

首先，你遇到的Cannot re-initialize CUDA in forked subprocess错误是因为Gunicorn默认用fork方式创建子进程，而CUDA上下文没办法在fork出的子进程里重新初始化。结合你的需求，我整理了一套完整的可运行配置，包含Dockerfile优化、代码调整和启动命令规范。

一、修正Dockerfile（适配CUDA 10.2 + Torch环境） #

你的原Dockerfile存在两个核心问题：一是PyTorch安装未指定CUDA版本，默认会安装CPU版本；二是依赖安装步骤可以更精简。以下是优化后的版本：

FROM nvidia/cuda:10.2-base-ubuntu18.04

# 设置环境变量避免交互提示
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# 安装基础依赖与Python3.7
RUN apt-get update && apt-get install -y --no-install-recommends \
    curl \
    ca-certificates \
    sudo \
    git \
    bzip2 \
    libx11-6 \
    python3.7 \
    python3.7-dev \
    python3.7-distutils \
    && rm -rf /var/lib/apt/lists/*

# 配置Python默认版本为3.7
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.7 1
RUN update-alternatives --set python /usr/bin/python3.7

# 升级pip到最新版本
RUN curl -s https://bootstrap.pypa.io/get-pip.py -o get-pip.py && \
    python get-pip.py --force-reinstall && \
    rm get-pip.py

# 创建工作目录并切换
WORKDIR /usr/src/app

# 安装Python依赖
COPY requirements.txt .
RUN pip --no-cache-dir install -r requirements.txt

# 安装适配CUDA10.2的PyTorch版本
# 注：Torch1.4.0官方仅支持到CUDA10.1，若必须用1.4.0，建议将基础镜像换成nvidia/cuda:10.1-base-ubuntu18.04，再用下面的命令：
# RUN pip install torch==1.4.0 torchvision==0.5.0 -f https://download.pytorch.org/whl/cu101/torch_stable.html
# 若可以升级Torch，推荐用兼容CUDA10.2的1.6.0版本：
RUN pip install torch==1.6.0 torchvision==0.7.0 -f https://download.pytorch.org/whl/cu102/torch_stable.html

# 复制应用代码到容器内
COPY . ./

# 用Gunicorn直接启动应用，配置预加载参数
CMD ["gunicorn", "--bind", "0.0.0.0:8888", "--workers", "1", "--timeout", "300", "--preload", "new_main:app"]

二、调整Flask应用代码 #

1. 修改模型加载逻辑 #

把模型加载移到应用启动阶段（而非第一次请求时），配合Gunicorn的--preload参数让主进程提前加载模型，子进程共享已初始化的CUDA上下文：

# new_main.py
import argparse
from gunicorn.app.base import StandaloneApplication
from flask import Flask, request

app = Flask(__name__)
app.worker = None

def init():
    # 启动时直接初始化模型，避免请求时重复加载
    global app
    app.worker = worker()
    app.worker.load_models()

# 启动应用时执行初始化
init()

@app.route("/api/work", methods=["POST"])
def work():
    try:
        body = request.get_json()
        # 直接使用已加载好的模型处理请求
        ...
    except Exception as e:
        return {"error": str(e)}, 500

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", action='store_true')
    parser.add_argument("--workers", type=int, default=1)
    args = parser.parse_args()
    if args.test:
        # 测试模式直接启动Flask调试服务
        app.run(host='0.0.0.0', port=8888, debug=True)
    else:
        options = {
            'bind': '0.0.0.0:8888',
            'workers': args.workers,
            'timeout': 300,
            'preload_app': True  # 关键参数：主进程预加载模型后再fork子进程
        }
        StandaloneApplication(app, options).run()

2. 多GPU场景适配（可选） #

如果需要使用多GPU，建议每个worker绑定单独GPU，且不要使用--preload，而是在每个worker进程中用spawn方式初始化模型：

# 在模型初始化代码中添加
import torch.multiprocessing as mp
mp.set_start_method('spawn', force=True)

# 绑定指定GPU
device = torch.device(f"cuda:{worker_id}" if torch.cuda.is_available() else "cpu")
model = model.to(device)

三、正确的Docker启动命令 #

确保你的Docker环境已安装nvidia-docker2以支持GPU，启动命令如下：

sudo docker run -p 8889:8888 \
  -e MODELSLOCATION=/mnt/models \
  --gpus all \
  -v $MODELSLOCATION:/mnt/models \
  cc14ffc68256

关键注意事项 #

• --preload_app=True：让Gunicorn主进程先加载应用和模型，再创建子进程，所有子进程共享已初始化的CUDA上下文，避免重复初始化报错。
• 匹配CUDA与Torch版本：必须安装和基础镜像CUDA版本兼容的Torch，否则会出现CUDA兼容性问题。
• 避免请求时延迟加载模型：延迟加载会导致每个worker进程尝试初始化CUDA，触发fork子进程的CUDA错误。

内容的提问来源于stack exchange，提问作者artona