#!/bin/ksh93

# A wrapper around docker.io, which just adds --runtime=nvidia to run and create
# command and passes $NV_GPU as NVIDIA_VISIBLE_DEVICES env var if set.

typeset NV_DOCKER="${NV_DOCKER:-docker}" ARG VERSION='2.0.3.18.09.2'

function checkNvRuntime {
	typeset ARG='/usr/bin/nvidia-container-runtime' DV BEST
	[[ -x ${ARG} ]] && return 0

	typeset -a LINE ALL
	print -u2 "Missing the executable '${ARG}'."
	DV=${ dpkg-query -W -f='${Version}' docker.io ; }
	[[ -n ${DV} ]] && DV=${DV%%-*} || \
		print -u2 'Ooops, docker.io seems not to be installed.'
	apt-cache madison nvidia-container-runtime 2>/dev/null | while read -A L
	do
		if [[ -z ${BEST} ]]; then
			X=${L[2]%-*}
			[[ ${DV} == ${X#*docker} ]] && BEST="${L[2]}" && continue
		fi
		ALL+="    apt-get install nvidia-container-runtime=${L[2]}\n"
	done
	if [[ -z ${ALL} && -z ${BEST} ]]; then
		print -u2 'nvidia-container-runtime package infos are not available.'
		print -u2 "Please run '/usr/sbin/nvidia-Repos+Keys.sh -a' as root\n" \
			'to add the nvidia repository infos to your system.'
		return 2
	fi
	X="The package 'nvidia-container-runtime' should be [re-]installed on this"
	X+=' system.\n\nRun:
   apt-get install nvidia-container-runtime
to install the latest version'
	if [[ -n ${BEST} ]]; then
		X+=", or\n   apt-get install nvidia-container-runtime=${BEST}\nto "
		X+='install the version which matches the installed docker.io package'
	fi
	if [[ -n ${ALL} ]]; then
		X+=", or any of\n${ALL}to install the version you want (but may not"
		X+='work at all)'
	fi
	print "$X"'.

To prevent the automatic upgrade to the most recent version just run:
    apt-mark hold nvidia-container-runtime
and the same command but with "unhold" to release the lock.
'
	return 3
}

checkNvRuntime || return 99

typeset -a SEEN
typeset -a NEW_ARGS

while (( $# )); do
	SEEN+=( "$1" )
	ARG="$1"
	shift
	case "${ARG}" in
		run|create)
			NEW_ARGS+=( '--runtime=nvidia' )
			[[ -n ${NV_GPU} ]] && \
				NEW_ARGS+=( '-e' 'NVIDIA_VISIBLE_DEVICES='"${NV_GPU// /,}" )
			break
			;;
		version)
			print "NVIDIA Docker: ${VERSION}"
			break
			;;
		--)
			break
			;;
	esac
done

[[ -n ${NV_DEBUG} ]] && set -x

exec ${NV_DOCKER} "${SEEN[@]}" "${NEW_ARGS[@]}" "$@"