Feature/docker (#496)

* Refactor Dockerfile. * setup.py install * Multi-stage docker build. * Update docker readme. * Add omp deps. * Fix pip install warning. * Auto update version * Add lightning-kokkos stage. * Add LGPU stage. * Fix codefactor warnings. * Add LKokkos+CUDA stage. * Fix wheel-lightning-kokkos-cuda runtime env. * Add make docker & dockerhub readme. * Build kokkos-cuda from --branch fix-rtld-deepbind https://github.com/cz4rs/kokkos.git for bugfix, waiting kokkos-4.2 * Try fixing codefactor * Add plugin doc links. Add HIP image. Refactor makefile. * Auto update version * Update readme. * Update readme. * Update changelog. * wget --progress. * Update docker/README.md Co-authored-by: Amintor Dusko <[email protected]> * Auto update version * trigger CI * Add GCC_VERSION build arg. * Auto update version * Download lightning in base-build-python. Build LQubit without Kokkos. * Add PL versions docker. * Auto update version * Change CUDA/ROCM ARGs for installer urls. Clear *.deb after install. * Auto update version * Ship wheel-lightning-kokkos-rocm on top of rocm/dev-ubuntu-22.04. * Auto update version * Add rocm/dev-ubuntu-22.04 tag. --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: Amintor Dusko <[email protected]>
PennyLaneAI · Sep 22, 2023 · 8f744aa · 8f744aa
1 parent 17812fb
commit 8f744aa
Show file tree

Hide file tree

Showing 5 changed files with 362 additions and 27 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add `Dockerfile` in `docker` and `make docker` workflow in `Makefile`. The Docker images and documentation are available on [DockerHub](https://hub.docker.com/repository/docker/pennylaneai/pennylane).
+  [(#496)](https://github.com/PennyLaneAI/pennylane-lightning/pull/496)
+
 * Add mid-circuit state preparation operation tests.
   [(#495)](https://github.com/PennyLaneAI/pennylane-lightning/pull/495)
 

diff --git a/Makefile b/Makefile
@@ -45,6 +45,9 @@ help:
 	@echo "  check-tidy [verbose=1]   to build PennyLane-Lightning with ENABLE_CLANG_TIDY=ON (requires clang-tidy & CMake)"
 	@echo "                           use with 'verbose=1' for building with verbose flag"
 	@echo "  check-tidy [target=?]    to build a specific PennyLane-Lightning target with ENABLE_CLANG_TIDY=ON (requires clang-tidy & CMake)"
+	@echo "  docker-build [target=?]  to build a Docker image for a specific PennyLane-Lightning target"
+	@echo "  docker-push  [target=?]  to push a Docker image to the PennyLaneAI Docker Hub repo"
+	@echo "  docker-all  		      to build and push Docker images for all PennyLane-Lightning targets"
 
 .PHONY : clean
 clean:
@@ -125,3 +128,34 @@ docs:
 .PHONY : clean-docs
 clean-docs:
 	$(MAKE) -C doc clean
+
+.PHONY : docker
+ifdef target
+    TARGET := $(target)
+else
+    TARGET := lightning-qubit
+endif
+ifdef version
+    VERSION := $(version)
+else
+    VERSION := 0.32.0
+endif
+docker-build:
+	docker build -f docker/Dockerfile --tag=pennylaneai/pennylane:$(VERSION)-$(TARGET) --target wheel-$(TARGET) .
+docker-push:
+	docker push pennylaneai/pennylane:$(VERSION)-$(TARGET)
+docker-build-all:
+	$(MAKE) docker-build target=lightning-gpu
+	$(MAKE) docker-build target=lightning-kokkos-cuda
+	$(MAKE) docker-build target=lightning-kokkos-rocm
+	$(MAKE) docker-build target=lightning-kokkos-openmp
+	$(MAKE) docker-build target=lightning-qubit
+docker-push-all:
+	$(MAKE) docker-push target=lightning-gpu
+	$(MAKE) docker-push target=lightning-kokkos-cuda
+	$(MAKE) docker-push target=lightning-kokkos-rocm
+	$(MAKE) docker-push target=lightning-kokkos-openmp
+	$(MAKE) docker-push target=lightning-qubit
+docker-all:
+	$(MAKE) docker-build-all
+	$(MAKE) docker-push-all
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -1,4 +1,4 @@
-# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,43 +12,266 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-FROM ubuntu:focal AS base
+# Create basic runtime environment base on Ubuntu 22.04 (jammy)
+# Define build arguments
+# Create and activate runtime virtual environment
+FROM ubuntu:jammy AS base-runtime
+ARG AMD_ARCH=AMD_GFX90A
+ARG CUDA_ARCH=AMPERE80
+ARG CUDA_INSTALLER=https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run
+ARG GCC_VERSION=11
+ARG KOKKOS_VERSION=4.1.00
+ARG LIGHTNING_VERSION=0.32.0
+ARG ROCM_INSTALLER=https://repo.radeon.com/amdgpu-install/5.7/ubuntu/jammy/amdgpu-install_5.7.50700-1_all.deb
+RUN apt-get update \
+    && apt-get install --no-install-recommends -y \
+    apt-utils \
+    ca-certificates \
+    libgomp1 \
+    python3 \
+    python3-pip \
+    python3-venv \
+    tzdata \
+    wget \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
-# Setup and install basic packages
+# Create basic build environment with build tools and compilers
+FROM base-runtime AS base-build
 RUN apt-get update \
-    && apt-get install --no-install-recommends -y apt-utils \
-    && DEBIAN_FRONTEND="noninteractive" \
-    apt-get install --no-install-recommends -y tzdata \
+    && apt-get install --no-install-recommends -y \
     build-essential \
-    ca-certificates \
     ccache \
     cmake \
     curl \
     git \
+    ninja-build \
+    python3-dev \
+    gcc-${GCC_VERSION} g++-${GCC_VERSION} cpp-${GCC_VERSION} \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+RUN update-alternatives \
+    --install /usr/bin/gcc gcc /usr/bin/gcc-${GCC_VERSION} 100 \
+    --slave /usr/bin/g++ g++ /usr/bin/g++-${GCC_VERSION} \
+    --slave /usr/bin/gcov gcov /usr/bin/gcov-${GCC_VERSION}
+RUN /usr/sbin/update-ccache-symlinks
+RUN mkdir /opt/ccache
+RUN ccache --set-config=cache_dir=/opt/ccache
+
+# Create and activate build virtual environment
+# Install Lightning dev requirements
+FROM base-build AS base-build-python
+WORKDIR /opt/pennylane-lightning
+ENV VIRTUAL_ENV=/opt/venv-build
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+RUN wget --progress=dot:giga https://github.com/PennyLaneAI/pennylane-lightning/archive/refs/tags/v${LIGHTNING_VERSION}.tar.gz
+RUN tar -zxf v${LIGHTNING_VERSION}.tar.gz && mv pennylane-lightning-${LIGHTNING_VERSION}/* /opt/pennylane-lightning
+RUN pip install --no-cache-dir wheel
+RUN pip install --no-cache-dir -r requirements-dev.txt
+RUN pip install --no-cache-dir -r doc/requirements.txt
+
+# Download Lightning release and build lightning-qubit backend
+FROM base-build-python AS build-wheel-lightning-qubit
+WORKDIR /opt/pennylane-lightning
+RUN pip uninstall -y pennylane-lightning
+RUN python setup.py build_ext --define="ENABLE_KOKKOS=OFF"
+RUN python setup.py bdist_wheel
+
+# Install lightning-qubit backend
+FROM base-runtime AS wheel-lightning-qubit
+COPY --from=build-wheel-lightning-qubit /opt/pennylane-lightning/dist/ /
+RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
+
+# Download Kokkos release and build OpenMP backend
+FROM base-build as build-kokkos-openmp
+WORKDIR /opt/kokkos-${KOKKOS_VERSION}
+RUN wget --progress=dot:giga https://github.com/kokkos/kokkos/archive/refs/tags/${KOKKOS_VERSION}.tar.gz
+RUN tar -zxf ${KOKKOS_VERSION}.tar.gz && mv /opt/kokkos-${KOKKOS_VERSION}/kokkos-${KOKKOS_VERSION}/* /opt/kokkos-${KOKKOS_VERSION}
+RUN cmake -S /opt/kokkos-${KOKKOS_VERSION} -B /opt/kokkos-${KOKKOS_VERSION}/build \
+    -G Ninja \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX=/opt/kokkos \
+    -DBUILD_SHARED_LIBS:BOOL=ON \
+    -DBUILD_TESTING:BOOL=OFF \
+    -DKokkos_ENABLE_SERIAL:BOOL=ON \
+    -DKokkos_ENABLE_OPENMP:BOOL=ON \
+    -DCMAKE_CXX_STANDARD=20 \
+    -DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
+    -DKokkos_ENABLE_TESTS:BOOL=OFF \
+    -DKokkos_ENABLE_LIBDL:BOOL=OFF
+RUN cmake --build /opt/kokkos-${KOKKOS_VERSION}/build
+RUN cmake --install /opt/kokkos-${KOKKOS_VERSION}/build
+
+# Download Lightning release and build lightning-kokkos backend with Kokkos-OpenMP
+FROM base-build-python AS build-wheel-lightning-kokkos-openmp
+WORKDIR /opt/pennylane-lightning
+COPY --from=build-kokkos-openmp /opt/kokkos /opt/kokkos
+ENV CMAKE_PREFIX_PATH=/opt/kokkos:$CMAKE_PREFIX_PATH
+RUN pip uninstall -y pennylane-lightning
+RUN PL_BACKEND=lightning_kokkos python setup.py build_ext
+RUN PL_BACKEND=lightning_kokkos python setup.py bdist_wheel
+
+# Install lightning-kokkos OpenMP backend
+FROM base-runtime AS wheel-lightning-kokkos-openmp
+COPY --from=build-kokkos-openmp /opt/kokkos /opt/kokkos
+COPY --from=build-wheel-lightning-kokkos-openmp /opt/pennylane-lightning/dist/ /
+RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
+
+# Install CUDA-12 in build venv image
+FROM base-build-python as base-build-cuda
+WORKDIR /opt/cuda-build
+RUN curl -o cuda-install.run ${CUDA_INSTALLER}
+RUN chmod a+x cuda-install.run
+RUN ./cuda-install.run --silent --toolkit --toolkitpath=/usr/local/cuda-$(echo ${CUDA_INSTALLER} | grep -o -P '/cuda/.{0,4}' | cut -d / -f 3)
+ENV PATH=/usr/local/cuda/bin:${PATH}
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}
+
+# Download Kokkos release and build CUDA backend
+FROM base-build-cuda as build-kokkos-cuda
+WORKDIR /opt/kokkos-${KOKKOS_VERSION}
+# RUN wget --progress=dot:giga https://github.com/kokkos/kokkos/archive/refs/tags/${KOKKOS_VERSION}.tar.gz
+# RUN tar -zxf ${KOKKOS_VERSION}.tar.gz && mv /opt/kokkos-${KOKKOS_VERSION}/kokkos-${KOKKOS_VERSION}/* /opt/kokkos-${KOKKOS_VERSION}
+RUN git clone --branch fix-rtld-deepbind https://github.com/cz4rs/kokkos.git /opt/kokkos-${KOKKOS_VERSION}
+RUN cmake -S /opt/kokkos-${KOKKOS_VERSION} -B /opt/kokkos-${KOKKOS_VERSION}/build \
+    -G Ninja \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX=/opt/kokkos \
+    -DBUILD_SHARED_LIBS:BOOL=ON \
+    -DBUILD_TESTING:BOOL=OFF \
+    -DKokkos_ENABLE_SERIAL:BOOL=ON \
+    -DKokkos_ENABLE_OPENMP:BOOL=ON \
+    -DCMAKE_CXX_STANDARD=20 \
+    -DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
+    -DKokkos_ENABLE_TESTS:BOOL=OFF \
+    -DKokkos_ENABLE_LIBDL:BOOL=OFF \
+    -DKokkos_ENABLE_CUDA:BOOL=ON \
+    -DKokkos_ARCH_${CUDA_ARCH}=ON
+RUN cmake --build /opt/kokkos-${KOKKOS_VERSION}/build
+RUN cmake --install /opt/kokkos-${KOKKOS_VERSION}/build
+
+# Download Lightning release and build lightning-kokkos backend with Kokkos-CUDA
+FROM base-build-cuda AS build-wheel-lightning-kokkos-cuda
+WORKDIR /opt/pennylane-lightning
+COPY --from=build-kokkos-cuda /opt/kokkos /opt/kokkos
+ENV CMAKE_PREFIX_PATH=/opt/kokkos:$CMAKE_PREFIX_PATH
+RUN pip uninstall -y pennylane-lightning
+RUN PL_BACKEND=lightning_kokkos python setup.py build_ext
+RUN PL_BACKEND=lightning_kokkos python setup.py bdist_wheel
+
+# Install python3 and setup runtime virtual env in CUDA-12-runtime image (includes CUDA runtime and math libraries)
+# Install lightning-kokkos CUDA backend
+FROM nvidia/cuda:12.2.0-base-ubuntu22.04 AS wheel-lightning-kokkos-cuda
+RUN apt-get update \
+    && apt-get install --no-install-recommends -y \
+    libgomp1 \
     python3 \
     python3-pip \
     python3-venv \
-    python3-dev \
-    gcc-11 g++-11 cpp-11 \
     && apt-get clean \
-    && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 --slave /usr/bin/gcov gcov /usr/bin/gcov-10 \
-    && /usr/sbin/update-ccache-symlinks \
-    && mkdir /opt/ccache \
-    && ccache --set-config=cache_dir=/opt/ccache \
-    && python3 -m venv /opt/venv
+    && rm -rf /var/lib/apt/lists/*
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+COPY --from=build-kokkos-cuda /opt/kokkos /opt/kokkos
+COPY --from=build-wheel-lightning-kokkos-cuda /opt/pennylane-lightning/dist/ /
+RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
 
-ENV PATH="/opt/venv/bin:$PATH"
+# Download and build Lightning-GPU release
+FROM base-build-cuda AS build-wheel-lightning-gpu
+WORKDIR /opt/pennylane-lightning-gpu
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+RUN wget --progress=dot:giga https://github.com/PennyLaneAI/pennylane-lightning-gpu/archive/refs/tags/v${LIGHTNING_VERSION}.tar.gz
+RUN tar -zxf v${LIGHTNING_VERSION}.tar.gz && ls -l
+RUN mv pennylane-lightning-gpu-${LIGHTNING_VERSION}/* /opt/pennylane-lightning-gpu
+RUN pip uninstall -y pennylane-lightning-gpu
+RUN pip install --no-cache-dir wheel custatevec-cu12
+RUN --mount=type=bind,source=requirements.txt,target=requirements.txt \
+    pip install --no-cache-dir -r requirements.txt
+RUN python setup.py build_ext --cuquantum=$(python -c "import site; print( f'{site.getsitepackages()[0]}/cuquantum/lib')")
+RUN python setup.py bdist_wheel
 
-# Setup and build pennylane-lightning
-WORKDIR /opt/pennylane-lightning
+# Install python3 and setup runtime virtual env in CUDA-12-runtime image (includes CUDA runtime and math libraries)
+# Install lightning-kokkos CUDA backend
+FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS wheel-lightning-gpu
+RUN apt-get update \
+    && apt-get install --no-install-recommends -y \
+    libgomp1 \
+    python3 \
+    python3-pip \
+    python3-venv \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+RUN pip install --no-cache-dir custatevec-cu12
+COPY --from=build-wheel-lightning-gpu /opt/pennylane-lightning-gpu/dist/ /
+RUN pip install --no-cache-dir --force-reinstall PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl
 
-COPY . .
+# Install ROCm in build venv image
+FROM base-build-python as base-build-rocm
+RUN wget --progress=dot:giga ${ROCM_INSTALLER}
+RUN apt-get update \
+    && apt-get install --no-install-recommends -y \
+    "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)" \
+    ./$(echo ${ROCM_INSTALLER} | xargs -I {} basename {}) \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+RUN amdgpu-install -y --accept-eula --usecase=rocm,hiplibsdk --no-dkms
+
+# Download Kokkos release and build HIP backend
+FROM base-build-rocm as build-kokkos-rocm
+WORKDIR /opt/kokkos-${KOKKOS_VERSION}
+ENV LD_LIBRARY_PATH=/opt/rocm/lib:$LD_LIBRARY_PATH
+# RUN wget --progress=dot:giga https://github.com/kokkos/kokkos/archive/refs/tags/${KOKKOS_VERSION}.tar.gz
+# RUN tar -zxf ${KOKKOS_VERSION}.tar.gz && mv /opt/kokkos-${KOKKOS_VERSION}/kokkos-${KOKKOS_VERSION}/* /opt/kokkos-${KOKKOS_VERSION}
+RUN git clone --branch fix-rtld-deepbind https://github.com/cz4rs/kokkos.git /opt/kokkos-${KOKKOS_VERSION}
+RUN CXX=hipcc cmake -S /opt/kokkos-${KOKKOS_VERSION} -B /opt/kokkos-${KOKKOS_VERSION}/build \
+    -G Ninja \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX=/opt/kokkos \
+    -DBUILD_SHARED_LIBS:BOOL=ON \
+    -DBUILD_TESTING:BOOL=OFF \
+    -DKokkos_ENABLE_SERIAL:BOOL=ON \
+    -DKokkos_ENABLE_OPENMP:BOOL=ON \
+    -DCMAKE_CXX_STANDARD=20 \
+    -DKokkos_ENABLE_EXAMPLES:BOOL=OFF \
+    -DKokkos_ENABLE_TESTS:BOOL=OFF \
+    -DKokkos_ENABLE_LIBDL:BOOL=OFF \
+    -DKokkos_ENABLE_HIP:BOOL=ON \
+    -DKokkos_ARCH_${AMD_ARCH}=ON
+RUN cmake --build /opt/kokkos-${KOKKOS_VERSION}/build
+RUN cmake --install /opt/kokkos-${KOKKOS_VERSION}/build
+
+# Download Lightning release and build lightning-kokkos backend with Kokkos-ROCm
+FROM base-build-rocm AS build-wheel-lightning-kokkos-rocm
+WORKDIR /opt/pennylane-lightning
+COPY --from=build-kokkos-rocm /opt/kokkos /opt/kokkos
+ENV CXX=hipcc
+ENV CMAKE_PREFIX_PATH=/opt/kokkos:$CMAKE_PREFIX_PATH
+RUN pip uninstall -y pennylane-lightning
+RUN PL_BACKEND=lightning_kokkos python setup.py build_ext
+RUN PL_BACKEND=lightning_kokkos python setup.py bdist_wheel
 
-RUN pip install --no-cache-dir wheel \
-    sphinx \
-    && pip install --no-cache-dir -r requirements-dev.txt \
-    && pip uninstall -y PennyLane_Lightning \
-    && make install \
-    && make test-python \
-    && make test-cpp
+# Install lightning-kokkos HIP backend
+FROM rocm/dev-ubuntu-22.04:5.7 AS wheel-lightning-kokkos-rocm
+RUN apt-get update \
+    && apt-get install --no-install-recommends -y \
+    libgomp1 \
+    python3 \
+    python3-pip \
+    python3-venv \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+COPY --from=build-kokkos-rocm /opt/kokkos /opt/kokkos
+COPY --from=build-wheel-lightning-kokkos-rocm /opt/pennylane-lightning/dist/ /
+RUN pip install --force-reinstall --no-cache-dir PennyLane_Lightning*.whl && rm PennyLane_Lightning*.whl