xsuite-containers/run_jupyter.sh at main · ekatralis/xsuite-containers · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#!/usr/bin/env bash
set -euo pipefail

IMAGE="ghcr.io/ekatralis/xsuite-containers:latest"
PORT="${PORT:-8888}"
ENGINE="${ENGINE:-}"
JUPYTER_TOKEN="${JUPYTER_TOKEN:-"xsuite"}"
GPU="${GPU:-}"

usage() {
  echo "Usage: $0 /PATH/TO/NOTEBOOKS"
  echo "Env: PORT=8888 (optional), ENGINE=docker|podman (optional, auto-detected), JUPYTER_TOKEN=auto|<token> (optional, default 'xsuite'), GPU=NVIDIA|AMD (optional, default empty)"
}

NOTEBOOKS_DIR="${1:-}"
if [[ -z "${NOTEBOOKS_DIR}" ]]; then
  usage
  exit 1
fi

if [[ ! -d "${NOTEBOOKS_DIR}" ]]; then
  echo "Error: '${NOTEBOOKS_DIR}' is not a directory."
  exit 1
fi

# Promote to absolute path for best compatibility
NOTEBOOKS_DIR="$(cd "${NOTEBOOKS_DIR}" && pwd -P)"

if [[ -z "${ENGINE}" ]]; then
    echo "Automatically selecting container engine..."
    # Prefer podman if present, otherwise docker
    if command -v podman >/dev/null 2>&1; then
    ENGINE="podman"
    elif command -v docker >/dev/null 2>&1; then
    ENGINE="docker"
    else
    echo "Error: neither 'podman' nor 'docker' found in PATH."
    exit 1
    fi
fi

echo "Using container engine: ${ENGINE}"
ENGINE_ARGS=("-e" "HOME=/home/xsuiteuser/")
if [[ "${GPU}" == "NVIDIA" ]]; then
  IMAGE+="-cuda"
  echo "GPU mode enabled: NVIDIA; using image ${IMAGE}"
  if [[ "${ENGINE}" == "docker" ]]; then
    ENGINE_ARGS+=( "--gpus" "all" "--runtime" "nvidia" )
  else
    ENGINE_ARGS+=( "--device" "nvidia.com/gpu=all" )
  fi
elif [[ "${GPU}" == "AMD" ]]; then
  echo "ROCm container not implemented yet"
  exit 1
elif [[ -n "${GPU}" ]]; then
  echo "Unsupported GPU type: ${GPU}"
  exit 1
fi

echo "Pulling image: ${IMAGE}"
"${ENGINE}" pull "${IMAGE}"

JUPYTER_CMD="jupyter lab --ip=0.0.0.0 --no-browser --notebook-dir=/workspace"

if [[ "${JUPYTER_TOKEN}" != "auto" ]]; then
    JUPYTER_CMD+=" --ServerApp.token='${JUPYTER_TOKEN}'"
    echo "Starting Jupyter Lab on http://localhost:${PORT}/lab?token=${JUPYTER_TOKEN}"
else
    echo "Starting Jupyter Lab on http://localhost:${PORT}"
fi

# Detect OS
OS="$(uname -s)"

# Build engine args depending on OS + podman mode rules
[[ -e /sys/fs/selinux/enforce ]] && MOUNT_SUFFIX=":Z" || MOUNT_SUFFIX=""
VOLUME_ARG=( -v "${NOTEBOOKS_DIR}:/workspace${MOUNT_SUFFIX}" )
PORT_ARG=( -p "${PORT}:8888" )

if [[ "${ENGINE}" == "docker" ]]; then
  # Docker is the same as podman rootful mode but requires specifying the home directory
  ENGINE_ARGS+=(  "--user" "$(id -u):$(id -g)" "--group-add" "2020")
elif [[ "${OS}" == "Darwin" ]]; then
  # macOS podman runs in a VM (rootful), use macOS-specific user/group setup
  ENGINE_ARGS+=( "--user" "$(id -u):$(id -g)" "--group-add" "2020" )
else
  # Linux + podman: choose rootless vs rootful
  ROOTLESS="$("${ENGINE}" info --format '{{.Host.Security.Rootless}}' 2>/dev/null || echo "false")"
  if [[ "${ROOTLESS}" == "true" ]]; then
    ENGINE_ARGS+=( "--userns=keep-id" "--user" "$(id -u):$(id -g)" "--group-add" "2020" )
  else
    ENGINE_ARGS+=( "--user" "$(id -u):$(id -g)" "--group-add" "2020")
  fi
fi

echo "Mounting notebooks: ${NOTEBOOKS_DIR} -> /workspace"

exec "${ENGINE}" run --rm -it \
  "${ENGINE_ARGS[@]}" \
  "${PORT_ARG[@]}" \
  "${VOLUME_ARG[@]}" \
  "${IMAGE}" \
  bash -lc "${JUPYTER_CMD}"