Docker容器中llama-cpp-python无GPU支持问题排查与解决

背景 #

尝试在Docker容器内部署集成llama.cpp的LlamaIndex，采用llama_cpp_python==0.2.6版本，按官方仓库指引安装。

初始Dockerfile #

# 使用官方Python 3.11镜像
FROM python:3.11

# 设置容器内工作目录
WORKDIR /app

# 将当前目录内容复制到容器的/app目录下
COPY . /app

# ARG FORCE_CMAKE=1

# ARG CMAKE_ARGS="-DLLAMA_CUBLAS=on"

# 安装项目依赖
RUN FORCE_CMAKE=1 CMAKE_ARGS="-DLLAMA_CUBLAS=on" python -m pip install -r requirements.txt

# 启动服务器的命令
CMD ["python", "./server.py"]

运行命令 #

docker build -t llm_server ./llm
docker run -it -p 2023:2023 --gpus all llm_server

问题描述 #

llama.cpp文档中指定的环境变量在Docker容器内未按预期生效，当前LLM仅使用CPU（BLAS=0），预期应启用GPU加速（BLAS=1）。容器内nvidia-smi输出显示GPU状态正常：

# nvidia-smi
Thu Nov 23 05:48:30 2023       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.29.01              Driver Version: 546.01       CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce GTX 1660 Ti     On  | 00000000:01:00.0  On |                  N/A |
| N/A   48C    P8               4W /  80W |   1257MiB /  6144MiB |      7%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                          
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A        20      G   /Xwayland                                 N/A      |
|    0   N/A  N/A        20      G   /Xwayland                                 N/A      |
|    0   N/A  N/A       392      G   /Xwayland                                 N/A      |
+---------------------------------------------------------------------------------------+
#

尝试的解决方法 #

• 在Dockerfile中尝试多种环境变量配置方式：

  # ARG FORCE_CMAKE=1
  
  # ARG CMAKE_ARGS="-DLLAMA_CUBLAS=on"
  # ENV FORCE_CMAKE=1
  
  # ENV CMAKE_ARGS="-DLLAMA_CUBLAS=on"
  
  # 安装项目依赖
  RUN FORCE_CMAKE=1 CMAKE_ARGS="-DLLAMA_CUBLAS=on" python -m pip install -r requirements.txt
  

• 在运行中的容器内重新安装llama-cpp-python：

  CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python
  

最终有效Dockerfile #

切换到基于nvidia/cuda的基础镜像后解决了问题，Dockerfile内容如下：

FROM nvidia/cuda:11.7.1-devel-ubuntu22.04

# 设置容器内工作目录
WORKDIR /app

# 将当前目录内容复制到容器的/app目录下
COPY . /app

# 安装Python和pip
RUN apt-get update && apt-get install -y python3 python3-pip

# 设置环境变量
ENV CMAKE_ARGS="-DLLAMA_CUBLAS=ON"

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

# 启动服务器的命令
CMD ["python3", "./server.py"]

解决说明 #

使用nvidia/cuda开发镜像作为基础，确保容器内具备完整的CUDA开发环境，编译llama-cpp-python时能正确链接CUDA依赖；通过ENV设置全局环境变量，避免了RUN命令中环境变量传递可能出现的问题，确保编译过程中LLAMA_CUBLAS参数生效。

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