如何修改dockerized Flask代码以利用GPU

Question

我有一个 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()

我的 Docker 文件：

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"]

Docker 文件（初始）：

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

Answer 1

我建议使用基于 nvidia/cuda 的图像。

   FROM nvidia/cuda:11.2.2-cudnn8-runtime-ubuntu20.04
    
    ENV DEBIAN_FRONTEND=noninteractive
    
    # Install Python and other dependencies
    RUN apt-get update && apt-get install -y \
        python3.8 \
        python3-pip \
        && rm -rf /var/lib/apt/lists/*
    
    # Set Python 3.8 as the default
    RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 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 支持下运行 docker :

docker run --gpus all <image_name>

如果您想访问所有可用的 GPU。