#! /bin/bash
# Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.

NV_DOCKER=${NV_DOCKER:-"docker"}

DOCKER_ARGS=()
NV_DOCKER_ARGS=()
while [ $# -gt 0 ]; do
    arg=$1
    shift
    DOCKER_ARGS+=("$arg")
    case $arg in
        run|create)
            NV_DOCKER_ARGS+=("--runtime=nvidia")
            if [ ! -z "${NV_GPU}" ]; then
                NV_DOCKER_ARGS+=(-e NVIDIA_VISIBLE_DEVICES="${NV_GPU// /,}")
            fi
            break
        ;;
        version)
            printf "NVIDIA Docker: 2.5.0\n"
            break
        ;;
        --)
            break
        ;;
    esac
done

if [ ! -z $NV_DEBUG ]; then
    set -x
fi

exec $NV_DOCKER "${DOCKER_ARGS[@]}" "${NV_DOCKER_ARGS[@]}" "$@"