docker_iblgf.sh

#!/usr/bin/env bash
# docker_iblgf.sh — enter IBLGF Docker env with optional CPU limit and Python
set -euo pipefail

CPU_BASE_IMAGE="ccardina/iblgf:cpu"

GPU_BASE_IMAGE="ccardina/iblgf:gpu"

USE_GPU=0
FORCE_PULL=0

# Where the repo lives
CONTAINER_ROOT="/workspace2"
CONTAINER_REPO_DIR="$CONTAINER_ROOT/IBLGF-AMR"

detect_host_timezone() {
  if [[ -n "${TZ:-}" ]]; then
    echo "$TZ"
    return 0
  fi

  if command -v timedatectl >/dev/null 2>&1; then
    timedatectl show --value --property=Timezone 2>/dev/null && return 0
  fi

  if command -v systemsetup >/dev/null 2>&1; then
    systemsetup -gettimezone 2>/dev/null | awk -F': ' 'NF > 1 { print $2; exit }' && return 0
  fi

  if [[ -L /etc/localtime ]]; then
    local tz_path
    tz_path="$(readlink /etc/localtime)"
    if [[ "$tz_path" == *"/zoneinfo/"* ]]; then
      echo "${tz_path##*/zoneinfo/}"
      return 0
    fi
  fi

  return 1
}

host_utc_offset() {
  date +"%z"
}

usage() {
  cat <<EOF
Usage:
  ./docker_iblgf.sh
  ./docker_iblgf.sh -c N
  ./docker_iblgf.sh -g
  ./docker_iblgf.sh -g -c N
  ./docker_iblgf.sh --pull

Options:
  -c N    Limit Docker container to N CPU cores
  -g      Use GPU image 
  --pull  Pull the latest version of the selected image before starting

Examples:
  ./docker_iblgf.sh
  ./docker_iblgf.sh -c 4
  ./docker_iblgf.sh -g
  ./docker_iblgf.sh -g -c 4
  ./docker_iblgf.sh --pull
EOF
}

CPUS=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --)
      shift
      break   # <-- stop option parsing here
      ;;
    -c|--cpus)
      shift
      CPUS="${1:-}"
      [[ -n "$CPUS" ]] || { echo "Missing value for -c"; exit 1; }
      # basic numeric validation
      [[ "$CPUS" =~ ^[0-9]+$ ]] || { echo "CPU count must be an integer"; exit 1; }
      [[ "$CPUS" -ge 1 ]] || { echo "CPU count must be >= 1"; exit 1; }
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    -g|--gpu)
      USE_GPU=1
      shift
      ;;
    -p|--pull)
      FORCE_PULL=1
      shift
      ;;
    *)
      echo "Unknown option: $1"
      usage
      exit 1
      ;;
  esac
done

if [[ "$USE_GPU" -eq 1 ]]; then
  BASE_IMAGE="$GPU_BASE_IMAGE"
else
  BASE_IMAGE="$CPU_BASE_IMAGE"
fi

echo "==> Base image: $BASE_IMAGE"

# Pull base image if requested or missing
if [[ "$FORCE_PULL" -eq 1 ]]; then
  echo "==> Pulling latest base image..."
  docker pull "$BASE_IMAGE"
elif ! docker image inspect "$BASE_IMAGE" >/dev/null 2>&1; then
  echo "==> Pulling base image..."
  docker pull "$BASE_IMAGE"
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Case A: script sits in repo root and repo has CMakeLists.txt
if [[ -f "$SCRIPT_DIR/CMakeLists.txt" ]]; then
  HOST_REPO_DIR="$SCRIPT_DIR"
# Case B: script is in a subdir of the repo 
  SEARCH_DIR="$SCRIPT_DIR"
  HOST_REPO_DIR=""
  while [[ "$SEARCH_DIR" != "/" ]]; do
    if [[ -f "$SEARCH_DIR/CMakeLists.txt" ]]; then
      HOST_REPO_DIR="$SEARCH_DIR"
      break
    fi
    SEARCH_DIR="$(dirname "$SEARCH_DIR")"
  done
  [[ -n "$HOST_REPO_DIR" ]] || { echo "ERROR: Could not locate repo root (CMakeLists.txt)"; exit 1; }
fi

echo "==> Host repo dir: $HOST_REPO_DIR"
echo "==> Container repo dir: $CONTAINER_REPO_DIR"

INTERACTIVE=1
if [[ $# -gt 0 ]]; then
  INTERACTIVE=0
fi

DOCKER_ARGS=(docker run --rm)

DOCKER_ARGS+=( --user "$(id -u):$(id -g)" )

if [[ "$INTERACTIVE" -eq 1 ]]; then
  DOCKER_ARGS+=( -it )
fi

# Mount repo and set working directory
DOCKER_ARGS+=(
  -v "$HOST_REPO_DIR:$CONTAINER_REPO_DIR"
  -w "$CONTAINER_REPO_DIR"
)

DOCKER_ARGS+=( -e LD_LIBRARY_PATH="/usr/local/lib:/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-}" )

if HOST_TZ="$(detect_host_timezone)"; then
  echo "==> Using timezone: $HOST_TZ"
  DOCKER_ARGS+=( -e TZ="$HOST_TZ" )
fi

HOST_UTC_OFFSET="$(host_utc_offset)"
echo "==> Using UTC offset: $HOST_UTC_OFFSET"
DOCKER_ARGS+=( -e IBLGF_HOST_UTC_OFFSET="$HOST_UTC_OFFSET" )

#check if env var CI is set to true or if GITHUB_ACTIONS is set to non-empty string, and if so, set OMPI_MCA_rmaps_base_oversubscribe=1
if [[ "${CI:-false}" == "true" ]]; then
  echo "==> Detected CI environment, setting OMPI_MCA_rmaps_base_oversubscribe=1"
  DOCKER_ARGS+=( -e OMPI_MCA_rmaps_base_oversubscribe=1 )
fi


if [[ "$USE_GPU" -eq 1 ]]; then
  DOCKER_ARGS+=(--gpus all)
fi

if [[ -n "$CPUS" ]]; then
  echo "==> Limiting container to $CPUS CPU(s)"
  DOCKER_ARGS+=(--cpuset-cpus="0-$((CPUS-1))")
fi

DOCKER_ARGS+=("$BASE_IMAGE")

if [[ $# -gt 0 ]]; then
  # Non-interactive
  DOCKER_ARGS+=("$@")
else
  # Interactive shell
  DOCKER_ARGS+=("bash")
fi

exec "${DOCKER_ARGS[@]}"