Allow cmake vars in docker build (#100867)

Fixes #100866

Pull Request resolved: https://github.com/pytorch/pytorch/pull/100867
Approved by: https://github.com/malfet

Dockerfile

@@ -50,10 +50,12 @@ COPY . .
 RUN git submodule update --init --recursive
 
 FROM conda as build
+ARG CMAKE_VARS
 WORKDIR /opt/pytorch
 COPY --from=conda /opt/conda /opt/conda
 COPY --from=submodule-update /opt/pytorch /opt/pytorch
 RUN --mount=type=cache,target=/opt/ccache \
+    export eval ${CMAKE_VARS} && \
     TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
     CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
     python setup.py install

README.md

@@ -362,11 +362,19 @@ should increase shared memory size either with `--ipc=host` or `--shm-size` comm
 The `Dockerfile` is supplied to build images with CUDA 11.1 support and cuDNN v8.
 You can pass `PYTHON_VERSION=x.y` make variable to specify which Python version is to be used by Miniconda, or leave it
 unset to use the default.
+
 ```bash
 make -f docker.Makefile
 # images are tagged as docker.io/${your_docker_username}/pytorch
 ```
 
+You can also pass the `CMAKE_VARS="..."` environment variable to specify additional CMake variables to be passed to CMake during the build.
+See [setup.py](./setup.py) for the list of available variables.
+
+```bash
+CMAKE_VARS="BUILD_CAFFE2=ON BUILD_CAFFE2_OPS=ON" make -f docker.Makefile
+```
+
 ### Building the Documentation
 
 To build documentation in various formats, you will need [Sphinx](http://www.sphinx-doc.org) and the

docker.Makefile

@@ -12,6 +12,7 @@ CUDA_VERSION              = 11.7.0
 CUDNN_VERSION             = 8
 BASE_RUNTIME              = ubuntu:18.04
 BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu18.04
+CMAKE_VARS               ?=
 
 # The conda channel to use to install cudatoolkit
 CUDA_CHANNEL              = nvidia
@@ -31,7 +32,8 @@ BUILD_ARGS                = --build-arg BASE_IMAGE=$(BASE_IMAGE) \
 							--build-arg CUDA_CHANNEL=$(CUDA_CHANNEL) \
 							--build-arg PYTORCH_VERSION=$(PYTORCH_VERSION) \
 							--build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL) \
-							--build-arg TRITON_VERSION=$(TRITON_VERSION)
+							--build-arg TRITON_VERSION=$(TRITON_VERSION) \
+							--build-arg CMAKE_VARS="$(CMAKE_VARS)"
 EXTRA_DOCKER_BUILD_FLAGS ?=
 
 BUILD                    ?= build