nVidia 525 + Cuda 11.8 + Python 3.10 + pyTorch GPU Docker image

It is a base environment for torch with GPU support (including 3090Ti!) that can be used for working with AI models. This image requires nvidia-driver-525 and nvidia-docker2 installed on host. It needs 30GB on disk!

Dockerfile

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

ENV PYTHONUNBUFFERED=1 

# SYSTEM
RUN apt-get update --yes --quiet && DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \
    software-properties-common \
    build-essential apt-utils \
    wget curl vim git ca-certificates kmod \
    nvidia-driver-525 \
 && rm -rf /var/lib/apt/lists/*

# PYTHON 3.10
RUN add-apt-repository --yes ppa:deadsnakes/ppa && apt-get update --yes --quiet
RUN DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \
    python3.10 \
    python3.10-dev \
    python3.10-distutils \
    python3.10-lib2to3 \
    python3.10-gdbm \
    python3.10-tk \
    pip

RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 999 \
    && update-alternatives --config python3 && ln -s /usr/bin/python3 /usr/bin/python

RUN pip install --upgrade pip

# ANACONDA
RUN wget -O /tmp/anaconda.sh https://repo.anaconda.com/archive/Anaconda3-2022.10-Linux-x86_64.sh \
    && bash /tmp/anaconda.sh -b -p /anaconda \
    && eval "$(/anaconda/bin/conda shell.bash hook)" \
    && conda init \
    && conda update -n base -c defaults conda \
    && conda create --name env \
    && conda activate env \
    && conda install -y pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch-nightly -c nvidia

Build

$ docker build -t nvidia-cuda .

Run

$ docker run --gpus all -it nvidia-cuda

Test

Run inside container:

$ conda activate env
$ python

>>> include torch
>>> torch.cuda.is_available() 

Comments

[oranusa]

how large is this image? the offical torch-cuda image seems to be around 18G

[Dharma]

My host machine has driver version 525.147.05, and the version I installed inside the container is nvidia-driver-515. Why is the nvidia-smi command inside the container still showing the host machine's driver version?