docker_setup.sh

#!/bin/bash
# Ensure that you have installed docker(API >= 1.40) and the nvidia graphics driver on host!
# Copyright 2018-2022, Kolin Guo, All rights reserved.

############################################################
# Section 0: Project-Specific Settings                     #
############################################################
REPONAME="maya-slam"
declare -a CONT_NAMES=("realsense_capture" "maplab" "orbslam3" "maya_recon" "openchisel")
#IMG_TAG="latest"
IMG_TAG="v1.0"
declare -a IMG_NAMES=("kolinguo/realsense:${IMG_TAG}"
                      "kolinguo/maplab:${IMG_TAG}"
                      "kolinguo/orbslam3:${IMG_TAG}"
                      "kolinguo/maya_recon:${IMG_TAG}"
                      "kolinguo/openchisel:${IMG_TAG}")
declare -a DOCKERFILE_PATHS=("./docker/Dockerfile_realsense"
                             "./docker/Dockerfile_maplab"
                             "./docker/Dockerfile_orbslam3"
                             "./docker/Dockerfile_maya_recon"
                             "./docker/Dockerfile_openchisel")
declare -a JUPYTER_PORTS=("" "" "9000" "9001" "")
# Additional port forwarding, define as variable that ends with "PORT"
#TENSORBOARDPORT="6007"

############################################################
# argbash configuration section                            #
############################################################
# Created by argbash-init v2.10.0
# Rearrange the order of options below according to what you would like to see in the help message.
# ARG_POSITIONAL_SINGLE([container-name],[container to setup: ${all_cont_names}])
# ARG_OPTIONAL_BOOLEAN([build-local],[l],[build docker image locally or pull from DockerHub],[off])
# ARG_OPTIONAL_BOOLEAN([build-only],[],[build image/container only and do not start the container],[off])
# ARG_OPTIONAL_BOOLEAN([rm-img],[],[remove previously built docker image],[off])
# ARG_OPTIONAL_BOOLEAN([bashrc-only],[],[generate only the custom bashrc],[off])
# ARGBASH_SET_INDENT([  ])
# ARGBASH_SET_DELIM([ =])
# ARG_OPTION_STACKING([getopt])
# ARG_RESTRICT_VALUES([no-local-options])
# ARG_HELP([Generic docker image/container setup script, adapted for maya-slam],[Copyright 2018-2022, Kolin Guo, All rights reserved.\n])
# ARG_VERSION_AUTO([v0.2.2],['Copyright 2018-2022, Kolin Guo, All rights reserved.'],[V])
# ARGBASH_GO()
# needed because of Argbash --> m4_ignore([
### START OF CODE GENERATED BY Argbash v2.10.0 one line above ###
# Argbash is a bash code generator used to get arguments parsing right.
# Argbash is FREE SOFTWARE, see https://argbash.io for more info


die()
{
  local _ret="${2:-1}"
  test "${_PRINT_HELP:-no}" = yes && print_help >&2
  echo "$1" >&2
  exit "${_ret}"
}


evaluate_strictness()
{
  [[ "$2" =~ ^-(-(container-name|build-local|build-only|rm-img|bashrc-only|help|version)$|[lhV]) ]] && die "You have passed '$2' as a value of argument '$1', which makes it look like that you have omitted the actual value, since '$2' is an option accepted by this script. This is considered a fatal error."
}


begins_with_short_option()
{
  local first_option all_short_options='lhV'
  first_option="${1:0:1}"
  test "$all_short_options" = "${all_short_options/$first_option/}" && return 1 || return 0
}

# THE DEFAULTS INITIALIZATION - POSITIONALS
_positionals=()
# THE DEFAULTS INITIALIZATION - OPTIONALS
_arg_build_local="off"
_arg_build_only="off"
_arg_rm_img="off"
_arg_bashrc_only="off"


print_help()
{
  local all_cont_names=$(printf "'%s' " "${CONT_NAMES[@]}")
  local usage="Usage: ${0} [OPTION]... <container-name>\n"
  usage+="Generic docker image/container setup script, adapted for maya-slam\n\n"
  usage+="  <container-name>        container to setup: ${all_cont_names}\n"
  usage+="  -l, --build-local       build docker image locally with Dockerfile\n"
  usage+="      --no-build-local    pull image from DockerHub (default)\n"
  usage+="  --build-only            build image and container only\n"
  usage+="      --no-build-only     start and attach the container (default)\n"
  usage+="  --rm-img                remove previously built docker image\n"
  usage+="      --no-rm-img         (off by default)\n"
  usage+="  --bashrc-only           generate only the custom bashrc\n"
  usage+="      --no-bashrc-only    (off by default)\n"
  usage+="      -h, --help     display this help and exit\n"
  usage+="      -V, --version  output version information and exit\n"
  usage+="\nCopyright 2018-2022, Kolin Guo, All rights reserved.\n"

  echo -e "$usage"
}


parse_commandline()
{
  _positionals_count=0
  while test $# -gt 0
  do
    _key="$1"
    case "$_key" in
      -l|--no-build-local|--build-local)
        _arg_build_local="on"
        test "${1:0:5}" = "--no-" && _arg_build_local="off"
        ;;
      -l*)
        _arg_build_local="on"
        _next="${_key##-l}"
        if test -n "$_next" -a "$_next" != "$_key"
        then
          { begins_with_short_option "$_next" && shift && set -- "-l" "-${_next}" "$@"; } || die "The short option '$_key' can't be decomposed to ${_key:0:2} and -${_key:2}, because ${_key:0:2} doesn't accept value and '-${_key:2:1}' doesn't correspond to a short option."
        fi
        ;;
      --no-build-only|--build-only)
        _arg_build_only="on"
        test "${1:0:5}" = "--no-" && _arg_build_only="off"
        ;;
      --no-rm-img|--rm-img)
        _arg_rm_img="on"
        test "${1:0:5}" = "--no-" && _arg_rm_img="off"
        ;;
      --no-bashrc-only|--bashrc-only)
        _arg_bashrc_only="on"
        test "${1:0:5}" = "--no-" && _arg_bashrc_only="off"
        ;;
      -h|--help)
        print_help
        exit 0
        ;;
      -h*)
        print_help
        exit 0
        ;;
      -V|--version)
        printf '%s %s\n\n%s\n%s\n' "docker_setup.sh" "v0.2.2" 'Generic docker image/container setup script, adapted for maya-slam' 'Copyright 2018-2022, Kolin Guo, All rights reserved.'
        exit 0
        ;;
      -V*)
        printf '%s %s\n\n%s\n%s\n' "docker_setup.sh" "v0.2.2" 'Generic docker image/container setup script, adapted for maya-slam' 'Copyright 2018-2022, Kolin Guo, All rights reserved.'
        exit 0
        ;;
      *)
        _last_positional="$1"
        _positionals+=("$_last_positional")
        _positionals_count=$((_positionals_count + 1))
        ;;
    esac
    shift
  done
}


handle_passed_args_count()
{
  local _required_args_string="'container-name'"
  test "${_positionals_count}" -ge 1 || _PRINT_HELP=yes die "FATAL ERROR: Not enough positional arguments - we require exactly 1 (namely: $_required_args_string), but got only ${_positionals_count}." 1
  test "${_positionals_count}" -le 1 || _PRINT_HELP=yes die "FATAL ERROR: There were spurious positional arguments --- we expect exactly 1 (namely: $_required_args_string), but got ${_positionals_count} (the last one was: '${_last_positional}')." 1
}


assign_positional_args()
{
  local _positional_name _shift_for=$1
  _positional_names="_arg_container_name "

  shift "$_shift_for"
  for _positional_name in ${_positional_names}
  do
    test $# -gt 0 || break
    eval "$_positional_name=\${1}" || die "Error during argument parsing, possibly an Argbash bug." 1
    shift
  done
}

parse_commandline "$@"
handle_passed_args_count
assign_positional_args 1 "${_positionals[@]}"

# OTHER STUFF GENERATED BY Argbash

### END OF CODE GENERATED BY Argbash (sortof) ### ])
# [ <-- needed because of Argbash

### Save the nice-looking usage ###
#print_help()
#{
#  local all_cont_names=$(printf "'%s' " "${CONT_NAMES[@]}")
#  local usage="Usage: ${0} [OPTION]... <container-name>\n"
#  usage+="Generic docker image/container setup script, adapted for maya-slam\n\n"
#  usage+="  <container-name>        container to setup: ${all_cont_names}\n"
#  usage+="  -l, --build-local       build docker image locally with Dockerfile\n"
#  usage+="      --no-build-local    pull image from DockerHub (default)\n"
#  usage+="  --build-only            build image and container only\n"
#  usage+="      --no-build-only     start and attach the container (default)\n"
#  usage+="  --rm-img                remove previously built docker image\n"
#  usage+="      --no-rm-img         (off by default)\n"
#  usage+="  --bashrc-only           generate only the custom bashrc\n"
#  usage+="      --no-bashrc-only    (off by default)\n"
#  usage+="      -h, --help     display this help and exit\n"
#  usage+="      -V, --version  output version information and exit\n"
#  usage+="\nCopyright 2018-2022, Kolin Guo, All rights reserved.\n"
#
#  echo -e "$usage"
#}

# vvv  PLACE YOUR CODE HERE  vvv
############################################################
# Section 1: Helper Function Definition                    #
############################################################
SCRIPT=$(readlink -f "$0")
SCRIPTPATH=$(dirname "$SCRIPT")
cd "$SCRIPTPATH"

test_retval() {
  if [ $? -ne 0 ] ; then
    echo -e "\nFailed to ${*}... Exiting...\n"
    exit 1
  fi
}

check_cont_name() {
  for (( i=0; i<${#CONT_NAMES[@]}; i++ )); do
    if [ "$_arg_container_name" = "${CONT_NAMES[$i]}" ] ; then
      IMGNAME="${IMG_NAMES[$i]}"
      CONTNAME="${CONT_NAMES[$i]}"
      DOCKERFILEPATH="${DOCKERFILE_PATHS[$i]}"
      JUPYTERPORT="${JUPYTER_PORTS[$i]}"
      break
    fi
  done

  if [ -z "${CONTNAME+x}" ] ; then
    _PRINT_HELP=yes die "FATAL ERROR: Unknown <container-name> '${_arg_container_name}'" 2
  fi
}

print_setup_info() {
  local ts="    "  # Tab spaces
  # Echo the set up information
  echo -e "\n"
  echo -e "################################################################################\n"
  echo -e "${ts}Docker Set Up Information\n"

  echo -e "${ts}${ts}container: ${CONTNAME}"
  if [ "$_arg_build_local" = "on" ] ; then
    echo -e "${ts}${ts}image: ${IMGNAME}"
    echo -e "${ts}${ts}Dockerfile: ${DOCKERFILEPATH}"
  else
    echo -e "${ts}${ts}Remote image: ${IMGNAME}"
  fi

  echo -e ""
  if [ "$_arg_bashrc_only" = "on" ] ; then
    echo -e "${ts}${ts}* Only generate the custom bashrc"
  fi
  if [ "$_arg_rm_img" = "on" ] ; then
    echo -e "${ts}${ts}* Cautious!! Remove previously built Docker image"
  else
    echo -e "${ts}${ts}* Keep previously built Docker image"
  fi
  echo -e ""
  echo -e "################################################################################\n"
}

remove_prev_docker_image() {
  # Remove previously built Docker image
  if [ "$_arg_rm_img" = "on" ] ; then
    echo -e "\nRemoving previously built image..."
    docker rmi -f "$IMGNAME"
  fi
}

create_custom_bashrc() {
  cat > bashrc <<- "EOF"
# If not running interactively, don't do anything
[ -z "$PS1" ] && return

# check the window size after each command and, if necessary,
# update the values of LINES and COLUMNS.
shopt -s checkwinsize

# enable bash completion in interactive shells
if ! shopt -oq posix; then
  if [ -f /usr/share/bash-completion/bash_completion ]; then
    . /usr/share/bash-completion/bash_completion
  elif [ -f /etc/bash_completion ]; then
    . /etc/bash_completion
  fi
fi

. /etc/bash_completion

# if the command-not-found package is installed, use it
if [ -x /usr/lib/command-not-found -o -x /usr/share/command-not-found/command-not-found ]; then
  function command_not_found_handle {
    # check because c-n-f could've been removed in the meantime
    if [ -x /usr/lib/command-not-found ]; then
      /usr/lib/command-not-found -- "$1"
      return $?
    elif [ -x /usr/share/command-not-found/command-not-found ]; then
      /usr/share/command-not-found/command-not-found -- "$1"
      return $?
    else
      printf "%s: command not found\n" "$1" >&2
      return 127
    fi
  }
fi

# Change PS1 and terminal color
export PS1="\[\e[31m\]${CONTNAME}-docker\[\e[m\] \[\e[33m\]\w\[\e[m\] > "  # match square bracket for argbash: ]]]]
sed -i "s/PS1=.*$/PS1=\$PS1/" /root/.bashrc  # Override PS1
export TERM=xterm-256color
alias grep="grep --color=auto"
alias ls="ls --color=auto"

# some more ls aliases
alias ll="ls -alF"
alias la="ls -A"
alias l="ls -CF"

# >>> conda initialize >>>
export CONDA_AUTO_ACTIVATE_BASE=false
# !! Contents within this block are managed by 'conda init' !!
__conda_setup="$('/miniconda/bin/conda' 'shell.bash' 'hook' 2> /dev/null)"
if [ $? -eq 0 ]; then
    eval "$__conda_setup"
else
    if [ -f "/miniconda/etc/profile.d/conda.sh" ]; then
        . "/miniconda/etc/profile.d/conda.sh"
    else
        export PATH="$PATH:/miniconda/bin"
    fi
fi
unset __conda_setup
# <<< conda initialize <<<

# Cyan color
echo -en "\e[1;36m"  # match square bracket for argbash: ]

# If any environment variable that starts with COMMANDTO
if printenv | egrep -q "^COMMANDTO.+="; then
  echo -e ""
  echo -e "################################################################################\n"
  if [ ! -z "$COMMANDTOCOMPILE" ] ; then
    echo -e "\tCommand to compile:\n\t\t${COMMANDTOCOMPILE}\n"
  fi
  if [ ! -z "$COMMANDTORUNJUPYTER" ] ; then
    echo -e "\tCommand to run jupyter notebook:\n\t\t${COMMANDTORUNJUPYTER}\n"
  fi
  if [ ! -z "$COMMANDTORUNTENSORBOARD" ] ; then
    echo -e "\tCommand to run TensorBoard:\n\t\t${COMMANDTORUNTENSORBOARD}\n"
  fi
  echo -e "################################################################################\n"
fi

# Turn off colors
echo -en "\e[m"  # match square bracket for argbash: ]

EOF

  # Overwrite CONTNAME
  sed -i "s/\${CONTNAME}/${CONTNAME}/" bashrc
}

update_port_forward() {
  PORT_FORWARD_DOCKER_CMD=""

  # Iterate over all local variables with names ending in PORT
  while IFS= read -r line; do
    port_name="${line%=*}"
    port_num="${line#*=}"
    #echo "Line: ${port_name} ${port_num}"
    PORT_FORWARD_DOCKER_CMD+=" -p ${port_num}:${port_num} "

    # Update Dockerfile environment variables
    sed -i "s/^ENV ${port_name}=.*$/ENV ${port_name}=\"${port_num}\"/" "$DOCKERFILEPATH"
  done < <(set -o posix ; set | egrep -o "^.+PORT=.+$")
}

build_docker_image() {
  # Set REPOPATH for WORKDIR
  sed -i "s/^ENV REPOPATH.*$/ENV REPOPATH \/${REPONAME}/" "$DOCKERFILEPATH"

  # Build and run the image
  echo -e "\nBuilding image ${IMGNAME}..."
  docker build -f "$DOCKERFILEPATH" -t "$IMGNAME" .
  test_retval "build Docker image ${IMGNAME}"
  rm -rf bashrc
}

build_docker_container() {
  # Build a container from the image
  echo -e "\nRemoving older container ${CONTNAME}..."
  if [ 1 -eq $(docker container ls -a | grep -w "${CONTNAME}$" | wc -l) ] ; then
    docker rm -f "$CONTNAME"
  fi

  echo -e "\nBuilding a container ${CONTNAME} from the image ${IMGNAME}..."
  docker create -it --name="$CONTNAME" \
    -v "$SCRIPTPATH":/"$REPONAME" \
    -v /tmp/.X11-unix:/tmp/.X11-unix \
    -v /etc/localtime:/etc/localtime:ro \
    -e DISPLAY="$DISPLAY" \
    -e QT_X11_NO_MITSHM=1 \
    -e XAUTHORITY \
    -e NVIDIA_DRIVER_CAPABILITIES=all \
    --net=host \
    --ipc=host \
    --gpus all \
    $PORT_FORWARD_DOCKER_CMD \
    --privileged=true \
    --cap-add=CAP_SYS_ADMIN \
    "$IMGNAME" /bin/bash
  test_retval "create Docker container"
}

start_docker_container() {
  docker start -ai "$CONTNAME"

  if [ 0 -eq $(docker container ls -a | grep -w "${CONTNAME}$" | wc -l) ] ; then
    echo -e "\nFailed to start/attach Docker container... Exiting...\n"
    exit 1
  fi
}

print_exit_command() {
  local command_to_start_container="docker start -ai ${CONTNAME}"
  # Echo command to start container
  echo -e "\n"
  echo -e "################################################################################\n"
  echo -e "\tCommand to start Docker container:\n\t\t${command_to_start_container}\n"
  echo -e "################################################################################\n"
}


############################################################
# Section 2: Call Helper Functions                         #
############################################################
# Check if _arg_container_name can be found
check_cont_name

# Print the setup info
print_setup_info
# Print usage of the script
print_help

echo -e ".......... Set up will start in 5 seconds .........."
sleep 5

if [ "$_arg_bashrc_only" = "on" ] ; then
  create_custom_bashrc
  exit 0
fi

remove_prev_docker_image

### Build docker image ###
update_port_forward

if [ "$_arg_build_local" = "on" ] ; then
  create_custom_bashrc
  build_docker_image
else
  docker pull "$IMGNAME"
fi

build_docker_container
if [ "$_arg_build_only" = "off" ] ; then
  start_docker_container
fi

# When exit from docker container
print_exit_command


# ^^^  TERMINATE YOUR CODE BEFORE THE BOTTOM ARGBASH MARKER  ^^^

# ] <-- needed because of Argbash