tools/compute-images/scripts/usr-local-bin-detect-gpu.sh

   1 #!/bin/sh
   2 #
   3 # Copyright (C) The Arvados Authors. All rights reserved.
   4 #
   5 # SPDX-License-Identifier: Apache-2.0
   6
   7 set -e
   8 set -u
   9
  10 usage() {
  11     cat <<EOF
  12 usage: $0 SUBCOMMAND [options ...]
  13
  14 Subcommands:
  15   $0 enable [extension]
  16 EOF
  17 }
  18
  19 # Enumerate GPU devices on the host and output a standard "driver" name for
  20 # each one found. Currently only detects `nvidia`.
  21 detect_gpus() {
  22     # -vmm sets a machine-readable output format.
  23     # The -d option queries 3D controllers only.
  24     # You must stick with standard awk - no GNU extensions.
  25     lspci -vmm -d ::0302 | awk '
  26 BEGIN { FS="\t"; }
  27 ($1 != "Vendor:") { next; }
  28 (tolower($2) ~ /^nvidia/) { print "nvidia"; }
  29 '
  30 }
  31
  32 case "${1:-}" in
  33     "-?"|-h|--help|help)
  34         usage
  35         ;;
  36
  37     enable)
  38         src_ext="${2:-avail}"
  39         # Ensure src_ext starts with a dot
  40         src_ext=".${src_ext#.}"
  41         dst_dir=/run/modules-load.d
  42         mkdir -p "$dst_dir"
  43         detect_gpus | while read driver; do
  44             src="/etc/modules-load.d/$driver$src_ext"
  45             if [ -e "$src" ]; then
  46                 ln -s "$src" "$dst_dir/$driver.conf"
  47             fi
  48         done
  49         ;;
  50
  51     *)
  52         echo "$0: unknown subcommand: ${1:-}" >&2
  53         usage >&2
  54         exit 2
  55         ;;
  56 esac