如何修改Docker化Flask代码以利用GPU运行多进程任务

我有一个Docker化的Flask应用，希望以多进程模式运行以下代码，但当前代码仅占用CPU核心且容器出现崩溃情况。请问应如何操作才能确保应用利用GPU而非CPU？代码或容器需要做出哪些修改？

当前多进程代码 #

processes = []
for i in range(len(session['optitower_files'])):
    try:
        print("_________________________________Started Processess__________________________________________")
        p = Process(target= generate_optimizer_formula, kwargs={"data":optitower_data_frames[f'data_{i}'],"save_file_paths":save_file_paths[i], "return_data":return_data[f'data_{i}'],"brd_folder_path":brd_folder_path,"user_session_path":user_session_path,"meta_data":meta_data,"legacy":legacy})
        processes.append(p)
        p.start()
        # data,data2     = return_data

    except Exception as e:
        print(e)
        pass
        # data_dict               = read_from_json(loc=save_file_paths["Optimizer Status"])
        # data_dict["Optimizer"]  = "Error Occured; Optimizer Stopped"
        # write_to_json(data_dict,save_file_paths["Optimizer Status"])  
for p in processes:
    p.join()

当前Dockerfile #

FROM company_flask_app_v2_backup
RUN bash
ENV DEBIAN_FRONTEND=noninteractive

RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
WORKDIR /company

COPY ./requirements/requirements.txt .
COPY ./requirements/requirements.py .

RUN pip3 install -r requirements.txt
RUN python3 requirements.py

CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "app_dev.py"]

初始基础镜像Dockerfile（company_flask_app_v2_backup来源） #

FROM ubuntu_image_v2

WORKDIR /company

ENV DEBIAN_FRONTEND=noninteractive

COPY ./requirements/requirements.txt .
COPY ./requirements/requirements.py .
RUN apt-get update

RUN apt-get install -y \
    software-properties-common

RUN apt-get update &&  add-apt-repository universe
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt-get update &&  apt-get install -y \
    python3.11 \
    python3-pip

RUN apt install -y wkhtmltopdf
RUN pip3 install -r requirements.txt

RUN apt-get update &&  apt-get install latexmk -y --fix-missing
RUN apt-get install texlive-latex-extra -y
RUN apt-get install texlive-fonts-recommended -y

RUN apt-get update && \
    apt-get install -y curl gnupg unixodbc

RUN curl https://packages.microsoft.com/keys/microsoft.asc | apt-key add - && \
    curl https://packages.microsoft.com/config/ubuntu/$(lsb_release -rs)/prod.list > /etc/apt/sources.list.d/mssql-release.list

RUN apt-get update && \
    ACCEPT_EULA=Y apt-get install -y msodbcsql18
    
RUN apt-get clean && \
    rm -rf /var/lib/apt/lists/*


RUN apt-get update &&  apt-get install libreoffice -y --fix-missing

一、容器层面：适配GPU运行环境 #

1. 替换基础镜像为GPU兼容版本 #

当前基础镜像为普通Ubuntu，需切换到NVIDIA官方提供的CUDA镜像，确保与宿主机CUDA版本匹配。修改初始Dockerfile的基础镜像：

# 示例：选择CUDA 12.1 + Ubuntu 22.04的运行时镜像，根据宿主机CUDA版本调整
FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04

2. 安装GPU依赖与适配Python库 #

• 在requirements.txt中添加GPU版本的计算库，比如：

  torch>=2.0.0+cu121
  cupy-cuda12x>=12.0.0
  cudf>=23.10.0
  # 根据业务依赖替换为对应GPU版本的库
  

• 初始Dockerfile中无需额外安装CUDA驱动（镜像已包含），只需确保Python环境正常即可。

3. 修改应用Dockerfile #

添加GPU识别环境变量，确保容器能感知GPU资源：

FROM company_flask_app_v2_gpu_backup # 基于新GPU镜像构建的基础镜像
ENV DEBIAN_FRONTEND=noninteractive
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
WORKDIR /company

COPY ./requirements/requirements.txt .
COPY ./requirements/requirements.py .

RUN pip3 install -r requirements.txt
RUN python3 requirements.py

CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "app_dev.py"]

二、代码层面：迁移计算逻辑到GPU #

1. 改造核心计算函数generate_optimizer_formula #

将函数中的CPU计算逻辑替换为GPU加速实现，以NumPy转cuPy为例：

import cupy as cp
# 如果用Pandas处理数据，替换为cudf：import cudf as pd

def generate_optimizer_formula(data, save_file_paths, return_data, brd_folder_path, user_session_path, meta_data, legacy):
    # 将CPU数据转移到GPU内存
    gpu_data = cp.array(data)
    
    # 替换所有NumPy计算为cuPy操作
    optimized_result = cp.matmul(gpu_data, gpu_data.T) # 示例GPU计算逻辑
    
    # 若需保存到磁盘，将结果转回CPU内存
    cpu_result = cp.asnumpy(optimized_result)
    
    # 后续保存、返回逻辑保持不变
    save_to_file(cpu_result, save_file_paths)

2. 优化多进程管理，避免容器崩溃 #

• 限制进程数量：根据GPU核心数分批启动进程，避免一次性耗尽资源：

  import cupy as cp
  
  max_processes = cp.cuda.runtime.getDeviceCount() * 2 # 每个GPU最多跑2个进程
  total_files = len(session['optitower_files'])
  
  # 分批启动进程
  for batch_start in range(0, total_files, max_processes):
      batch_end = min(batch_start + max_processes, total_files)
      batch_processes = []
      
      for i in range(batch_start, batch_end):
          try:
              p = Process(
                  target=generate_optimizer_formula,
                  kwargs={"data":optitower_data_frames[f'data_{i}'],
                          "save_file_paths":save_file_paths[i],
                          "return_data":return_data[f'data_{i}'],
                          "brd_folder_path":brd_folder_path,
                          "user_session_path":user_session_path,
                          "meta_data":meta_data,
                          "legacy":legacy}
              )
              batch_processes.append(p)
              p.start()
          except Exception as e:
              print(e)
      
      # 等待当前批次进程完成再启动下一批
      for p in batch_processes:
          p.join()
  

• 进程绑定GPU：多GPU环境下为每个进程分配独立GPU，避免资源冲突：

  import os
  
  def worker(gpu_id, data, save_file_paths, ...):
      os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
      # 执行GPU计算逻辑
      generate_optimizer_formula(data, save_file_paths, ...)
  
  # 启动进程时分配GPU
  p = Process(target=worker, args=(i % max_processes, optitower_data_frames[f'data_{i}'], save_file_paths[i], ...))
  

三、容器启动配置 #

必须使用NVIDIA容器工具链启动容器，确保GPU被正确挂载：

# 挂载所有可用GPU
docker run --gpus all -p 5000:5000 your-gpu-flask-app-image

# 或指定具体GPU
docker run --gpus "device=0,1" -p 5000:5000 your-gpu-flask-app-image

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