run_nim.sh

#!/usr/bin/env bash
#
# run_nim.sh — ergonomic launcher for NVIDIA NIM LLM containers
#
# Usage examples are at the bottom of --help.
#
set -euo pipefail

# ---------- defaults ----------
HOST_PORT=8000                 # what you hit from your browser/curl
CONTAINER_NAME="nim-${HOST_PORT}"
IMG_NAME="nvcr.io/nim/meta/llama-3.1-8b-instruct-dgx-spark:latest"
CONTAINER_PORT=""              # defaults to HOST_PORT unless set
LOCAL_NIM_CACHE="${LOCAL_NIM_CACHE:-$HOME/.cache/nim}"
SHM_SIZE="16GB"
USER_FLAG="-u $(id -u)"        # match your UID to avoid root-owned files
EXTRA_ENVS=()

# ---------- helpers ----------
die() { echo "Error: $*" >&2; exit 1; }

print_help() {
  cat <<'EOF'
run_nim.sh — Run a NIM LLM container with sane defaults.

USAGE:
  ./run_nim.sh [options] [-- <extra docker args>]

OPTIONS:
  -i, --image IMAGE         Container image to run
                            (default: nvcr.io/nim/meta/llama-3.1-8b-instruct-dgx-spark:latest)
  -p, --port PORT           Host port to expose for the service (default: 8000)
      --container-port P    Internal container port (default: same as --port)
  -n, --name NAME           Container name (default: nim-llm)
  -c, --cache PATH          Local NIM cache dir to mount at /opt/nim/.cache
                            (default: ~/.cache/nim or \$LOCAL_NIM_CACHE)
  -s, --shm SIZE            --shm-size (default: 16GB)
  -e, --env KEY=VALUE       Additional env var(s) to pass (repeatable)
      --as-root             Run as root in the container (omit user mapping)
  -h, --help                Show this help and exit

NOTES:
- You must have Docker with GPU support and a valid \$NGC_API_KEY to pull NIM images.
- By default, the script maps HOST:CONTAINER ports as PORT:PORT and exports LLM_PORT=PORT
  into the container. If your image listens on a fixed port (often 8000), use --container-port 8000.
- The NIM cache is mounted at /opt/nim/.cache to reuse model artifacts across runs.

EXAMPLES:
  # Use Meta Llama 3.1 8B, port 8000 (defaults)
  ./run_nim.sh

  # Switch to Qwen 32B and expose on host :8080 (container also 8080)
  ./run_nim.sh -i nvcr.io/nim/qwen/qwen3-32b-dgx-spark:latest -p 8080

  # Nemotron Nano 9B, host :9000 mapped to container :8000
  ./run_nim.sh -i nvcr.io/nim/nvidia/nvidia-nemotron-nano-9b-v2-dgx-spark:1.0.0-variant \
               -p 9000 --container-port 8000

  # Pass extra docker args after --
  ./run_nim.sh -p 8000 -- --pull=always

EOF
}

# ---------- arg parsing ----------
ARGS=()
while [[ $# -gt 0 ]]; do
  case "$1" in
    -i|--image)          IMG_NAME="${2:-}"; shift 2 ;;
    -p|--port)           HOST_PORT="${2:-}"; shift 2 ;;
        --container-port) CONTAINER_PORT="${2:-}"; shift 2 ;;
    -n|--name)           CONTAINER_NAME="${2:-}"; shift 2 ;;
    -c|--cache)          LOCAL_NIM_CACHE="${2:-}"; shift 2 ;;
    -s|--shm)            SHM_SIZE="${2:-}"; shift 2 ;;
    -e|--env)            EXTRA_ENVS+=("$2"); shift 2 ;;
        --as-root)       USER_FLAG=""; shift ;;
    -h|--help)           print_help; exit 0 ;;
    --)                  shift; ARGS+=("$@"); break ;;
    *)                   die "Unknown argument: $1 (see --help)";;
  esac
done

# ---------- validations ----------
command -v docker >/dev/null 2>&1 || die "Docker not found in PATH."
[[ -n "${IMG_NAME}" ]]    || die "--image is empty."
[[ -n "${HOST_PORT}" ]]   || die "--port is empty."
if [[ -z "${CONTAINER_PORT}" ]]; then CONTAINER_PORT="${HOST_PORT}"; fi

# NGC key check (required for most nvcr.io pulls)
if [[ -z "${NGC_API_KEY:-}" ]]; then
  cat <<'EOF'

NGC_API_KEY is not set. This must be set in your environment or manually set by you:

export NGC_API_KEY="<your key here>"

EOF

   die "Set NGC_API_KEY and retry"
fi

mkdir -p "${LOCAL_NIM_CACHE}"

# ---------- assemble docker args ----------
DOCKER_ENVS=""
for kv in "${EXTRA_ENVS[@]:-}"; do
  if [[ ! -z "$kv" ]]; then
    DOCKER_ENVS+=(-e "$kv")
  fi
done

# ---------- show plan ----------
echo "Launching:"
echo "  Image           : ${IMG_NAME}"
echo "  Name            : ${CONTAINER_NAME}"
echo "  Host port       : ${HOST_PORT}"
echo "  Container port  : ${CONTAINER_PORT}"
echo "  NIM cache       : ${LOCAL_NIM_CACHE}"
echo "  SHM size        : ${SHM_SIZE}"
echo


# ---------- run ----------
exec docker run -it --rm \
  --name="${CONTAINER_NAME}" \
  --gpus all \
  ${USER_FLAG} \
  --shm-size="${SHM_SIZE}" \
  -v "${LOCAL_NIM_CACHE}:/opt/nim/.cache" \
  -p "${HOST_PORT}:${CONTAINER_PORT}" \
  ${DOCKER_ENVS[@]} ${ARGS[@]} ${IMG_NAME}