From c9aee3e35654929a15a308758415b96bcd2dd9cc Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Wed, 7 Dec 2022 12:31:46 -0800 Subject: [PATCH] conda: Add CUDA_HOME, cuda binaries to path Adds CUDA_HOME and the cuda binaries to our path. Users were experiencing issues building GPU binaries with these so this should make it so that the proper environment variables are setup by default Signed-off-by: Eli Uriegas --- conda/Dockerfile | 4 ++++ conda/build_docker.sh | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/conda/Dockerfile b/conda/Dockerfile index 70166bd4b..c65e1ad99 100644 --- a/conda/Dockerfile +++ b/conda/Dockerfile @@ -41,8 +41,12 @@ RUN bash ./install_conda.sh && rm install_conda.sh # Install CUDA FROM base as cuda +ARG CUDA_VERSION=10.2 RUN rm -rf /usr/local/cuda-* ADD ./common/install_cuda.sh install_cuda.sh +ENV CUDA_HOME=/usr/local/cuda-${CUDA_VERSION} +# Make things in our path by default +ENV PATH=/usr/local/cuda-${CUDA_VERSION}/bin:$PATH FROM cuda as cuda11.6 RUN bash ./install_cuda.sh 11.6 diff --git a/conda/build_docker.sh b/conda/build_docker.sh index db7c5be62..a3385ff38 100755 --- a/conda/build_docker.sh +++ b/conda/build_docker.sh @@ -52,6 +52,11 @@ if [[ "${DOCKER_TAG}" =~ ^cuda* ]]; then set -x docker tag ${DOCKER_IMAGE} "pytorch/conda-builder:cuda${CUDA_VERSION/./}" ) + # Test that we're using the right CUDA compiler + ( + set -x + docker run --rm "${DOCKER_IMAGE}" nvcc --version | grep "cuda_${CUDA_VERSION}" + ) fi if [[ -n ${GITHUB_REF} ]]; then