#!/bin/bash
#
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


set -o errexit
set -o nounset

_GPU_INSTALLER_ENV_KEY="cos-gpu-installer-env"
# Environment variable file should be stored in the same directory as
# this script.
_GPU_INSTALLER_ENV_PATH="$(dirname $0)/gpu-installer-env"

# The following environment variables may be changed by cos-gpu-installer-env.
COS_NVIDIA_INSTALLER_CONTAINER="gcr.io/cos-cloud/cos-gpu-installer:latest"
NVIDIA_INSTALL_DIR_HOST="/var/lib/nvidia"
NVIDIA_INSTALL_DIR_CONTAINER="/usr/local/nvidia"
ROOT_MOUNT_DIR="/root"

setup() {
  # Always use environment variable from metadata if provided.
  if /usr/share/google/get_metadata_value "attributes/${_GPU_INSTALLER_ENV_KEY}" \
    > /tmp/gpu-installer-env; then
    cp -f /tmp/gpu-installer-env "${_GPU_INSTALLER_ENV_PATH}"
  fi
  source "${_GPU_INSTALLER_ENV_PATH}"

  mkdir -p "${NVIDIA_INSTALL_DIR_HOST}"
  # Make NVIDIA_INSTALL_DIR_HOST executable by bind mounting it.
  mount --bind "${NVIDIA_INSTALL_DIR_HOST}" "${NVIDIA_INSTALL_DIR_HOST}"
  mount -o remount,exec "${NVIDIA_INSTALL_DIR_HOST}"
}

main() {
  setup
  docker run \
    --privileged \
    --net=host \
    --pid=host \
    --volume "${NVIDIA_INSTALL_DIR_HOST}":"${NVIDIA_INSTALL_DIR_CONTAINER}" \
    --volume /dev:/dev \
    --volume "/":"${ROOT_MOUNT_DIR}" \
    --env-file "${_GPU_INSTALLER_ENV_PATH}" \
    "${COS_NVIDIA_INSTALLER_CONTAINER}"
  # Verify installation.
  ${NVIDIA_INSTALL_DIR_HOST}/bin/nvidia-smi
}

main