run

#!/bin/bash

set -e

if [[ -n "${RASTER_VISION_DEBUG}" ]]; then
    set -x
fi

SOURCE="${BASH_SOURCE[0]}"
while [ -h "$SOURCE" ] ; do SOURCE="$(readlink "$SOURCE")"; done
DOCKER_DIR="$( cd -P "$( dirname "$SOURCE" )" && pwd )"
REPO_ROOT="$( cd -P "$( dirname "$DOCKER_DIR" )" && pwd )"

DOCKER_VERSION=$(docker version -f "{{.Server.Version}}")
DOCKER_MAJOR=$(echo "$DOCKER_VERSION"| cut -d'.' -f 1)

if [ "${DOCKER_MAJOR}" -ge 19 ]; then
    runtime="--gpus=all"
else
    runtime="--runtime=nvidia"
fi

function usage() {
    echo -n \
         "Usage: $(basename "$0") [<options>] [<additional docker run args>]

Run a console in a Raster Vision Docker image locally.
By default, the raster-vision-pytorch image is used in the CPU runtime.

Environment variables:
RASTER_VISION_DATA_DIR (directory for storing data; mounted to /opt/data)
RASTER_VISION_NOTEBOOK_DIR (optional directory for Jupyter notebooks; mounted to /opt/notebooks)
AWS_PROFILE (optional AWS profile)

Options:
--aws forwards AWS credentials (sets AWS_PROFILE env var and mounts ~/.aws to /root/.aws)
--tensorboard maps port 6006
--name sets the name of the running container
--jupyter forwards a port, mounts RASTER_VISION_NOTEBOOK_DIR to /opt/notebooks, and runs Jupyter
--jupyter-lab forwards a port, mounts RASTER_VISION_NOTEBOOK_DIR to /opt/notebooks, and runs Jupyter Lab
--docs runs the docs server and forwards port 8000
--debug forwards port 3000 for use with remote debugger
--gpu use nvidia runtime
--arm64 uses image built for arm64 architecture

All arguments after above options are passed to 'docker run'.
"
}

# finds an unused port in a given range and assigns it to the FREE_PORT global
# variable. Adapted from answers in:
# https://unix.stackexchange.com/questions/55913/whats-the-easiest-way-to-find-an-unused-local-port
function find_free_port_in_range() {
    LOWER=$1
    UPPER=$2
    for (( PORT = $LOWER ; PORT <= $UPPER ; PORT++ ))
    do
        ss -lpn | grep -q ":$PORT " || break
    done
    FREE_PORT="$PORT"
}

IMAGE="raster-vision-pytorch"
RASTER_VISION_DATA_DIR="${RASTER_VISION_DATA_DIR:-${REPO_ROOT}/data}"
RASTER_VISION_NOTEBOOK_DIR="${RASTER_VISION_NOTEBOOK_DIR:-${REPO_ROOT}/docs/usage/tutorials}"

# Parse options using scheme in
# https://stackoverflow.com/questions/192249/how-do-i-parse-command-line-arguments-in-bash
POSITIONAL=()
while [[ $# -gt 0 ]]
do
    key="$1"
    case $key in
        --help)
        usage
        exit 0
        shift
        ;;
        --aws)
        AWS="-e AWS_PROFILE=${AWS_PROFILE:-default} -v ${HOME}/.aws:/root/.aws:ro"
        shift # past argument
        ;;
        --arm64)
        IMAGE=${IMAGE}-arm64
        shift # past argument
        ;;
        --tensorboard)
        find_free_port_in_range 6006 7007
        TENSORBOARD="-p $FREE_PORT:$FREE_PORT"
        shift # past argument
        ;;
        --gpu)
        RUNTIME=$runtime
        shift # past argument
        ;;
        --name)
        shift
        NAME="--name $1"
        shift
        ;;
        --jupyter)
        find_free_port_in_range 8888 9999
        JUPYTER="-v ${RASTER_VISION_NOTEBOOK_DIR}:/opt/notebooks -p $FREE_PORT:$FREE_PORT"
        CMD=(jupyter notebook --ip 0.0.0.0 --port $FREE_PORT --no-browser --allow-root --notebook-dir=/opt/notebooks)
        echo "Starting Jupyter server at 0.0.0.0:$FREE_PORT. This may take a few seconds."
        shift
        ;;
        --jupyter-lab)
        find_free_port_in_range 8888 9999
        JUPYTER="-v ${RASTER_VISION_NOTEBOOK_DIR}:/opt/notebooks -v ${HOME}/.jupyter:/root/.jupyter -p $FREE_PORT:$FREE_PORT"
        # run jupyter lab in the background
        CMD=(/bin/bash -c "jupyter lab --ip 0.0.0.0 --port $FREE_PORT --no-browser --allow-root --notebook-dir=/opt/notebooks & bash")
        echo "Starting Jupyter Lab server at 0.0.0.0:$FREE_PORT. This may take a few seconds."
        shift
        ;;
        --docs)
        find_free_port_in_range 8000 8100
        DOCS="-p $FREE_PORT:$FREE_PORT"
        CMD=(/bin/bash -c "cd docs && make livehtml")
        shift
        ;;
        --debug)
        DEBUG="-p 3003:3000"
        shift
        ;;
        *)    # unknown option
        POSITIONAL+=("$1") # save it in an array for later
        shift # past argument
        ;;
    esac
done
set -- "${POSITIONAL[@]}" # restore positional parameters

if [ -z "${CMD}" ]
then
    CMD=(${@:1})
fi

if [ "${BASH_SOURCE[0]}" = "${0}" ]
then
    docker run \
        ${RUNTIME} ${NAME} --rm --ipc=host -it \
        -v "${HOME}"/.rastervision:/root/.rastervision \
        -v ${REPO_ROOT}/:/opt/src/ \
        -v ${RASTER_VISION_DATA_DIR}:/opt/data \
        ${TENSORBOARD} \
        ${AWS} \
        ${JUPYTER} \
        ${DOCS} \
        ${DEBUG} \
        ${IMAGE} \
        "${CMD[@]}"
fi