feat: Add GPU enabled Dockerfile (#748)

* Add NVIDIA GPU enabled Dockerfile for the machine learning framework backends
* Add .dockerignore to reduce build context size

Author: matthewfeickert

.dockerignore

@@ -0,0 +1,12 @@
+# Ignore everything
+**
+
+# Except the following
+!/.git/**
+!/docker/**
+!/src/**
+!/LICENSE
+!/pyproject.toml
+!/README.md
+!/setup.cfg
+!/setup.py

docker/gpu/Dockerfile

@@ -0,0 +1,33 @@
+FROM nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04 as base
+
+FROM base as builder
+# hadolint ignore=DL3015
+RUN apt-get update -y && \
+    apt-get install -y \
+      git \
+      python3 \
+      python3-pip && \
+    apt-get -y autoclean && \
+    apt-get -y autoremove && \
+    rm -rf /var/lib/apt-get/lists/*
+COPY . /code
+COPY ./docker/gpu/install_backend.sh /code/install_backend.sh
+WORKDIR /code
+ARG BACKEND=tensorflow
+RUN python3 -m pip install --upgrade --no-cache-dir pip setuptools wheel && \
+    /bin/bash install_backend.sh ${BACKEND} && \
+    python3 -m pip list
+
+FROM base
+# Use C.UTF-8 locale to avoid issues with ASCII encoding
+ENV LC_ALL=C.UTF-8
+ENV LANG=C.UTF-8
+COPY --from=builder /lib/x86_64-linux-gnu /lib/x86_64-linux-gnu
+COPY --from=builder /usr/local /usr/local
+COPY --from=builder /usr/bin/python3 /usr/bin/python3
+COPY --from=builder /usr/bin/python3.6 /usr/bin/python3.6
+COPY --from=builder /usr/bin/pip3 /usr/bin/pip3
+COPY --from=builder /usr/lib/python3 /usr/lib/python3
+COPY --from=builder /usr/lib/python3.6 /usr/lib/python3.6
+COPY --from=builder /usr/lib/x86_64-linux-gnu /usr/lib/x86_64-linux-gnu
+ENTRYPOINT ["/usr/local/bin/pyhf"]

docker/gpu/install_backend.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+set -e
+
+function get_JAXLIB_GPU_WHEEL {
+  # c.f. https://github.com/google/jax#pip-installation
+  local PYTHON_VERSION # alternatives: cp35, cp36, cp37, cp38
+  PYTHON_VERSION="cp"$(python3 --version | awk '{print $NF}' | awk '{split($0, rel, "."); print rel[1]rel[2]}')
+  local CUDA_VERSION # alternatives: cuda90, cuda92, cuda100, cuda101
+  CUDA_VERSION="cuda"$(< /usr/local/cuda/version.txt awk '{print $NF}' | awk '{split($0, rel, "."); print rel[1]rel[2]}')
+  local PLATFORM=linux_x86_64
+  local JAXLIB_VERSION=0.1.37
+  local BASE_URL="https://storage.googleapis.com/jax-releases"
+  local JAXLIB_GPU_WHEEL="${BASE_URL}/${CUDA_VERSION}/jaxlib-${JAXLIB_VERSION}-${PYTHON_VERSION}-none-${PLATFORM}.whl"
+  echo "${JAXLIB_GPU_WHEEL}"
+}
+
+function install_backend() {
+  # 1: the backend option name in setup.py
+  local backend="${1}"
+  if [[ "${backend}" == "tensorflow" ]]; then
+    # shellcheck disable=SC2102
+    python3 -m pip install --no-cache-dir .[xmlio,tensorflow]
+  elif [[ "${backend}" == "torch" ]]; then
+    # shellcheck disable=SC2102
+    python3 -m pip install --no-cache-dir .[xmlio,torch]
+  elif [[ "${backend}" == "jax" ]]; then
+    python3 -m pip install --no-cache-dir .[xmlio]
+    python3 -m pip install --no-cache-dir "$(get_JAXLIB_GPU_WHEEL)"
+    python3 -m pip install --no-cache-dir jax
+  fi
+}
+
+function main() {
+  # 1: the backend option name in setup.py
+  local BACKEND="${1}"
+  install_backend "${BACKEND}"
+}
+
+main "$@" || exit 1