docker: Make CUDA_VERSION configurable (pytorch#48199)

Summary:
makes CUDA_VERSION configurable for the docker images:

make CUDA_VERSION=10.2 CUDNN_VERSION=7 official-image

Signed-off-by: Eli Uriegas <[email protected]>

Pull Request resolved: pytorch#48199

Reviewed By: xuzhao9, janeyx99

Differential Revision: D25064256

Pulled By: seemethere

fbshipit-source-id: 25f52185097be647d11b5324f9f97cd41cdad75b

Dockerfile

@@ -51,8 +51,10 @@ RUN --mount=type=cache,target=/opt/ccache \
 
 FROM conda as conda-installs
 ARG PYTHON_VERSION=3.8
+ARG CUDA_VERSION=11.0
 ARG INSTALL_CHANNEL=pytorch-nightly
-RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION} pytorch torchvision cudatoolkit=11.0.221 && \
+ENV CONDA_OVERRIDE_CUDA=${CUDA_VERSION}
+RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION} pytorch torchvision torchtext "cudatoolkit=${CUDA_VERSION}" && \
     /opt/conda/bin/conda clean -ya
 RUN /opt/conda/bin/pip install torchelastic
 

docker.Makefile

@@ -8,8 +8,10 @@ $(warning WARNING: No docker user found using results from whoami)
 DOCKER_ORG       = $(shell whoami)
 endif
 
+CUDA_VERSION     = 11.0
+CUDNN_VERSION    = 8
 BASE_RUNTIME     = ubuntu:18.04
-BASE_DEVEL       = nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04
+BASE_DEVEL       = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu18.04
 
 # The conda channel to use to install pytorch / torchvision
 INSTALL_CHANNEL  = pytorch
@@ -18,7 +20,10 @@ PYTHON_VERSION   = 3.7
 # Can be either official / dev
 BUILD_TYPE       = dev
 BUILD_PROGRESS   = auto
-BUILD_ARGS       = --build-arg BASE_IMAGE=$(BASE_IMAGE) --build-arg PYTHON_VERSION=$(PYTHON_VERSION) --build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL)
+BUILD_ARGS       = --build-arg BASE_IMAGE=$(BASE_IMAGE) \
+				   --build-arg PYTHON_VERSION=$(PYTHON_VERSION) \
+				   --build-arg CUDA_VERSION=$(CUDA_VERSION) \
+				   --build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL)
 DOCKER_BUILD     = DOCKER_BUILDKIT=1 docker build --progress=$(BUILD_PROGRESS) --target $(BUILD_TYPE) -t $(DOCKER_FULL_NAME):$(DOCKER_TAG) $(BUILD_ARGS) .
 DOCKER_PUSH      = docker push $(DOCKER_FULL_NAME):$(DOCKER_TAG)