1#!/usr/bin/env bash 2 3# Licensed to the Apache Software Foundation (ASF) under one 4# or more contributor license agreements. See the NOTICE file 5# distributed with this work for additional information 6# regarding copyright ownership. The ASF licenses this file 7# to you under the Apache License, Version 2.0 (the 8# "License"); you may not use this file except in compliance 9# with the License. You may obtain a copy of the License at 10# 11# http://www.apache.org/licenses/LICENSE-2.0 12# 13# Unless required by applicable law or agreed to in writing, 14# software distributed under the License is distributed on an 15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 16# KIND, either express or implied. See the License for the 17# specific language governing permissions and limitations 18# under the License. 19 20# build and install are separated so changes to build don't invalidate 21# the whole docker cache for the image 22 23# Assumes base image is from nvidia/cuda 24 25set -ex 26 27if [ -z ${CUDNN_VERSION} ]; then 28 echo "Error: CUDNN_VERSION environment variable undefiend" 29 exit 1 30fi 31 32apt-get update || true 33 34case ${CUDA_VERSION} in 35 11\.0*) 36 export libcudnn_package="libcudnn8" 37 export libcudnn_version="${CUDNN_VERSION}-1+cuda11.0" 38 export libcudnn_dev_version="${CUDNN_VERSION}-1+cuda11.0" 39 ;; 40 10\.2*) 41 export libcudnn_package="libcudnn8" 42 export libcudnn_version="${CUDNN_VERSION}-1+cuda10.2" 43 export libcudnn_dev_version="${CUDNN_VERSION}-1+cuda10.2" 44 ;; 45 10\.1*) 46 export libcudnn_package="libcudnn7" 47 export libcudnn_version="${CUDNN_VERSION}-1+cuda10.1" 48 export libcudnn_dev_version="${CUDNN_VERSION}-1+cuda10.1" 49 ;; 50 10\.0*) 51 export libcudnn_package="libcudnn7" 52 export libcudnn_version="${CUDNN_VERSION}-1+cuda10.0" 53 export libcudnn_dev_version="${CUDNN_VERSION}-1+cuda10.0" 54 ;; 55 9\.0*) 56 export libcudnn_package="libcudnn7" 57 export libcudnn_version="${CUDNN_VERSION}-1+cuda9.0" 58 export libcudnn_dev_version="${CUDNN_VERSION}-1+cuda9.0" 59 ;; 60 9\.2*) 61 export libcudnn_package="libcudnn7" 62 export libcudnn_version="${CUDNN_VERSION}-1+cuda9.2" 63 export libcudnn_dev_version="${CUDNN_VERSION}-1+cuda9.2" 64 ;; 65 *) 66 echo "Unsupported CUDA version ${CUDA_VERSION}" 67 exit 1 68 ;; 69esac 70 71apt-get install -y --allow-downgrades ${libcudnn_package}=${libcudnn_version} ${libcudnn_package}-dev=${libcudnn_dev_version} 72 73