hooks/nvidia

   1 #! /bin/bash
   2
   3 # Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved.
   4
   5 set -eu
   6
   7 # NVIDIA_VISIBLE_DEVICES="" *or* NVIDIA_VISIBLE_DEVICES="void"
   8 # GPU support was explicitly disabled, exit early.
   9 if [ -z "${NVIDIA_VISIBLE_DEVICES-x}" ] || [ "${NVIDIA_VISIBLE_DEVICES:-}" = "void" ]; then
  10     exit 0
  11 fi
  12
  13 # https://github.com/nvidia/nvidia-container-runtime#cuda_version
  14 if [ -n "${CUDA_VERSION:-}" ] && [ -z "${NVIDIA_REQUIRE_CUDA:-}" ]; then
  15     # Legacy CUDA image: default to all devices and all driver capabilities.
  16     if [ -z "${NVIDIA_VISIBLE_DEVICES+x}" ]; then
  17         NVIDIA_VISIBLE_DEVICES="all"
  18     fi
  19     if [ -z "${NVIDIA_DRIVER_CAPABILITIES:-}" ]; then
  20         NVIDIA_DRIVER_CAPABILITIES="all"
  21     fi
  22     if [[ "${CUDA_VERSION}" =~ ^[0-9]+\.[0-9]+ ]]; then
  23         NVIDIA_REQUIRE_CUDA="cuda>=${BASH_REMATCH[0]}"
  24     fi
  25 else
  26     # NVIDIA_VISIBLE_DEVICES unset and it's not a legacy CUDA image.
  27     # This is not a GPU image, exit early.
  28     if [ -z "${NVIDIA_VISIBLE_DEVICES+x}" ]; then
  29         exit 0
  30     fi
  31 fi
  32
  33 export PATH=$PATH:/usr/sbin:/usr/bin:/sbin:/bin
  34 if ! which nvidia-container-cli >/dev/null; then
  35     echo "ERROR: Missing tool nvidia-container-cli, see https://github.com/NVIDIA/libnvidia-container" >&2
  36     exit 1
  37 fi
  38
  39 in_userns() {
  40     [ -e /proc/self/uid_map ] || { echo no; return; }
  41     while read line; do
  42         fields=$(echo $line | awk '{ print $1 " " $2 " " $3 }')
  43         [ "$fields" = "0 0 4294967295" ] && { echo no; return; } || true
  44         echo $fields | grep -q " 0 1$" && { echo userns-root; return; } || true
  45     done < /proc/self/uid_map
  46
  47     [ "$(cat /proc/self/uid_map)" = "$(cat /proc/1/uid_map)" ] && \
  48         { echo userns-root; return; }
  49     echo yes
  50 }
  51
  52 get_ldconfig() {
  53     which "ldconfig.real" || which "ldconfig"
  54     return $?
  55 }
  56
  57 capability_to_cli() {
  58     case "$1" in
  59         compute)  echo "--compute";;
  60         compat32) echo "--compat32";;
  61         display)  echo "--display";;
  62         graphics) echo "--graphics";;
  63         utility)  echo "--utility";;
  64         video)    echo "--video";;
  65         *)        exit 1;;
  66     esac
  67     return
  68 }
  69
  70 # Same behavior as strconv.ParseBool in golang
  71 parse_bool() {
  72     case "$1" in
  73         1|t|T|TRUE|true|True)    echo "true";;
  74         0|f|F|FALSE|false|False) echo "false";;
  75         *)                       exit 1;;
  76     esac
  77     return
  78 }
  79
  80 usage() {
  81     cat <<EOF
  82 nvidia-container-cli hook for LXC
  83
  84 Special arguments:
  85 [ -h | --help ]: Print this help message and exit.
  86
  87 Optional arguments:
  88 [ --no-load-kmods ]: Do not try to load the NVIDIA kernel modules.
  89 [ --disable-require ]: Disable all the constraints of the form NVIDIA_REQUIRE_*.
  90 [ --debug <path> ]: The path to the log file.
  91 [ --ldcache <path> ]: The path to the host system's DSO cache.
  92 [ --root <path> ]: The path to the driver root directory.
  93 [ --ldconfig <path> ]: The path to the ldconfig binary, use a '@' prefix for a host path.
  94 EOF
  95     return 0
  96 }
  97
  98 options=$(getopt -o h -l help,no-load-kmods,disable-require,debug:,ldcache:,root:,ldconfig: -- "$@")
  99 if [ $? -ne 0 ]; then
 100     usage
 101     exit 1
 102 fi
 103 eval set -- "$options"
 104
 105 CLI_LOAD_KMODS="true"
 106 CLI_DISABLE_REQUIRE="false"
 107 CLI_DEBUG=
 108 CLI_LDCACHE=
 109 CLI_ROOT=
 110 CLI_LDCONFIG=
 111
 112 while :; do
 113     case "$1" in
 114         --help)             usage && exit 1;;
 115         --no-load-kmods)    CLI_LOAD_KMODS="false"; shift 1;;
 116         --disable-require)  CLI_DISABLE_REQUIRE="true"; shift 1;;
 117         --debug)            CLI_DEBUG=$2; shift 2;;
 118         --ldcache)          CLI_LDCACHE=$2; shift 2;;
 119         --root)             CLI_ROOT=$2; shift 2;;
 120         --ldconfig)         CLI_LDCONFIG=$2; shift 2;;
 121         --)                 shift 1; break;;
 122         *)                  break;;
 123     esac
 124 done
 125
 126 HOOK_SECTION=
 127 HOOK_TYPE=
 128 case "${LXC_HOOK_VERSION:-0}" in
 129     0) HOOK_SECTION="${2:-}"; HOOK_TYPE="${3:-}";;
 130     1) HOOK_SECTION="${LXC_HOOK_SECTION:-}"; HOOK_TYPE="${LXC_HOOK_TYPE:-}";;
 131     *) echo "ERROR: Unsupported hook version: ${LXC_HOOK_VERSION}." >&2; exit 1;;
 132 esac
 133
 134 if [ "${HOOK_SECTION}" != "lxc" ]; then
 135     echo "ERROR: Not running through LXC." >&2
 136     exit 1
 137 fi
 138
 139 if [ "${HOOK_TYPE}" != "mount" ]; then
 140     echo "ERROR: This hook must be used as a \"mount\" hook." >&2
 141     exit 1
 142 fi
 143
 144 USERNS=$(in_userns)
 145 if [ "${USERNS}" != "yes" ]; then
 146     # This is a limitation of libnvidia-container.
 147     echo "FIXME: This hook currently only works in unprivileged mode." >&2
 148     exit 1
 149 fi
 150
 151 if [ "${USERNS}" = "yes" ]; then
 152     CLI_LOAD_KMODS="false"
 153     if ! grep -q nvidia_uvm /proc/modules; then
 154         echo "WARN: Kernel module nvidia_uvm is not loaded, nvidia-container-cli might fail. Make sure the NVIDIA device driver is installed and loaded." >&2
 155     fi
 156 fi
 157
 158 # https://github.com/nvidia/nvidia-container-runtime#nvidia_disable_require
 159 if [ -n "${NVIDIA_DISABLE_REQUIRE:-}" ]; then
 160     if [ "$(parse_bool "${NVIDIA_DISABLE_REQUIRE}")" = "true" ]; then
 161         CLI_DISABLE_REQUIRE="true"
 162     fi
 163 fi
 164
 165 if [ -z "${CLI_DEBUG}" ]; then
 166     if [ "${LXC_LOG_LEVEL}" = "DEBUG" ] || [ "${LXC_LOG_LEVEL}" = "TRACE" ]; then
 167         rootfs_path="${LXC_ROOTFS_PATH#*:}"
 168         hookdir="${rootfs_path/%rootfs/hook}"
 169         if mkdir -p "${hookdir}"; then
 170             CLI_DEBUG="${hookdir}/nvidia.log"
 171         fi
 172     fi
 173 fi
 174
 175 # A '@' prefix means a host path.
 176 if [ -z "${CLI_LDCONFIG}" ]; then
 177     if host_ldconfig=$(get_ldconfig); then
 178         CLI_LDCONFIG="@${host_ldconfig}"
 179     fi
 180 fi
 181
 182 # https://github.com/nvidia/nvidia-container-runtime#nvidia_visible_devices
 183 CLI_DEVICES="${NVIDIA_VISIBLE_DEVICES}"
 184
 185 # https://github.com/nvidia/nvidia-container-runtime#nvidia_driver_capabilities
 186 CLI_CAPABILITIES=
 187 if [ -n "${NVIDIA_DRIVER_CAPABILITIES:-}" ]; then
 188      CLI_CAPABILITIES="${NVIDIA_DRIVER_CAPABILITIES//,/ }"
 189 fi
 190
 191 # https://github.com/nvidia/nvidia-container-runtime#nvidia_require_
 192 CLI_REQUIREMENTS=
 193 for req in $(compgen -e "NVIDIA_REQUIRE_"); do
 194     CLI_REQUIREMENTS="${CLI_REQUIREMENTS} ${!req}"
 195 done
 196
 197 if [ "${CLI_CAPABILITIES}" = "all" ]; then
 198     CLI_CAPABILITIES="compute compat32 display graphics utility video"
 199 fi
 200
 201 if [ -z "${CLI_CAPABILITIES}" ]; then
 202     CLI_CAPABILITIES="utility"
 203 fi
 204
 205 global_args=("")
 206 configure_args=("")
 207
 208 if [ -n "${CLI_DEBUG}" ]; then
 209     echo "INFO: Writing nvidia-container-cli log at ${CLI_DEBUG}." >&2
 210     global_args+=("--debug=${CLI_DEBUG}")
 211 fi
 212
 213 if [ "${CLI_LOAD_KMODS}" = "true" ]; then
 214     global_args+=(--load-kmods)
 215 fi
 216
 217 if [ "${USERNS}" = "yes" ]; then
 218     global_args+=(--user)
 219     configure_args+=(--no-cgroups)
 220 fi
 221
 222 if [ -n "${CLI_LDCACHE}" ]; then
 223     global_args+=(--ldcache="${CLI_LDCACHE}")
 224 fi
 225
 226 if [ -n "${CLI_ROOT}" ]; then
 227     global_args+=(--root="${CLI_ROOT}")
 228 fi
 229
 230 if [ -n "${CLI_LDCONFIG}" ]; then
 231     configure_args+=(--ldconfig="${CLI_LDCONFIG}")
 232 fi
 233
 234 if [ -n "${CLI_DEVICES}" ] && [ "${CLI_DEVICES}" != "none" ]; then
 235     configure_args+=(--device="${CLI_DEVICES}")
 236 fi
 237
 238 for cap in ${CLI_CAPABILITIES}; do
 239     if arg=$(capability_to_cli "${cap}"); then
 240         configure_args+=("${arg}")
 241     else
 242         echo "ERROR: Unknown driver capability \"${cap}\"." >&2
 243         exit 1
 244     fi
 245 done
 246
 247 if [ "${CLI_DISABLE_REQUIRE}" = "false" ]; then
 248     for req in ${CLI_REQUIREMENTS}; do
 249         configure_args+=(--require="${req}")
 250     done
 251 fi
 252
 253 if [ -d "/sys/kernel/security/apparmor" ]; then
 254     # Try to transition to the unconfined AppArmor profile.
 255     echo "changeprofile unconfined" > /proc/self/attr/current || true
 256 fi
 257
 258 set -x
 259 exec nvidia-container-cli ${global_args[@]} configure ${configure_args[@]} "${LXC_ROOTFS_MOUNT}"