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