#!/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='2.0.3.18.09.2'

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[@]}" "$@"