| #!/bin/bash |
| # |
| # Copyright 2020 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| # |
| # Utility to manage COS extensions. |
| |
| set -o errexit |
| set -o pipefail |
| set -o nounset |
| |
| PROG_NAME=$(basename "$0") |
| readonly PROG_NAME |
| readonly OS_RELEASE="/etc/os-release" |
| readonly EXTENSIONS_CACHE="/var/lib/cos-extensions" |
| readonly ARTIFACT_REGISTRY_REGIONS=("us" "eu" "asia") |
| DEFAULT_GPU_INSTALLER="gcr.io/cos-cloud/cos-gpu-installer:v2.0.26" |
| |
| COS_GPU_INSTALLER="${COS_GPU_INSTALLER:-}" |
| |
| set_cos_gpu_installer() { |
| if [[ -n "${COS_GPU_INSTALLER}" ]]; then |
| return |
| fi |
| |
| # Change gpu container image name if location of the vm is available |
| # and it's from cos-cloud project. |
| local instance_zone |
| instance_zone="$(curl -s "http://metadata.google.internal/computeMetadata/v1/instance/zone?alt=text" -H "Metadata-Flavor: Google" || true)" |
| if [[ -n "${instance_zone}" ]]; then |
| # Syntax of instance_zone: projects/<project-id>/zones/<region> |
| # Example: projects/475556798229/zones/us-central1-a |
| local instance_region |
| instance_region="$(echo "${instance_zone}"|cut -d/ -f4|cut -d- -f1)" |
| for region in "${ARTIFACT_REGISTRY_REGIONS[@]}"; do |
| if [[ "${region}" == "${instance_region}" ]]; then |
| DEFAULT_GPU_INSTALLER="${instance_region}.${DEFAULT_GPU_INSTALLER}" |
| break |
| fi |
| done |
| fi |
| COS_GPU_INSTALLER="${DEFAULT_GPU_INSTALLER}" |
| } |
| |
| usage() { |
| cat <<EOF |
| |
| ${PROG_NAME}: Utility to manage COS extensions. |
| |
| Usage: |
| ${PROG_NAME} [OPTIONS] COMMAND [ARGS]... |
| |
| Options: |
| -h, --help print help message |
| |
| Commands: |
| list [-- --gpu-installer] list all available COS extensions and |
| their versions. |
| |
| install <extension> [-- -<version>] install a COS extension. If no version |
| is given, then the default version |
| will be installed. |
| |
| Additional Description: |
| ${PROG_NAME} install gpu -- --prepare-build-tools |
| |
| The gpu extension can be invoked with a --prepare-build-tools optional |
| argument that can be used to cache the toolchain for the installer. |
| Caching the toolchain carries the overhead of ~1GB disk space on the |
| stateful partition. |
| Using this command only populates the cache and does NOT install the GPU |
| drivers thus may saves time on downloading the toolchain during subsequent |
| installations. |
| |
| ${PROG_NAME} install gpu -- --clean-build--tools |
| Use this optional command to delete the cache for the toolchain present on |
| the stateful partition. |
| } |
| EOF |
| exit "${1}" |
| } |
| |
| parse_args() { |
| local args |
| if ! args=$(getopt --options "h" --longoptions "help" -- "$@"); then |
| usage 1 |
| fi |
| |
| eval set -- "${args}" |
| while true; do |
| case "$1" in |
| -h|--help) |
| usage 0 |
| ;; |
| --) |
| shift |
| break |
| ;; |
| *) |
| usage 1 |
| ;; |
| esac |
| done |
| |
| if [[ "$#" -eq 0 ]]; then |
| usage 1 |
| fi |
| |
| case "$1" in |
| list) |
| shift |
| list "$@" |
| ;; |
| install) |
| if [[ "$#" -eq 2 ]]; then |
| install "$2" |
| elif [[ "$#" -ge 3 ]]; then |
| extension="$2" |
| shift 2 |
| install "${extension}" "$@" |
| else |
| usage 1 |
| fi |
| ;; |
| *) |
| usage 1 |
| ;; |
| esac |
| } |
| |
| list() { |
| if [[ "$#" -eq 0 ]]; then |
| # shellcheck disable=SC2154 |
| printf "Available extensions for COS version %s-%s:\n\n" \ |
| "${VERSION_ID}" "${BUILD_ID}" |
| echo "[gpu]" |
| echo "gpu installer: ${DEFAULT_GPU_INSTALLER}" |
| run_gpu_installer list 2>/dev/null |
| elif [[ "$#" -eq 1 ]]; then |
| case "$1" in |
| --gpu-installer) |
| echo "${DEFAULT_GPU_INSTALLER}" |
| ;; |
| *) |
| echo "Unsupported argument $1" |
| usage 1 |
| ;; |
| esac |
| else |
| usage 1 |
| fi |
| } |
| |
| install() { |
| case "$1" in |
| gpu) |
| shift |
| run_gpu_installer install "-host-dir=/var/lib/nvidia" "$@" |
| ;; |
| *) |
| echo "Unsupported extension $1" |
| exit 1 |
| ;; |
| esac |
| } |
| |
| check_arch() { |
| arch=$(uname -m) |
| if [[ ${arch} != "x86_64" ]]; then |
| echo "GPU installation is only supported on X86 for now. |
| Current architecture detected: ${arch}" |
| exit 1 |
| fi |
| } |
| |
| run_gpu_installer() { |
| check_arch |
| local use_build_cache=false |
| local clean_cache=false |
| |
| local installer_args=() |
| for i in "$@"; do |
| if [[ ${i} == '--clean-build-tools' ]]; then |
| clean_cache=true |
| else |
| if [[ ${i} == '--prepare-build-tools' ]]; then |
| use_build_cache=true |
| fi |
| installer_args+=("${i}") |
| fi |
| done |
| |
| local docker_args=( |
| --rm |
| --name="cos-gpu-installer" |
| --privileged |
| --net=host |
| --pid=host |
| --volume /dev:/dev |
| --volume /:/root |
| ) |
| |
| if [[ "${clean_cache}" = true ]]; then |
| echo "Cleaning cache present at: ${EXTENSIONS_CACHE}" |
| rm -rf "${EXTENSIONS_CACHE}" |
| else |
| # use extensions cache(if it exists) by default |
| if [[ "${use_build_cache}" = true || -d ${EXTENSIONS_CACHE} ]]; then |
| docker_args+=(--volume "${EXTENSIONS_CACHE}/:/build/") |
| fi |
| |
| /usr/bin/docker run "${docker_args[@]}" "${COS_GPU_INSTALLER}"\ |
| "${installer_args[@]}" |
| fi |
| } |
| |
| main() { |
| # shellcheck source=/etc/os-release |
| source "${OS_RELEASE}" |
| set_cos_gpu_installer |
| parse_args "$@" |
| } |
| |
| main "$@" |