m105 contains multiple cuda versions, select the one we are building against

Author: djherbis

Dockerfile.tmpl

@@ -25,6 +25,8 @@ ARG CUDA_MAJOR_VERSION
 ARG CUDA_MINOR_VERSION
 ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION}
 ENV CUDA_MINOR_VERSION=${CUDA_MINOR_VERSION}
+# Make sure we are on the right version of CUDA
+RUN update-alternatives --set cuda /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION
 # NVIDIA binaries from the host are mounted to /opt/bin.
 ENV PATH=/opt/bin:${PATH}
 # Add CUDA stubs to LD_LIBRARY_PATH to support building the GPU image on a CPU machine.

Jenkinsfile

@@ -52,7 +52,12 @@ pipeline {
               set -exo pipefail
               source config.txt
               cd packages/
-              ./build_package --base-image $BASE_IMAGE_REPO/$GPU_BASE_IMAGE_NAME:$BASE_IMAGE_TAG --package lightgbm --version $LIGHTGBM_VERSION --push
+              ./build_package --base-image $BASE_IMAGE_REPO/$GPU_BASE_IMAGE_NAME:$BASE_IMAGE_TAG \
+                --package lightgbm \
+                --version $LIGHTGBM_VERSION \
+                --build-arg CUDA_MAJOR_VERSION=$CUDA_MAJOR_VERSION \
+                --build-arg CUDA_MINOR_VERSION=$CUDA_MINOR_VERSION \
+                --push
             '''
           }
         }

packages/lightgbm.Dockerfile

@@ -3,6 +3,11 @@ ARG BASE_IMAGE
 FROM ${BASE_IMAGE} AS builder
 
 ARG PACKAGE_VERSION
+ARG CUDA_MAJOR_VERSION
+ARG CUDA_MINOR_VERSION
+
+# Make sure we are on the right version of CUDA
+RUN update-alternatives --set cuda /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION
 
 # Build instructions: https://lightgbm.readthedocs.io/en/latest/GPU-Tutorial.html#build-lightgbm
 RUN apt-get update && \

packages/torch.Dockerfile

@@ -9,6 +9,9 @@ ARG TORCHVISION_VERSION
 ARG CUDA_MAJOR_VERSION
 ARG CUDA_MINOR_VERSION
 
+# Make sure we are on the right version of CUDA
+RUN update-alternatives --set cuda /usr/local/cuda-$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION
+
 # TORCHVISION_VERSION is mandatory
 RUN test -n "$TORCHVISION_VERSION"