Docker image for MadGraph5_aMC@NLO + MadMiner + Jupyter. #docker #python #jupyter

!my-madminer-jupyter-env.md

my-madminer-jupyter-env

• Docker is required.
• NVIDIA Container Toolkit is required to use GPU. See this.

Usage:

git clone https://gist.github.com/ac4a004de14b34e4e746e6412cf70249.git my-madminer-jupyter-env
cd my-madminer-jupyter-env

./run.sh  # starts a Bash shell session; you can run python, root, mg5_aMC, etc.

./run.sh lab  # launches JupyterLab

The run.sh script builds a Docker image and then runs it within the data directory located in the current directory, creating it if it does not exist.

Tutorial

Clone the repository in the data directory (either on the host or inside the container):

git clone --depth=1 https://github.com/madminer-tool/madminer.git

Then, open and run notebook files in madminer/examples/tutorial_particle_physics.

Running on remote host

You need something like:

ssh -L 8888:localhost:8888 example.com

Development

pre-commit install

.dockerignore

*
!docker-entrypoint.sh

.editorconfig

root = true

[*]
charset = utf-8
end_of_line = lf
insert_final_newline = true
trim_trailing_whitespace = true

[*.md]
trim_trailing_whitespace = false

[*.{sh,yml}]
indent_size = 2
indent_style = space

[{.pre-commit-config.yaml,Dockerfile}]
indent_size = 4
indent_style = space

.gitlint

[general]
ignore=body-min-length,body-is-missing
contrib=contrib-title-conventional-commits

[title-max-length]
line-length=80

[body-max-line-length]
line-length=80

[ignore-by-title]
regex=^Merge(.*)into(.*)

[contrib-title-conventional-commits]
types = build,chore,ci,docs,feat,fix,perf,refactor,revert,style,test

.pre-commit-config.yaml

default_install_hook_types: [pre-commit, commit-msg]

repos:
-   repo: https://github.com/pre-commit/pre-commit-hooks
    rev: v5.0.0
    hooks:
    -   id: check-added-large-files
    -   id: check-case-conflict
    -   id: check-executables-have-shebangs
    -   id: check-json
    -   id: check-merge-conflict
    -   id: check-shebang-scripts-are-executable
    -   id: check-symlinks
    -   id: check-toml
    -   id: check-yaml
    -   id: debug-statements
    -   id: destroyed-symlinks
    -   id: end-of-file-fixer
    -   id: fix-byte-order-marker
    -   id: mixed-line-ending
    -   id: trailing-whitespace

-   repo: https://github.com/hadolint/hadolint
    rev: v2.12.0
    hooks:
    -   id: hadolint-docker

-   repo: https://github.com/shellcheck-py/shellcheck-py
    rev: v0.10.0.1
    hooks:
    -   id: shellcheck
        args: [-x]

-   repo: https://github.com/jorisroovers/gitlint
    rev: v0.19.1
    hooks:
    -   id: gitlint

docker-entrypoint.sh

# shellcheck shell=bash

export USER=user
export HOME=/home/$USER

uid=$(stat -c "%u" .)
gid=$(stat -c "%g" .)

if [ "$uid" -ne 0 ]; then
    if [ "$(id -g $USER)" -ne "$gid" ]; then
        getent group "$gid" >/dev/null 2>&1 || groupmod -g "$gid" $USER
        chgrp -R "$gid" $HOME
    fi
    if [ "$(id -u $USER)" -ne "$uid" ]; then
        usermod -u "$uid" $USER
    fi
fi

exec setpriv --reuid=$USER --regid=$USER --init-groups "$@"

Dockerfile

ARG TARGET_ARCH=amd64
ARG ROOT_VERSION=6.32.02
ARG ROOT_PRIV=root

################################################################################

FROM rootproject/root:$ROOT_VERSION-ubuntu22.04 AS base-amd64

################################################################################

FROM ubuntu:22.04 AS base-arm64v8

RUN echo "This Dockerfile does not yet support arm64v8" >&2; exit 1

################################################################################

# hadolint ignore=DL3006
FROM base-$TARGET_ARCH AS base

################################################################################

FROM base AS mg_builder

SHELL ["/bin/bash", "-o", "pipefail", "-c"]

# Running mg5_aMC requires: python3-six
# Building pythia8 requires: rsync
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    python3-six=1.16.* \
    rsync=3.2.* \
    && rm -rf /var/lib/apt/lists/*

ARG MG5_URL=https://launchpad.net/mg5amcnlo/3.0/3.5.x/+download/MG5_aMC_v3.5.6.tar.gz

RUN curl -sSL $MG5_URL | tar -xz -C /opt && mv /opt/MG5_* MG5_aMC
RUN echo "n" | python3 /opt/MG5_aMC/bin/mg5_aMC
RUN echo "set auto_update 0; save options auto_update" | python3 /opt/MG5_aMC/bin/mg5_aMC

# Install Pythia8 and Delphes
RUN echo "install pythia8" | python3 /opt/MG5_aMC/bin/mg5_aMC
RUN grep -q "^pythia8_path" /opt/MG5_aMC/input/mg5_configuration.txt
RUN echo "install Delphes" | python3 /opt/MG5_aMC/bin/mg5_aMC
# RUN grep -q "^delphes_path" /opt/MG5_aMC/input/mg5_configuration.txt

# Turn ON Python2 -> Python3 models conversion
RUN echo "set auto_convert_model T" | python3 /opt/MG5_aMC/bin/mg5_aMC
RUN echo "import model EWdim6-full" | python3 /opt/MG5_aMC/bin/mg5_aMC

################################################################################

FROM base AS py_builder

ARG TORCH_WHL=cpu

RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    python3-pip=22.* \
    && rm -rf /var/lib/apt/lists/*

RUN python3 -m pip install --no-cache-dir \
    torch==1.13.1+$TORCH_WHL \
    torchvision==0.14.1+$TORCH_WHL \
    --extra-index-url https://download.pytorch.org/whl/$TORCH_WHL

ARG MADMINER_VERSION=0.9.6

# Running mg5_aMC requires: six
RUN python3 -m pip install --no-cache-dir \
    bqplot==0.12.43 \
    madminer==$MADMINER_VERSION \
    matplotlib==3.2.2 \
    pandas==2.0.3 \
    scikit-learn==1.5.1 \
    scipy==1.10.1 \
    seaborn==0.12.2 \
    six==1.16.0 \
    git+https://github.com/tueda/[email protected]

# hadolint ignore=DL3059
RUN python3 -m pip install --no-cache-dir \
    black==24.8.0 \
    isort==5.13.2 \
    jupyterlab-code-formatter==3.0.0 \
    jupyterlab-lsp==5.1.0 \
    jupyterlab==4.2.4 \
    lckr-jupyterlab-variableinspector==3.2.1 \
    python-lsp-server==1.11.0

################################################################################

FROM base AS runner-base

COPY --from=mg_builder /opt/MG5_aMC /opt/MG5_aMC
COPY --from=py_builder /usr/local/bin /usr/local/bin
COPY --from=py_builder /usr/local/lib/python3.10 /usr/local/lib/python3.10
COPY --from=py_builder /usr/local/etc/jupyter /usr/local/etc/jupyter
COPY --from=py_builder /usr/local/share/jupyter /usr/local/share/jupyter

RUN jupyter labextension disable "@jupyterlab/apputils-extension:announcements"

ENV MG_FOLDER_PATH "/opt/MG5_aMC"

ENV PATH $MG_FOLDER_PATH/bin:$MG_FOLDER_PATH/HEPTools/lhapdf6_py3/bin:$PATH
ENV LD_LIBRARY_PATH $MG_FOLDER_PATH/HEPTools/lhapdf6_py3/lib:$ROOTSYS/lib
ENV PYTHONPATH $MG_FOLDER_PATH/HEPTools/lhapdf6_py3/local/lib/python3.10/dist-packages:$ROOTSYS/lib
ENV ROOT_INCLUDE_PATH $MG_FOLDER_PATH/Delphes/external

ARG EXTRA_APT_PACKAGES=

# hadolint ignore=DL3008
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    sudo=1.9.* \
    $EXTRA_APT_PACKAGES \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /data
CMD ["/bin/bash"]

################################################################################

FROM runner-base AS runner-rootless

RUN ln -s /root/.cache /home/user/.cache \
    && ln -s /root/.config /home/user/.config \
    && ln -s /root/.ipython /home/user/.ipython \
    && ln -s /root/.jupyter /home/user/.jupyter \
    && ln -s /root/.local /home/user/.local

################################################################################

FROM runner-base AS runner-root

RUN useradd -m user \
    && echo 'user ALL=(ALL) NOPASSWD:ALL' >>/etc/sudoers \
    && usermod -aG sudo user \
    && ln -s /data/.cache /home/user/.cache \
    && ln -s /data/.config /home/user/.config \
    && ln -s /data/.ipython /home/user/.ipython \
    && ln -s /data/.jupyter /home/user/.jupyter \
    && ln -s /data/.local /home/user/.local \
    && chown -h user:user /home/user/.cache /home/user/.config /home/user/.ipython /home/user/.jupyter /home/user/.local

COPY docker-entrypoint.sh /
ENTRYPOINT ["/bin/bash", "/docker-entrypoint.sh"]

################################################################################

# hadolint ignore=DL3006
FROM runner-$ROOT_PRIV AS runner

run.sh

#!/bin/bash
set -eu
set -o pipefail

ROOT_VERSION=6.32.02
MG5_VERSION=3.5.6
MADMINER_VERSION=0.9.6
EXTRA_APT_PACKAGES=''
DATA_DIR=data

# get_mg5_url <VERSION> prints the download URL for the specified version of MG5.
get_mg5_url() {
  local version=$1

  local series
  case $version in
    2.?*.?*)
      series=lts
      ;;
    3.?*.?*)
      series=3.0
      ;;
    *)
      echo "error: invalid mg5 version: $version">&2
      exit 1
      ;;
  esac

  local release
  release=${version%.*}.x

  echo "https://launchpad.net/mg5amcnlo/$series/$release/+download/MG5_aMC_v$version.tar.gz"
}

# find_unused_port <START_PORT> <END_PORT> prints an unused port within the given range.
find_unused_port() {
  local start_port=$1
  local end_port=$2

  if command -v nc &>/dev/null; then
    for port in $(seq "$start_port" "$end_port"); do
      if ! nc -z 127.0.0.1 "$port" &>/dev/null; then
        echo "$port"
        return 0
      fi
    done
  elif command -v telnet &>/dev/null; then
    for port in $(seq "$start_port" "$end_port"); do
      if ! printf "\035" | telnet 127.0.0.1 "$port" &>/dev/null; then
        echo "$port"
        return 0
      fi
    done
  fi

  echo "error: can't find unused port in the range $start_port-$end_port">&2
  exit 1
}

# Entry point.
src_dir=$(dirname "$0")
data_dir="$(pwd)/$DATA_DIR"

# Check for the machine architecture.
target_arch=amd64
case $(uname -m) in
  aarch64|arm64)
    target_arch=arm64v8
    ;;
esac

# Check for GPUs.
torch_whl=cpu
if command -v nvidia-smi &>/dev/null && command -v nvidia-container-toolkit &>/dev/null; then
  # torch 1.13.1: cu116 or cu117
  torch_whl=cu117
fi

# Construct the MG5 URL.
mg5_url=$(get_mg5_url $MG5_VERSION)

# Ensure the data directory exists, which is shared between the host and the container.
mkdir -p "$data_dir"
[ -f "$data_dir/.gitignore" ] || echo '*' >"$data_dir/.gitignore"
mkdir -p "$data_dir/.cache" "$data_dir/.config" "$data_dir/.ipython" "$data_dir/.jupyter" "$data_dir/.local"
[ -L "$data_dir/.lsp_symlink" ] || ln -s / "$data_dir/.lsp_symlink"  # for jupyterlab-lsp

# Check for root privileges in containers. The result is cached.
if [ "$(id -u)" -eq 0 ]; then
  # already root
  root_priv=rootless
else
  if [ -f "$data_dir/.docker-root-priv" ]; then
    # cached
    root_priv=$(cat "$data_dir/.docker-root-priv")
  else
    docker build "$src_dir" \
      --target base \
      --build-arg TARGET_ARCH=$target_arch \
      --build-arg ROOT_VERSION=$ROOT_VERSION \
      --build-arg MG5_URL="$mg5_url" \
      --build-arg MADMINER_VERSION=$MADMINER_VERSION \
      --build-arg EXTRA_APT_PACKAGES="$EXTRA_APT_PACKAGES" \
      --build-arg TORCH_WHL=$torch_whl \
      -t my-madminer-jupyter-env-base
    docker run --rm \
      -v "$data_dir:/data" \
      -w /data \
      my-madminer-jupyter-env-base \
      touch .docker-root-priv-test
    if [ "$(stat -c %u "$data_dir/.docker-root-priv-test")" -eq 0 ]; then
      root_priv=root
    else
      root_priv=rootless
    fi
    echo $root_priv >"$data_dir/.docker-root-priv"
    docker run --rm \
      -v "$data_dir:/data" \
      -w /data \
      my-madminer-jupyter-env-base \
      rm .docker-root-priv-test
  fi
fi

# Build the image.
docker build "$src_dir" \
  --build-arg TARGET_ARCH=$target_arch \
  --build-arg ROOT_VERSION=$ROOT_VERSION \
  --build-arg MG5_URL="$mg5_url" \
  --build-arg MADMINER_VERSION=$MADMINER_VERSION \
  --build-arg EXTRA_APT_PACKAGES="$EXTRA_APT_PACKAGES" \
  --build-arg TORCH_WHL=$torch_whl \
  --build-arg ROOT_PRIV="$root_priv" \
  -t my-madminer-jupyter-env

# Run the container.
run_cmd=bash
extra_run_opts=
if [[ $torch_whl != cpu ]]; then
  extra_run_opts="$extra_run_opts --gpus all"
fi
if [ $# -ge 1 ]; then
  case $1 in
    bash)
      run_cmd="bash"
      ;;
    lab)
      port=$(find_unused_port 8888 9999)
      run_cmd="jupyter $1 --allow-root --no-browser --ip 0.0.0.0 --port $port --ContentsManager.allow_hidden=True"
      extra_run_opts="$extra_run_opts -p $port:$port"
      ;;
    *)
      echo "error: invalid command $1" >&2
      exit 1
      ;;
  esac
fi
# shellcheck disable=SC2086
docker run -it --rm \
  -v "$data_dir:/data" \
  -w /data \
  $extra_run_opts \
  my-madminer-jupyter-env \
  $run_cmd