#!/bin/bash

set -ex

[ -n "$UBUNTU_VERSION" ]
[ -n "$CUDA_VERSION" ]

# The NCCL version is not encoded in the build environment.
# This file installs the latest version that works.

# There are only NCCL packages for Ubuntu 16.04 and 14.04
if [[ "$UBUNTU_VERSION" == 16.04 ]]; then
  NCCL_UBUNTU_VER=ubuntu1604
  NCCL_DEB='nvidia-machine-learning-repo-ubuntu1604_1.0.0-1_amd64.deb'
elif [[ "$UBUNTU_VERSION" == 14.04 ]]; then
  NCCL_UBUNTU_VER=ubuntu1404
  NCCL_DEB='nvidia-machine-learning-repo-ubuntu1404_4.0-2_amd64.deb'
else
  echo "There is no NCCL package for Ubuntu version ${UBUNTU_VERSION}."
  echo "    NCCL will not be installed."
fi

if [ -n "$NCCL_UBUNTU_VER" ]; then

  # The deb is agnostic of CUDA version
  curl -LO "http://developer.download.nvidia.com/compute/machine-learning/repos/${NCCL_UBUNTU_VER}/x86_64/${NCCL_DEB}"

  # This dpkg call needs wget
  apt-get update
  apt-get install -y wget
  dpkg -i "${NCCL_DEB}"

  if [[ "$CUDA_VERSION" =~ ^10.* ]]; then
    NCCL_LIB_VERSION="2.4.8-1+cuda${CUDA_VERSION:0:4}"
  else
    NCCL_LIB_VERSION="2.1.15-1+cuda${CUDA_VERSION:0:3}"
  fi

  apt update
  apt install -y --allow-downgrades --allow-change-held-packages libnccl2=$NCCL_LIB_VERSION libnccl-dev=$NCCL_LIB_VERSION
fi
