Create a GPU-enabled image for ACI

Question

I do not have access to V100 GPU locally. So I cannot build/push the image that is then needed to run in ACI container.

How can I build the image, that is to be pushed to ACR and then used by GPU-enabled container? I was thinking of using the ACI itself to build the image, but calling docker inside a container is not easy.

Answer 1

One solution I found (for my Pytorch extension case) is to build the docker without GPU. I match the cuda version and hard code the architecture.

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

ENV CUDA_HOME /usr/local/cuda-11.7/
ENV TORCH_CUDA_ARCH_LIST "Volta;Turing"
ENV FORCE_CUDA "1"

RUN pip install torch==1.13.1 torchvision==0.14.1 --extra-index-url https://download.pytorch.org/whl/cu117

Should make sure troch.version.cuda returns the version string, otherwise the CPU version is picked up during build.