docker-launch

#!/bin/bash

# Get the directory of the script
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PWD=`pwd`

# Check if the argument (tag) is provided
if [ -z "$1" ]; then
  echo "Usage: $0 <tag> [--rebuild]"
  echo
  echo "This script checks if a Docker image with the given tag exists."
  echo "If the image exists, it runs bash on the container."
  echo "If the image does not exist, it builds the image from <tag>.dockerfile and then runs bash on the container."
  echo "Optional [--rebuild] flag forcefully rebuilds the image from <tag>.dockerfile even if the image already exists."
  exit 1
fi

TAG=$1

REBUILD=false

# Check for the --rebuild option
if [ "$2" == "--rebuild" ]; then
  REBUILD=true
fi

# Function to build the Docker image
build_image() {
  echo "Building Docker image from $TAG.dockerfile..."
  if [ -f "$SCRIPT_DIR/dockerfiles/$TAG.dockerfile" ]; then
    cd $SCRIPT_DIR/dockerfiles
    docker build -t "docker_launch_$TAG" -f "$SCRIPT_DIR/dockerfiles/$TAG.dockerfile" .
    cd $PWD
    if [ $? -ne 0 ]; then
      echo "Failed to build Docker image with tag '$TAG'."
      exit 1
    fi
  else
    echo "Dockerfile '$TAG.dockerfile' does not exist."
    exit 1
  fi
}

# Function to check if the Dockerfile uses an NVIDIA base image
uses_nvidia_runtime() {
  if [ -f "$SCRIPT_DIR/dockerfiles/$TAG.dockerfile" ]; then
    if grep -qi '^FROM.*nvidia' "$SCRIPT_DIR/dockerfiles/$TAG.dockerfile"; then
      return 0
    else
      return 1
    fi
  else
    echo "Dockerfile '$TAG.dockerfile' does not exist."
    exit 1
  fi
}

# Check if the Docker image with the given tag exists or if rebuild is requested
if $REBUILD || ! docker image inspect "docker_launch_$TAG" > /dev/null 2>&1; then
  build_image
fi

if docker ps | grep -q docker_launch_$TAG
then
  docker exec -it docker_launch_$TAG /bin/bash
else
  echo "(starting)"
  docker rm docker_launch_$TAG 2>/dev/null

  RUNTIME_FLAG=""
  if uses_nvidia_runtime; then
    RUNTIME_FLAG="--runtime nvidia"
  fi

  DOCKER_RUN_CMD="docker run \
      -e DISPLAY \
      -e DBUS_SESSION_BUS_ADDRESS \
      -v /tmp/.X11-unix:/tmp/.X11-unix \
      -v /home/dheera:/home/dheera \
      -v $XAUTHORITY:/root/.Xauthority \
      --net=host \
      --name docker_launch_$TAG \
      --privileged \
      $RUNTIME_FLAG \
      -it docker_launch_$TAG /bin/bash"

  eval $DOCKER_RUN_CMD
fi