Archive home directory using multi-stage build

.github/workflows/docker-build-test-upload.yml

@@ -62,7 +62,7 @@ jobs:
                   context: .
                   platforms: linux/amd64
                   cache-to: |
-                      type=gha,scope=${{ github.workflow }},mode=max
+                      type=gha,scope=${{ github.workflow }},mode=min
                   cache-from: |
                       type=gha,scope=${{ github.workflow }}
 

Dockerfile

@@ -4,50 +4,100 @@ ARG UV_VER=0.2.27
 ARG QE_VER=7.2
 ARG QE_DIR=/opt/conda/envs/quantum-espresso-${QE_VER}
 
-FROM ghcr.io/astral-sh/uv:0.2.27 AS uv
+ARG UV_CACHE_DIR=/tmp/uv_cache
+ARG QE_APP_SRC=/tmp/quantum-espresso
+
+FROM ghcr.io/astral-sh/uv:${UV_VER} AS uv
 
 # STAGE 1
 # Install QE into conda environment in /opt/conda
-# This step is independent from the others and can be run in parallel.
+# This step is largely independent from the others and can run in parallel.
+# However, it needs to be done before running `python -m aiidalab_qe install-qe`,
+# otherwise QE gets installed into ~/.conda folder.
 FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER} AS qe_conda_env
 ARG QE_VER
 ARG QE_DIR
-RUN mamba create -p "${QE_DIR}" --yes qe="${QE_VER}" && \
-     mamba clean --all -f -y && \
-     fix-permissions "${CONDA_DIR}"
 
+USER ${NB_USER}
+RUN mamba create -p ${QE_DIR} --yes qe=${QE_VER} && \
+    mamba clean --all -f -y
 
 # STAGE 2
-FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER}
-ARG QE_VER
+# Install python dependencies needed to run aiidalab_qe CLI commands
+# uv package cache from this stage is reused in the final stage as well.
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER} AS build_deps
 ARG QE_DIR
-# Copy whole repo and pre-install the dependencies and app to the tmp folder.
-# In the before notebook scripts the app will be re-installed by moving it to the app folder.
-ENV PREINSTALL_APP_FOLDER=${CONDA_DIR}/aiidalab-qe
-COPY --chown=${NB_UID}:${NB_GID} . ${PREINSTALL_APP_FOLDER}
+ARG UV_CACHE_DIR
+ARG QE_APP_SRC
 
-USER ${NB_USER}
+WORKDIR ${QE_APP_SRC}
+COPY --chown=${NB_UID}:${NB_GID} src/ ${QE_APP_SRC}/src
+COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml LICENSE README.md ${QE_APP_SRC}
 
-# Using uv to speed up installation, per docs:
+# Use uv instead of pip to speed up installation, per docs:
 # https://github.com/astral-sh/uv/blob/main/docs/guides/docker.md#using-uv-temporarily
-# Use the same constraint file as PIP
+# Use the same constraint file as pip
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
-RUN  --mount=from=uv,source=/uv,target=/bin/uv \
-     cd ${PREINSTALL_APP_FOLDER} && \
-     # Remove all untracked files and directories. For example the setup lock flag file.
-     git clean -fx && \
-     uv pip install --system --no-cache . && \
-     fix-permissions "${CONDA_DIR}"
-
-# Download the QE pseudopotentials to the folder for afterware installation.
-ENV PSEUDO_FOLDER=${CONDA_DIR}/pseudo
+RUN --mount=from=uv,source=/uv,target=/bin/uv \
+    uv pip install --strict --system --cache-dir=${UV_CACHE_DIR} .
+
+# STAGE 3
+# - Prepare AiiDA profile and localhost computer
+# - Install QE codes and pseudopotentials
+# - Archive home folder
+FROM build_deps AS home_build
+ARG QE_DIR
+ENV PSEUDO_FOLDER=/tmp/pseudo
 RUN mkdir -p ${PSEUDO_FOLDER} && \
     python -m aiidalab_qe download-pseudos --dest ${PSEUDO_FOLDER}
 
-COPY --from=qe_conda_env "${QE_DIR}" "${QE_DIR}"
-# TODO: Remove this once we get rid of 70_prepare-qe-executable.sh
-ENV QE_VERSION="$QE_VER"
+# TODO: Remove PGSQL and daemon log files, and other unneeded files
+RUN --mount=from=qe_conda_env,source=${QE_DIR},target=${QE_DIR} \
+    bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
+    bash /usr/local/bin/before-notebook.d/40_prepare-aiida.sh && \
+    python -m aiidalab_qe install-qe && \
+    python -m aiidalab_qe install-pseudos --source ${PSEUDO_FOLDER} && \
+    verdi daemon stop && \
+    mamba run -n aiida-core-services pg_ctl stop && \
+    cd /home/${NB_USER} && tar -cf /opt/conda/home.tar .
 
-COPY before-notebook.d/* /usr/local/bin/before-notebook.d/
+# STAGE 3 - Final stage
+# - Install python dependencies
+# - Copy QE env environment
+# - Remove all content of home folder
+# - Copy the whole repo content into the container
+# - Copy home folder archive
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER}
+ARG QE_DIR
+ARG QE_APP_SRC
+ARG UV_CACHE_DIR
+USER ${NB_USER}
+
+WORKDIR /tmp
+# Install python dependencies
+# Use uv cache from the previous build step
+ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
+RUN --mount=from=uv,source=/uv,target=/bin/uv \
+    --mount=from=build_deps,source=${UV_CACHE_DIR},target=${UV_CACHE_DIR},rw \
+    --mount=from=build_deps,source=${QE_APP_SRC},target=${QE_APP_SRC},rw \
+    uv pip install --strict --system --compile-bytecode --cache-dir=${UV_CACHE_DIR} ${QE_APP_SRC}
+
+COPY --from=qe_conda_env ${QE_DIR} ${QE_DIR}
 
+USER root
+COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
+RUN fix-permissions "${CONDA_DIR}"
+# Remove content of $HOME
+# '-mindepth=1' ensures that we do not remove the home directory itself.
+RUN find /home/${NB_USER}/ -mindepth 1 -delete
+
+ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
+COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
+# Remove all untracked files and directories.
+RUN git clean -dffx || true
+
+ENV HOME_TAR="/opt/home.tar"
+COPY --from=home_build /opt/conda/home.tar "$HOME_TAR"
+
+USER ${NB_USER}
 WORKDIR "/home/${NB_USER}"

before-notebook.d/00_untar_home.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -eux
+
+home="/home/${NB_USER}"
+
+# Untar home archive file to restore home directory if it is empty
+if [[ $(ls -A ${home} | wc -l) = "0" ]]; then
+  if [[ ! -f $HOME_TAR ]]; then
+    echo "File $HOME_TAR does not exist!"
+    exit 1
+  fi
+  if [[ ! -d ${QE_APP_FOLDER} ]]; then
+    echo "Folder $QE_APP_FOLDER does not exist!"
+    exit 1
+  fi
+
+  echo "Extracting $HOME_TAR to $home"
+  tar -xf $HOME_TAR -C "$home"
+
+  echo "Copying directory '$QE_APP_FOLDER' to '$AIIDALAB_APPS'"
+  cp -r "$QE_APP_FOLDER" "$AIIDALAB_APPS"
+else
+  echo "$home folder is not empty!"
+  ls -lrta "$home"
+fi
+set +eux

tests_integration/conftest.py

@@ -59,7 +59,7 @@ def notebook_service(docker_compose, docker_ip, docker_services):
     token = os.environ.get("JUPYTER_TOKEN", "testtoken")
     try:
         docker_services.wait_until_responsive(
-            timeout=180.0,
+            timeout=60.0,
             pause=1.0,
             check=lambda: is_responsive(url),
         )