构建含cuML与TensorFlow-GPU的Conda环境遇依赖冲突及GPU识别问题

问题根源 #

RAPIDS 23.04硬性依赖cudatoolkit=11.2，旧版tensorflow-gpu包的版本管理逻辑会与conda提供的cudatoolkit产生依赖冲突；直接替换为tensorflow后，若未正确配置容器GPU运行时环境或TF版本不匹配CUDA，会导致GPU无法被识别。

解决方案步骤 #

1. 修正environment.yml，对齐CUDA与TensorFlow版本 #

选用兼容cudatoolkit11.2的TensorFlow版本（2.9.x或2.10.x），移除tensorflow-gpu（TF 2.10+已统一为tensorflow包，包含GPU支持）。示例配置：

name: rapids-tf
channels:
  - rapidsai
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.10
  - cudatoolkit=11.2
  - cudf=23.04
  - cuml=23.04
  - cugraph=23.04
  - tensorflow=2.10.0  # 与cudatoolkit11.2完全兼容
  - pip
  - pip:
    - nvidia-cudnn-cu11==8.1.0.77  # 匹配CUDA11.2的cuDNN版本

2. 调整Dockerfile确保环境正确初始化 #

确保conda环境激活，且CUDA相关库路径被添加到系统环境变量：

FROM continuumio/miniconda3:latest

# 复制环境配置文件
COPY environment.yml .

# 创建并激活conda环境
RUN conda env create -f environment.yml && conda clean -afy

# 设置环境变量，确保TF能找到CUDA库
ENV PATH /opt/conda/envs/rapids-tf/bin:$PATH
ENV LD_LIBRARY_PATH /opt/conda/envs/rapids-tf/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64

# 设置默认激活环境
RUN echo "source activate rapids-tf" > ~/.bashrc
CMD ["/bin/bash"]

3. 配置docker-compose启用NVIDIA Runtime #

确保docker-compose.yml指定NVIDIA运行时，并映射GPU资源：

version: '3.8'
services:
  rapids-tf:
    build: .
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./workspace:/workspace
    working_dir: /workspace
    command: bash -c "python -c 'import tensorflow as tf; print(tf.config.list_physical_devices(\"GPU\"))' && python -c 'import cuml; print(cuml.__version__)'"

4. 验证环境 #

构建并启动容器后，检查输出：

• TensorFlow GPU识别成功的输出示例：[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
• cuML版本输出：23.04.00

关键注意事项 #

• 必须确保cudatoolkit、tensorflow、nvidia-cudnn-cu11三者版本严格匹配：CUDA11.2对应TF2.9/2.10，cuDNN8.1.x
• 避免同时用conda和pip安装CUDA相关库，防止版本冲突
• Docker主机需已安装NVIDIA Docker Toolkit，否则容器无法访问GPU

内容的提问来源于stack exchange，提问作者raymond.mh.ng