From 0bc327762e53af1641e74a7290e766b5763bd7f0 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Fri, 19 Jul 2024 10:16:03 +0100
Subject: [PATCH 01/23] Tar home build

---
 Dockerfile                                    | 79 ++++++++++---------
 before-notebook.d/00_untar_home.sh            | 16 ++++
 before-notebook.d/70_prepare-qe-executable.sh | 16 ----
 before-notebook.d/71_install-qeapp.sh         | 21 -----
 4 files changed, 59 insertions(+), 73 deletions(-)
 create mode 100644 before-notebook.d/00_untar_home.sh
 delete mode 100644 before-notebook.d/70_prepare-qe-executable.sh
 delete mode 100644 before-notebook.d/71_install-qeapp.sh

diff --git a/Dockerfile b/Dockerfile
index c53456900..5e3574b8a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,48 +1,55 @@
 # syntax=docker/dockerfile:1
 FROM ghcr.io/astral-sh/uv:0.2.18 AS uv
-FROM ghcr.io/aiidalab/full-stack:2024.1019
-
-# 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}
+FROM ghcr.io/aiidalab/full-stack:2024.1021
 
 USER ${NB_USER}
 
-# Using uv to speed up installation, per docs:
+ENV QE_VERSION="7.2"
+
+# 1. Install Quantum Espresso into a conda environment
+# (we do this here for better caching during Docker build)
+RUN mamba create -p /opt/conda/envs/quantum-espresso --yes \
+    qe=${QE_VERSION} && \
+    mkdir -p /home/${NB_USER}/.conda/envs && \
+    ln -s /opt/conda/envs/quantum-espresso /home/${NB_USER}/.conda/envs/quantum-espresso-${QE_VERSION} && \
+    mamba clean --all -f -y
+
+ENV QE_APP_FOLDER=${AIIDALAB_APPS}/quantum-espresso
+# 2. Copy files needed for installing stuff, the rest is copied at the end for better caching
+COPY --chown=${NB_UID}:${NB_GID} src/ ${QE_APP_FOLDER}/src
+COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml *yaml README.md ${QE_APP_FOLDER}
+
+WORKDIR "${QE_APP_FOLDER}"
+# 3.Install python dependencies
+# 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 && \
-     # It is important to install from `aiidalab install` to mimic the exact installation operation as
-     # from the app store.
-     # The command wil first install the dependencies from list by parsing setup config files,
-     # (for `aiidalab/aiidalab<23.03.2` the `setup.py` should be in the root folder of the app https://github.com/aiidalab/aiidalab/pull/382).
-     # and then the app and restart the daemon in the end.
-     # But since the aiida profile not yet exists, the daemon restart will fail but it is not a problem.
-     # Because we only need the dependencies to be installed.
-     # aiidalab install --yes --python ${CONDA_DIR}/bin/python  "quantum-espresso@file://${PREINSTALL_APP_FOLDER}" && \
-     # However, have to use `pip install` explicitly because `aiidalab install` call `pip install --user` which will install the app to `/home/${NB_USER}/.local`.
-     # It won't cause issue for docker but for k8s deployment the home folder is not bind mounted to the pod and the dependencies won't be found. (see issue in `jupyter/docker-stacks` https://github.com/jupyter/docker-stacks/issues/815)
-     uv pip install --system --no-cache . && \
-     fix-permissions "${CONDA_DIR}" && \
-     fix-permissions "/home/${NB_USER}"
+RUN --mount=from=uv,source=/uv,target=/bin/uv \
+    uv pip install --system --no-cache . && \
+    rm -rf build/ src/aiidalab_qe.egg-info/
+
+# 4. Prepare AiiDA profile and localhost computer
+# 5. Install the QE pseudopotentials and codes
+RUN 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 && \
+    verdi daemon stop && \
+    mamba run -n aiida-core-services pg_ctl stop
 
-ENV QE_VERSION="7.2"
-RUN mamba create -p /opt/conda/envs/quantum-espresso --yes \
-        qe=${QE_VERSION} \
-     && mamba clean --all -f -y && \
-     fix-permissions "${CONDA_DIR}" && \
-     fix-permissions "/home/${NB_USER}"
 
-# Download the QE pseudopotentials to the folder for afterware installation.
-ENV PSEUDO_FOLDER=${CONDA_DIR}/pseudo
-RUN mkdir -p ${PSEUDO_FOLDER} && \
-    python -m aiidalab_qe download-pseudos --dest ${PSEUDO_FOLDER}
+# 6. Copy the whole repo
+COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
+# Remove all untracked files and directories.
+RUN git clean -dffx || true
 
-COPY before-notebook.d/* /usr/local/bin/before-notebook.d/
+USER root
+COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
+RUN fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"
 
 WORKDIR "/home/${NB_USER}"
+RUN tar -cf /opt/home.tar .
+
+USER ${NB_USER}
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
new file mode 100644
index 000000000..5cf11cf8b
--- /dev/null
+++ b/before-notebook.d/00_untar_home.sh
@@ -0,0 +1,16 @@
+#!/bin/bash
+#
+home="/home/${NB_USER}"
+
+if [[ ! -d ${home} ]]; then
+    echo "Directory $home does not exist!"
+    exit 1
+fi
+
+# Untar home archive file to restore home directory if it is empty
+if [[ $(ls -A ${home} | wc -l) = "0" ]];then
+  if [ -f /opt/home.tar.gz ]; then
+    echo "Extracting /opt/home.tar to /home/${NB_USER}"
+    tar -xf /opt/home.tar -C /home/${NB_USER}
+  fi
+fi
diff --git a/before-notebook.d/70_prepare-qe-executable.sh b/before-notebook.d/70_prepare-qe-executable.sh
deleted file mode 100644
index 41ba2d391..000000000
--- a/before-notebook.d/70_prepare-qe-executable.sh
+++ /dev/null
@@ -1,16 +0,0 @@
-#!/bin/bash -e
-
-# Debugging.
-set -x
-
-# Copy quantum espresso env to user space.
-mkdir -p /home/${NB_USER}/.conda/envs
-if [ ! -d /home/${NB_USER}/.conda/envs/quantum-espresso-${QE_VERSION} ]; then
-  ln -s /opt/conda/envs/quantum-espresso /home/${NB_USER}/.conda/envs/quantum-espresso-${QE_VERSION}
-
-  # Install qe so the progress bar not shown in the notebook when first time using app.
-  echo "Installing qe."
-  python -m aiidalab_qe install-qe
-else
-  echo "Quantum ESPRESSO app is already installed."
-fi
diff --git a/before-notebook.d/71_install-qeapp.sh b/before-notebook.d/71_install-qeapp.sh
deleted file mode 100644
index 849d92849..000000000
--- a/before-notebook.d/71_install-qeapp.sh
+++ /dev/null
@@ -1,21 +0,0 @@
-#!/bin/bash -e
-
-# Debugging.
-set -x
-
-# Install qeapp if it is not already installed.
-if aiidalab list | grep -q quantum-espresso; then
-    echo "Quantum ESPRESSO app is already installed."
-else
-    echo "Installing Quantum ESPRESSO app."
-    # Install by move the repo folder that is already in the image.
-    mv ${PREINSTALL_APP_FOLDER} /home/${NB_USER}/apps/quantum-espresso
-fi
-
-# Install the pseudo libraries if not already installed.
-if aiida-pseudo list | grep -q "no pseudo potential families"; then
-    echo "Installing pseudo potential families."
-    python -m aiidalab_qe install-pseudos --source ${PSEUDO_FOLDER}
-else
-    echo "Pseudo potential families are already installed."
-fi

From 9f30615894ab55c350dcfe035d136c420737f6cf Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Fri, 19 Jul 2024 17:46:42 +0100
Subject: [PATCH 02/23] Add .dockerignore

---
 .dockerignore | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 .dockerignore

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 000000000..050466577
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,14 @@
+**/*.egg-info
+**/*.pyc
+**/*.tar.gz
+*.code-workspace
+**/.*.ipynb
+**/.ipynb*
+.venv/
+build/
+export/
+.do-not-setup-on-localhost
+
+# Sphinx documentation
+docs/html
+screenshots/

From 99ee0ce5a8747f79bd9e0e079a7560353f2747b0 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Fri, 19 Jul 2024 17:58:25 +0100
Subject: [PATCH 03/23] WIP: Smaller image

---
 Dockerfile                         | 12 ++++++------
 before-notebook.d/00_untar_home.sh |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 5e3574b8a..d877c18b9 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -31,18 +31,20 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
 
 # 4. Prepare AiiDA profile and localhost computer
 # 5. Install the QE pseudopotentials and codes
+# TODO: Remove PGSQL and daemon log files, and other unneeded files
 RUN 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 && \
     verdi daemon stop && \
-    mamba run -n aiida-core-services pg_ctl stop
-
+    mamba run -n aiida-core-services pg_ctl stop && \
+    cd /home/${NB_USER} && tar -cf /opt/conda/home.tar . && \
+    rm -rf *
 
 # 6. Copy the whole repo
-COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
+#COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
 # Remove all untracked files and directories.
-RUN git clean -dffx || true
+#RUN git clean -dffx || true
 
 USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
@@ -50,6 +52,4 @@ RUN fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 
 WORKDIR "/home/${NB_USER}"
-RUN tar -cf /opt/home.tar .
-
 USER ${NB_USER}
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 5cf11cf8b..097f98b76 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -11,6 +11,6 @@ fi
 if [[ $(ls -A ${home} | wc -l) = "0" ]];then
   if [ -f /opt/home.tar.gz ]; then
     echo "Extracting /opt/home.tar to /home/${NB_USER}"
-    tar -xf /opt/home.tar -C /home/${NB_USER}
+    tar -xf /opt/conda/home.tar -C /home/${NB_USER}
   fi
 fi

From 45c459675ac964f6ecebb412855ea3a878196e16 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 02:01:23 +0100
Subject: [PATCH 04/23] Still doesn't work

---
 Dockerfile                         | 52 +++++++++++++++++-------------
 before-notebook.d/00_untar_home.sh |  2 ++
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index d877c18b9..6e48e319b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,29 +1,15 @@
 # syntax=docker/dockerfile:1
 FROM ghcr.io/astral-sh/uv:0.2.18 AS uv
-FROM ghcr.io/aiidalab/full-stack:2024.1021
 
+FROM ghcr.io/aiidalab/full-stack:2024.1021 AS home_build
 USER ${NB_USER}
 
-ENV QE_VERSION="7.2"
+ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
+WORKDIR ${QE_APP_FOLDER}
 
-# 1. Install Quantum Espresso into a conda environment
-# (we do this here for better caching during Docker build)
-RUN mamba create -p /opt/conda/envs/quantum-espresso --yes \
-    qe=${QE_VERSION} && \
-    mkdir -p /home/${NB_USER}/.conda/envs && \
-    ln -s /opt/conda/envs/quantum-espresso /home/${NB_USER}/.conda/envs/quantum-espresso-${QE_VERSION} && \
-    mamba clean --all -f -y
-
-ENV QE_APP_FOLDER=${AIIDALAB_APPS}/quantum-espresso
-# 2. Copy files needed for installing stuff, the rest is copied at the end for better caching
 COPY --chown=${NB_UID}:${NB_GID} src/ ${QE_APP_FOLDER}/src
 COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml *yaml README.md ${QE_APP_FOLDER}
 
-WORKDIR "${QE_APP_FOLDER}"
-# 3.Install python dependencies
-# 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
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
     uv pip install --system --no-cache . && \
@@ -38,13 +24,35 @@ RUN bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
     python -m aiidalab_qe install-pseudos && \
     verdi daemon stop && \
     mamba run -n aiida-core-services pg_ctl stop && \
-    cd /home/${NB_USER} && tar -cf /opt/conda/home.tar . && \
-    rm -rf *
+    cd /home/${NB_USER} && tar -cf /opt/conda/home.tar .
 
-# 6. Copy the whole repo
-#COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
+FROM ghcr.io/aiidalab/full-stack:2024.1021
+COPY --from=home_build /opt/conda/home.tar /opt/conda
+
+USER ${NB_USER}
+
+ENV QE_VERSION="7.2"
+# 1. Install Quantum Espresso into a conda environment
+# (we do this here for better caching during Docker build)
+RUN mamba create -p /opt/conda/envs/quantum-espresso --yes \
+    qe=${QE_VERSION} && \
+    mamba clean --all -f -y
+
+ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
+WORKDIR "${QE_APP_FOLDER}"
+
+COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
 # Remove all untracked files and directories.
-#RUN git clean -dffx || true
+RUN git clean -dffx || true
+
+# 3.Install python dependencies
+# 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
+ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
+RUN --mount=from=uv,source=/uv,target=/bin/uv \
+    uv pip install --system --no-cache . && \
+    rm -rf build/ src/aiidalab_qe.egg-info/
 
 USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 097f98b76..0f40254ac 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -12,5 +12,7 @@ if [[ $(ls -A ${home} | wc -l) = "0" ]];then
   if [ -f /opt/home.tar.gz ]; then
     echo "Extracting /opt/home.tar to /home/${NB_USER}"
     tar -xf /opt/conda/home.tar -C /home/${NB_USER}
+    mkdir -p /home/${NB_USER}/.conda/envs && \
+    ln -s /opt/conda/envs/quantum-espresso /home/${NB_USER}/.conda/envs/quantum-espresso-${QE_VERSION} && \
   fi
 fi

From a685221494b29c17b5f86df5d042222f3d823b3f Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 14:16:45 +0100
Subject: [PATCH 05/23] WIP: Multistage

---
 Dockerfile                            | 28 +++++++++++++--------
 before-notebook.d/00_untar_home.sh    |  8 +++---
 src/aiidalab_qe/common/setup_codes.py | 35 ++++++++++++++++-----------
 3 files changed, 43 insertions(+), 28 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 6e48e319b..a3cc44cfe 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,8 +1,18 @@
 # syntax=docker/dockerfile:1
 FROM ghcr.io/astral-sh/uv:0.2.18 AS uv
 
-FROM ghcr.io/aiidalab/full-stack:2024.1021 AS home_build
+FROM ghcr.io/aiidalab/full-stack:2024.1021 AS conda_build
+
 USER ${NB_USER}
+# Install QE into conda environment in /opt/conda
+# We need to do this first, otherwise QE gets installed into home folder as part of
+# python -m aiidalab_qe install-qe
+ENV QE_VERSION="7.2"
+RUN mamba create -p /opt/conda/envs/quantum-espresso-${QE_VERSION} --yes \
+    qe=${QE_VERSION} && \
+    mamba clean --all -f -y
+
+FROM conda_build AS home_build
 
 ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
 WORKDIR ${QE_APP_FOLDER}
@@ -21,26 +31,24 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
 RUN 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 && \
+    #python -m aiidalab_qe install-pseudos && \
     verdi daemon stop && \
     mamba run -n aiida-core-services pg_ctl stop && \
     cd /home/${NB_USER} && tar -cf /opt/conda/home.tar .
 
+
+# TODO: Deduplicate the name of the full-stack image
 FROM ghcr.io/aiidalab/full-stack:2024.1021
-COPY --from=home_build /opt/conda/home.tar /opt/conda
 
 USER ${NB_USER}
 
 ENV QE_VERSION="7.2"
-# 1. Install Quantum Espresso into a conda environment
-# (we do this here for better caching during Docker build)
-RUN mamba create -p /opt/conda/envs/quantum-espresso --yes \
-    qe=${QE_VERSION} && \
-    mamba clean --all -f -y
+COPY --from=conda_build /opt/conda/envs/quantum-espresso-${QE_VERSION}/ /opt/conda/envs/quantum-espresso-${QE_VERSION}
+
+COPY --from=home_build /opt/conda/home.tar /opt/conda
 
 ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
 WORKDIR "${QE_APP_FOLDER}"
-
 COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
 # Remove all untracked files and directories.
 RUN git clean -dffx || true
@@ -51,7 +59,7 @@ RUN git clean -dffx || true
 # Use the same constraint file as pip
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
-    uv pip install --system --no-cache . && \
+    uv pip install --system --compile-bytecode --no-cache . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 
 USER root
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 0f40254ac..93005ec55 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -1,5 +1,4 @@
 #!/bin/bash
-#
 home="/home/${NB_USER}"
 
 if [[ ! -d ${home} ]]; then
@@ -9,10 +8,11 @@ fi
 
 # Untar home archive file to restore home directory if it is empty
 if [[ $(ls -A ${home} | wc -l) = "0" ]];then
-  if [ -f /opt/home.tar.gz ]; then
+  if [[ -f /opt/home.tar.gz ]]; then
     echo "Extracting /opt/home.tar to /home/${NB_USER}"
     tar -xf /opt/conda/home.tar -C /home/${NB_USER}
-    mkdir -p /home/${NB_USER}/.conda/envs && \
-    ln -s /opt/conda/envs/quantum-espresso /home/${NB_USER}/.conda/envs/quantum-espresso-${QE_VERSION} && \
+  fi
+  if [[ -n ${QE_APP_FOLDER} && -f ${QE_APP_FOLDER} ]]; then
+    cp -r "$QE_APP_FOLDER" "$home/apps/"
   fi
 fi
diff --git a/src/aiidalab_qe/common/setup_codes.py b/src/aiidalab_qe/common/setup_codes.py
index 453e279fa..11ee58d1a 100644
--- a/src/aiidalab_qe/common/setup_codes.py
+++ b/src/aiidalab_qe/common/setup_codes.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 from pathlib import Path
 from shutil import which
 from subprocess import CalledProcessError, run
@@ -21,9 +20,15 @@
 
 QE_VERSION = "7.2"
 
-CONDA_ENV_PREFIX = Path.home().joinpath(
-    ".conda", "envs", f"quantum-espresso-{QE_VERSION}"
-)
+
+def get_qe_env():
+    # QE is already pre-installed in the QE image
+    path = Path(f"/opt/conda/envs/quantum-espresso-{QE_VERSION}")
+    if path.exists():
+        return path
+    else:
+        return Path.home().joinpath(".conda", "envs", f"quantum-espresso-{QE_VERSION}")
+
 
 # Add all QE codes with the calcjob entry point in the aiida-quantumespresso.
 CODE_NAMES = (
@@ -46,7 +51,7 @@
 
 
 def qe_installed():
-    return CONDA_ENV_PREFIX.exists()
+    return get_qe_env().exists()
 
 
 def install_qe():
@@ -59,7 +64,7 @@ def install_qe():
             "--channel",
             "conda-forge",
             "--prefix",
-            str(CONDA_ENV_PREFIX),
+            str(get_qe_env()),
             f"qe={QE_VERSION}",
         ],
         capture_output=True,
@@ -102,9 +107,9 @@ def _generate_string_to_setup_code(code_name, computer_name="localhost"):
     except NotExistent:
         label = f"{code_name}-{QE_VERSION}"
         description = f"{code_name}.x ({QE_VERSION}) setup by AiiDAlab."
-        filepath_executable = str(CONDA_ENV_PREFIX.joinpath("bin", f"{code_name}.x"))
+        filepath_executable = get_qe_env().joinpath("bin", f"{code_name}.x")
         default_calc_job_plugin = f"quantumespresso.{code_name}"
-        prepend_text = f'eval "$(conda shell.posix hook)"\\nconda activate {CONDA_ENV_PREFIX}\\nexport OMP_NUM_THREADS=1'
+        prepend_text = f'eval "$(conda shell.posix hook)"\\nconda activate {get_qe_env()}\\nexport OMP_NUM_THREADS=1'
         python_code = """
 computer = load_computer('{}')
 code = InstalledCode(computer=computer,
@@ -116,7 +121,7 @@ def _generate_string_to_setup_code(code_name, computer_name="localhost"):
                     )
 
 code.store()
-""".format(
+""".format(  # noqa: UP032
             computer_name,
             label,
             description,
@@ -134,7 +139,7 @@ def setup_codes():
     try:
         run(["python", "-c", python_code], capture_output=True, check=True)
     except CalledProcessError as error:
-        raise RuntimeError(f"Failed to setup codes: {error}")
+        raise RuntimeError(f"Failed to setup codes: {error}") from None
 
 
 def install(force=False):
@@ -175,7 +180,9 @@ def install(force=False):
                 try:
                     install_qe()
                 except CalledProcessError as error:
-                    raise RuntimeError(f"Failed to create conda environment: {error}")
+                    raise RuntimeError(
+                        f"Failed to create conda environment: {error}"
+                    ) from None
 
             # After installing QE, we install the corresponding
             # AiiDA codes:
@@ -189,7 +196,7 @@ def install(force=False):
                 yield "Setting up all codes..."
                 run(["python", "-c", python_code], capture_output=True, check=True)
             except CalledProcessError as error:
-                raise RuntimeError(f"Failed to setup codes: {error}")
+                raise RuntimeError(f"Failed to setup codes: {error}") from None
 
     except Timeout:
         # Assume that the installation was triggered by a different process.
@@ -198,7 +205,7 @@ def install(force=False):
             if not codes_are_setup():
                 raise RuntimeError(
                     "Installation process did not finish in the expected time."
-                )
+                ) from None
 
 
 class QESetupWidget(ipw.VBox):
@@ -311,7 +318,7 @@ def _toggle_error_view(self, change):
     @traitlets.observe("busy")
     @traitlets.observe("error")
     @traitlets.observe("installed")
-    def _update(self, change):
+    def _update(self, _change):
         with self.hold_trait_notifications():
             if self.hide_by_default:
                 self.layout.visibility = (

From fd1ef03e0ff007c76b532ce6bb05fe6e3dc42226 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 14:21:38 +0100
Subject: [PATCH 06/23] Pre-download pseudos

---
 Dockerfile | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index a3cc44cfe..62c243f08 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -25,13 +25,17 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
     uv pip install --system --no-cache . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 
+ENV PSEUDO_FOLDER=/tmp/pseudo
+RUN mkdir -p ${PSEUDO_FOLDER} && \
+    python -m aiidalab_qe download-pseudos --dest ${PSEUDO_FOLDER}
+
 # 4. Prepare AiiDA profile and localhost computer
 # 5. Install the QE pseudopotentials and codes
 # TODO: Remove PGSQL and daemon log files, and other unneeded files
 RUN 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 && \
+    python -m aiidalab_qe install-pseudos && \
     verdi daemon stop && \
     mamba run -n aiida-core-services pg_ctl stop && \
     cd /home/${NB_USER} && tar -cf /opt/conda/home.tar .

From 0521d29e00e700e4c8ebb39a4f899b8bc394b935 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 16:06:40 +0100
Subject: [PATCH 07/23] Fixes

---
 .dockerignore                      |  2 ++
 Dockerfile                         | 23 ++++++++++++++---------
 before-notebook.d/00_untar_home.sh | 28 +++++++++++++++++-----------
 3 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/.dockerignore b/.dockerignore
index 050466577..f45fee4b5 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,4 +1,6 @@
 **/*.egg-info
+**/__pycache__
+**/.*_cache
 **/*.pyc
 **/*.tar.gz
 *.code-workspace
diff --git a/Dockerfile b/Dockerfile
index 62c243f08..fea66ff88 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,13 +1,17 @@
 # syntax=docker/dockerfile:1
-FROM ghcr.io/astral-sh/uv:0.2.18 AS uv
+ARG FULL_STACK_VERSION=2024.1021
+ARG QE_VERSION=7.2
+ARG UV_VERSION=0.2.27
 
-FROM ghcr.io/aiidalab/full-stack:2024.1021 AS conda_build
+FROM ghcr.io/astral-sh/uv:${UV_VERSION} AS uv
+
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION} AS conda_build
+ARG QE_VERSION
 
 USER ${NB_USER}
 # Install QE into conda environment in /opt/conda
 # We need to do this first, otherwise QE gets installed into home folder as part of
 # python -m aiidalab_qe install-qe
-ENV QE_VERSION="7.2"
 RUN mamba create -p /opt/conda/envs/quantum-espresso-${QE_VERSION} --yes \
     qe=${QE_VERSION} && \
     mamba clean --all -f -y
@@ -35,18 +39,16 @@ RUN mkdir -p ${PSEUDO_FOLDER} && \
 RUN 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 && \
+    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 .
 
 
-# TODO: Deduplicate the name of the full-stack image
-FROM ghcr.io/aiidalab/full-stack:2024.1021
-
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION}
+ARG QE_VERSION
 USER ${NB_USER}
 
-ENV QE_VERSION="7.2"
 COPY --from=conda_build /opt/conda/envs/quantum-espresso-${QE_VERSION}/ /opt/conda/envs/quantum-espresso-${QE_VERSION}
 
 COPY --from=home_build /opt/conda/home.tar /opt/conda
@@ -63,7 +65,7 @@ RUN git clean -dffx || true
 # Use the same constraint file as pip
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
-    uv pip install --system --compile-bytecode --no-cache . && \
+    uv pip install --strict --system --compile-bytecode --no-cache . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 
 USER root
@@ -71,5 +73,8 @@ COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
 RUN fix-permissions "${CONDA_DIR}" && \
     fix-permissions "/home/${NB_USER}"
 
+# REMOVE HOME
+RUN find /home/${NB_USER}/ -delete
+
 WORKDIR "/home/${NB_USER}"
 USER ${NB_USER}
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 93005ec55..4974d2b03 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -1,18 +1,24 @@
 #!/bin/bash
-home="/home/${NB_USER}"
+set -u
 
-if [[ ! -d ${home} ]]; then
-    echo "Directory $home does not exist!"
-    exit 1
-fi
+home="/home/${NB_USER}"
+HOME_TAR="/opt/conda/home.tar"
 
 # Untar home archive file to restore home directory if it is empty
-if [[ $(ls -A ${home} | wc -l) = "0" ]];then
-  if [[ -f /opt/home.tar.gz ]]; then
-    echo "Extracting /opt/home.tar to /home/${NB_USER}"
-    tar -xf /opt/conda/home.tar -C /home/${NB_USER}
+if [[ $(ls -A ${home} | wc -l) = "0" ]]; then
+  if [[ ! -f $HOME_TAR ]]; then
+    echo "File $HOME_TAR does not exist!"
+    exit 1
   fi
-  if [[ -n ${QE_APP_FOLDER} && -f ${QE_APP_FOLDER} ]]; then
-    cp -r "$QE_APP_FOLDER" "$home/apps/"
+  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 '$home/apps/'"
+  cp -r "$QE_APP_FOLDER" "$home/apps/"
 fi
+set +u

From 68dc3c0257a57875ed9d2a35abe015d0090a65dd Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 16:57:48 +0100
Subject: [PATCH 08/23] Build deps stage

---
 Dockerfile | 33 ++++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index fea66ff88..f126c5589 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,20 +5,22 @@ ARG UV_VERSION=0.2.27
 
 FROM ghcr.io/astral-sh/uv:${UV_VERSION} AS uv
 
+# STAGE 1
+# Install QE into conda environment in /opt/conda
+# We need to do this first, otherwise QE gets installed into home folder as part of
+# python -m aiidalab_qe install-qe
 FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION} AS conda_build
 ARG QE_VERSION
 
 USER ${NB_USER}
-# Install QE into conda environment in /opt/conda
-# We need to do this first, otherwise QE gets installed into home folder as part of
-# python -m aiidalab_qe install-qe
 RUN mamba create -p /opt/conda/envs/quantum-espresso-${QE_VERSION} --yes \
     qe=${QE_VERSION} && \
     mamba clean --all -f -y
 
-FROM conda_build AS home_build
-
-ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
+# STAGE 2
+# Install python dependencies, needed to run aiidalab_qe CLI commands
+FROM conda_build AS build_deps
+ENV QE_APP_FOLDER=/tmp/quantum-espresso
 WORKDIR ${QE_APP_FOLDER}
 
 COPY --chown=${NB_UID}:${NB_GID} src/ ${QE_APP_FOLDER}/src
@@ -29,12 +31,17 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
     uv pip install --system --no-cache . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 
+
+# STAGE 3:
+# - Prepare AiiDA profile and localhost computer
+# - Install QE codes and pseudopotentials
+# - Archive home folder
+FROM build_deps AS home_build
+
 ENV PSEUDO_FOLDER=/tmp/pseudo
 RUN mkdir -p ${PSEUDO_FOLDER} && \
     python -m aiidalab_qe download-pseudos --dest ${PSEUDO_FOLDER}
 
-# 4. Prepare AiiDA profile and localhost computer
-# 5. Install the QE pseudopotentials and codes
 # TODO: Remove PGSQL and daemon log files, and other unneeded files
 RUN bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
     bash /usr/local/bin/before-notebook.d/40_prepare-aiida.sh && \
@@ -44,14 +51,18 @@ RUN bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
     mamba run -n aiida-core-services pg_ctl stop && \
     cd /home/${NB_USER} && tar -cf /opt/conda/home.tar .
 
-
+# STAGE 4 - Final stage
+# - Copy QE env environment
+# - Copy home folder archive
+# - Copy the whole repo into the container
+# - Install python dependencies
+# - Remove all content of home folder
 FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION}
 ARG QE_VERSION
 USER ${NB_USER}
 
 COPY --from=conda_build /opt/conda/envs/quantum-espresso-${QE_VERSION}/ /opt/conda/envs/quantum-espresso-${QE_VERSION}
-
-COPY --from=home_build /opt/conda/home.tar /opt/conda
+COPY --from=home_build /opt/conda/home.tar /opt/conda/home.tar
 
 ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
 WORKDIR "${QE_APP_FOLDER}"

From 94254ed783ff177e677b3c35f666ff829dce921e Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 17:13:02 +0100
Subject: [PATCH 09/23] Use uv cache

---
 Dockerfile | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index f126c5589..b835bba41 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,6 +3,7 @@ ARG FULL_STACK_VERSION=2024.1021
 ARG QE_VERSION=7.2
 ARG UV_VERSION=0.2.27
 
+ARG UV_CACHE_DIR=/tmp/uv_cache
 FROM ghcr.io/astral-sh/uv:${UV_VERSION} AS uv
 
 # STAGE 1
@@ -17,9 +18,11 @@ RUN mamba create -p /opt/conda/envs/quantum-espresso-${QE_VERSION} --yes \
     qe=${QE_VERSION} && \
     mamba clean --all -f -y
 
+
 # STAGE 2
 # Install python dependencies, needed to run aiidalab_qe CLI commands
 FROM conda_build AS build_deps
+ARG UV_CACHE_DIR
 ENV QE_APP_FOLDER=/tmp/quantum-espresso
 WORKDIR ${QE_APP_FOLDER}
 
@@ -28,8 +31,7 @@ COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml *yaml README.md ${QE_A
 
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
-    uv pip install --system --no-cache . && \
-    rm -rf build/ src/aiidalab_qe.egg-info/
+    uv pip install --strict --system --cache-dir=${UV_CACHE_DIR} .
 
 
 # STAGE 3:
@@ -59,6 +61,7 @@ RUN bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
 # - Remove all content of home folder
 FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION}
 ARG QE_VERSION
+ARG UV_CACHE_DIR
 USER ${NB_USER}
 
 COPY --from=conda_build /opt/conda/envs/quantum-espresso-${QE_VERSION}/ /opt/conda/envs/quantum-espresso-${QE_VERSION}
@@ -76,13 +79,13 @@ RUN git clean -dffx || true
 # Use the same constraint file as pip
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
-    uv pip install --strict --system --compile-bytecode --no-cache . && \
+    --mount=from=build_deps,source=${UV_CACHE_DIR},target=${UV_CACHE_DIR},rw \
+    uv pip install --strict --system --compile-bytecode --cache-dir=${UV_CACHE_DIR} . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 
 USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
-RUN fix-permissions "${CONDA_DIR}" && \
-    fix-permissions "/home/${NB_USER}"
+RUN fix-permissions "${CONDA_DIR}"
 
 # REMOVE HOME
 RUN find /home/${NB_USER}/ -delete

From 7c30427053ac30e3bc453ce60cc5d156a94e76ff Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 17:55:01 +0100
Subject: [PATCH 10/23] Indepedent build steps

---
 Dockerfile | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index b835bba41..31b80293e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -4,24 +4,26 @@ ARG QE_VERSION=7.2
 ARG UV_VERSION=0.2.27
 
 ARG UV_CACHE_DIR=/tmp/uv_cache
+ARG QE_DIR=/opt/conda/envs/quantum-espresso-${QE_VERSION}
 FROM ghcr.io/astral-sh/uv:${UV_VERSION} AS uv
 
 # STAGE 1
 # Install QE into conda environment in /opt/conda
 # We need to do this first, otherwise QE gets installed into home folder as part of
 # python -m aiidalab_qe install-qe
-FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION} AS conda_build
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION} AS qe_conda_env
 ARG QE_VERSION
+ARG QE_DIR
 
 USER ${NB_USER}
-RUN mamba create -p /opt/conda/envs/quantum-espresso-${QE_VERSION} --yes \
+RUN mamba create -p ${QE_DIR} --yes \
     qe=${QE_VERSION} && \
     mamba clean --all -f -y
 
 
 # STAGE 2
-# Install python dependencies, needed to run aiidalab_qe CLI commands
-FROM conda_build AS build_deps
+# Install python dependencies needed to run aiidalab_qe CLI commands
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION} AS build_deps
 ARG UV_CACHE_DIR
 ENV QE_APP_FOLDER=/tmp/quantum-espresso
 WORKDIR ${QE_APP_FOLDER}
@@ -29,6 +31,9 @@ WORKDIR ${QE_APP_FOLDER}
 COPY --chown=${NB_UID}:${NB_GID} src/ ${QE_APP_FOLDER}/src
 COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml *yaml README.md ${QE_APP_FOLDER}
 
+# 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
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
     uv pip install --strict --system --cache-dir=${UV_CACHE_DIR} .
@@ -39,13 +44,15 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
 # - 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}
 
 # TODO: Remove PGSQL and daemon log files, and other unneeded files
-RUN bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
+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} && \
@@ -56,15 +63,15 @@ RUN bash /usr/local/bin/before-notebook.d/20_start-postgresql.sh && \
 # STAGE 4 - Final stage
 # - Copy QE env environment
 # - Copy home folder archive
-# - Copy the whole repo into the container
+# - Copy the whole repo content into the container
 # - Install python dependencies
 # - Remove all content of home folder
 FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VERSION}
-ARG QE_VERSION
+ARG QE_DIR
 ARG UV_CACHE_DIR
 USER ${NB_USER}
 
-COPY --from=conda_build /opt/conda/envs/quantum-espresso-${QE_VERSION}/ /opt/conda/envs/quantum-espresso-${QE_VERSION}
+COPY --from=qe_conda_env ${QE_DIR} ${QE_DIR}
 COPY --from=home_build /opt/conda/home.tar /opt/conda/home.tar
 
 ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
@@ -74,9 +81,7 @@ COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
 RUN git clean -dffx || true
 
 # 3.Install python dependencies
-# 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 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 \

From dba9fbb5d879a49eaf21b8a7af133f2d27a0f872 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 18:26:59 +0100
Subject: [PATCH 11/23] Trace startup script

---
 Dockerfile                         | 3 ++-
 before-notebook.d/00_untar_home.sh | 9 ++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 31b80293e..3511a9f3c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -90,7 +90,8 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
 
 USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
-RUN fix-permissions "${CONDA_DIR}"
+RUN fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"
 
 # REMOVE HOME
 RUN find /home/${NB_USER}/ -delete
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 4974d2b03..32d519148 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-set -u
+set -eux
 
 home="/home/${NB_USER}"
 HOME_TAR="/opt/conda/home.tar"
@@ -16,9 +16,12 @@ if [[ $(ls -A ${home} | wc -l) = "0" ]]; then
   fi
 
   echo "Extracting $HOME_TAR to $home"
-  tar -xf $HOME_TAR -C $home
+  tar -xf $HOME_TAR -C "$home"
 
   echo "Copying directory '$QE_APP_FOLDER' to '$home/apps/'"
   cp -r "$QE_APP_FOLDER" "$home/apps/"
+else
+  echo "$home folder is not empty!"
+  ls -lrta "$home"
 fi
-set +u
+set +eux

From 192e5587d3191b97ccf89607fe3f87b031c7e8c0 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 19:22:01 +0100
Subject: [PATCH 12/23] Smaller timeout, print logs when startup fails

---
 .github/workflows/docker-build-test-upload.yml |  2 +-
 tests_integration/conftest.py                  | 17 ++++++++++-------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/docker-build-test-upload.yml b/.github/workflows/docker-build-test-upload.yml
index c285cf4d0..f8f0e3b55 100644
--- a/.github/workflows/docker-build-test-upload.yml
+++ b/.github/workflows/docker-build-test-upload.yml
@@ -78,7 +78,7 @@ jobs:
               run: echo "JUPYTER_TOKEN=$(openssl rand -hex 32)" >> $GITHUB_ENV
 
             - name: Run pytest for Chrome
-              run: pytest --driver Chrome tests_integration/
+              run: pytest -sv --driver Chrome tests_integration/
               env:
                   QE_IMAGE: ${{ env.IMAGE }}@${{ steps.build-upload.outputs.digest }}
 
diff --git a/tests_integration/conftest.py b/tests_integration/conftest.py
index cbd576174..ff0f9e036 100644
--- a/tests_integration/conftest.py
+++ b/tests_integration/conftest.py
@@ -50,19 +50,22 @@ def nb_user(aiidalab_exec):
 
 
 @pytest.fixture(scope="session")
-def notebook_service(docker_ip, docker_services):
+def notebook_service(docker_compose, docker_ip, docker_services):
     """Ensure that HTTP service is up and responsive."""
 
     # `port_for` takes a container port and returns the corresponding host port
     port = docker_services.port_for("aiidalab", 8888)
     url = f"http://{docker_ip}:{port}"
     token = os.environ.get("JUPYTER_TOKEN", "testtoken")
-    docker_services.wait_until_responsive(
-        # The timeout is very high for this test, because the installation of pseudo libraries.
-        timeout=180.0,
-        pause=0.1,
-        check=lambda: is_responsive(url),
-    )
+    try:
+        docker_services.wait_until_responsive(
+            timeout=30.0,
+            pause=0.5,
+            check=lambda: is_responsive(url),
+        )
+    except Exception as e:
+        print(docker_compose.execute("logs").decode().strip())
+        pytest.exit(e)
     return url, token
 
 

From f735c2bfa04274d6a9c9651dce7b661f3ed9ecb1 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 20:00:22 +0100
Subject: [PATCH 13/23] Use AIIDALAB_APPS

---
 before-notebook.d/00_untar_home.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 32d519148..62aa23b1b 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -18,8 +18,8 @@ if [[ $(ls -A ${home} | wc -l) = "0" ]]; then
   echo "Extracting $HOME_TAR to $home"
   tar -xf $HOME_TAR -C "$home"
 
-  echo "Copying directory '$QE_APP_FOLDER' to '$home/apps/'"
-  cp -r "$QE_APP_FOLDER" "$home/apps/"
+  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"

From 0bbe0b18c18177cd2458b8f2ca1fe6f45db56728 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 20:07:13 +0100
Subject: [PATCH 14/23] Fix permissions?

---
 Dockerfile | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 3511a9f3c..e90117fe2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -90,11 +90,11 @@ RUN --mount=from=uv,source=/uv,target=/bin/uv \
 
 USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
-RUN fix-permissions "${CONDA_DIR}" && \
-    fix-permissions "/home/${NB_USER}"
+RUN fix-permissions "${CONDA_DIR}"
 
 # REMOVE HOME
 RUN find /home/${NB_USER}/ -delete
 
 WORKDIR "/home/${NB_USER}"
+RUN fix-permissions "/home/${NB_USER}"
 USER ${NB_USER}

From 48b399097318f7ad8fc79e0ccd0904bd7dca86a4 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Mon, 22 Jul 2024 21:16:18 +0100
Subject: [PATCH 15/23] Don't remove home dir

---
 Dockerfile | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index e90117fe2..01ef9174b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -92,9 +92,9 @@ USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
 RUN fix-permissions "${CONDA_DIR}"
 
-# REMOVE HOME
-RUN find /home/${NB_USER}/ -delete
+# Remove content of $HOME
+# '-mindepth=1' ensures that we do not remove the home directory itself.
+RUN find /home/${NB_USER}/ -mindepth 1 -delete
 
-WORKDIR "/home/${NB_USER}"
-RUN fix-permissions "/home/${NB_USER}"
 USER ${NB_USER}
+WORKDIR "/home/${NB_USER}"

From 475635ec759ff41702c82e8cee948f343f86ad87 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 23 Jul 2024 16:04:47 +0100
Subject: [PATCH 16/23] Smaller diff

---
 Dockerfile | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 76afd9f65..fcb8da2df 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,23 +1,23 @@
 # syntax=docker/dockerfile:1
 ARG FULL_STACK_VER=2024.1021
-ARG QE_VER=7.2
 ARG UV_VER=0.2.27
-
-ARG UV_CACHE_DIR=/tmp/uv_cache
+ARG QE_VER=7.2
 ARG QE_DIR=/opt/conda/envs/quantum-espresso-${QE_VER}
+ARG UV_CACHE_DIR=/tmp/uv_cache
+
 FROM ghcr.io/astral-sh/uv:${UV_VER} AS uv
 
 # STAGE 1
 # Install QE into conda environment in /opt/conda
-# We need to do this first, otherwise QE gets installed into home folder as part of
-# python -m aiidalab_qe install-qe
+# 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
 
 USER ${NB_USER}
-RUN mamba create -p ${QE_DIR} --yes \
-    qe=${QE_VER} && \
+RUN mamba create -p ${QE_DIR} --yes qe=${QE_VER} && \
     mamba clean --all -f -y
 
 # STAGE 2

From 2804752e1142707732a75254328c936f45769c7a Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 23 Jul 2024 16:08:04 +0100
Subject: [PATCH 17/23] Fix

---
 Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index fcb8da2df..c872149f1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -80,7 +80,7 @@ RUN git clean -dffx || true
 # 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=home_build,source=${UV_CACHE_DIR},target=${UV_CACHE_DIR},rw \
     uv pip install --strict --system --compile-bytecode --cache-dir=${UV_CACHE_DIR} . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 

From e9cd720303ae331ffc65b260a87871b2e7f7226e Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 23 Jul 2024 16:36:51 +0100
Subject: [PATCH 18/23] Separate stage for build_deps

---
 Dockerfile | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index c872149f1..e514fe0fc 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -22,10 +22,7 @@ RUN mamba create -p ${QE_DIR} --yes qe=${QE_VER} && \
 
 # STAGE 2
 # - Install python dependencies needed to run aiidalab_qe CLI commands
-# - Prepare AiiDA profile and localhost computer
-# - Install QE codes and pseudopotentials
-# - Archive home folder
-FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER} AS home_build
+FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER} AS build_deps
 ARG UV_CACHE_DIR
 ARG QE_DIR
 
@@ -42,6 +39,12 @@ ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 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}
@@ -67,9 +70,6 @@ ARG QE_DIR
 ARG UV_CACHE_DIR
 USER ${NB_USER}
 
-COPY --from=qe_conda_env ${QE_DIR} ${QE_DIR}
-COPY --from=home_build /opt/conda/home.tar /opt/conda/home.tar
-
 ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
 WORKDIR "${QE_APP_FOLDER}"
 COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
@@ -80,10 +80,13 @@ RUN git clean -dffx || true
 # Use uv cache from the previous build step
 ENV UV_CONSTRAINT=${PIP_CONSTRAINT}
 RUN --mount=from=uv,source=/uv,target=/bin/uv \
-    --mount=from=home_build,source=${UV_CACHE_DIR},target=${UV_CACHE_DIR},rw \
+    --mount=from=build_deps,source=${UV_CACHE_DIR},target=${UV_CACHE_DIR},rw \
     uv pip install --strict --system --compile-bytecode --cache-dir=${UV_CACHE_DIR} . && \
     rm -rf build/ src/aiidalab_qe.egg-info/
 
+COPY --from=qe_conda_env ${QE_DIR} ${QE_DIR}
+COPY --from=home_build /opt/conda/home.tar /opt/conda/home.tar
+
 USER root
 COPY ./before-notebook.d/* /usr/local/bin/before-notebook.d/
 RUN fix-permissions "${CONDA_DIR}"

From 9a18ab707971913735af49dc4ab6c2147bfba80f Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 23 Jul 2024 17:09:40 +0100
Subject: [PATCH 19/23] Better caching behaviour

---
 Dockerfile                         | 38 +++++++++++++++++-------------
 before-notebook.d/00_untar_home.sh |  1 -
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index e514fe0fc..2e7ad19d5 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,7 +3,9 @@ ARG FULL_STACK_VER=2024.1021
 ARG UV_VER=0.2.27
 ARG QE_VER=7.2
 ARG QE_DIR=/opt/conda/envs/quantum-espresso-${QE_VER}
+
 ARG UV_CACHE_DIR=/tmp/uv_cache
+ARG QE_APP_SRC=/tmp/quantum-espresso
 
 FROM ghcr.io/astral-sh/uv:${UV_VER} AS uv
 
@@ -21,16 +23,16 @@ RUN mamba create -p ${QE_DIR} --yes qe=${QE_VER} && \
     mamba clean --all -f -y
 
 # STAGE 2
-# - Install python dependencies needed to run aiidalab_qe CLI commands
+# 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 UV_CACHE_DIR
 ARG QE_DIR
+ARG UV_CACHE_DIR
+ARG QE_APP_SRC
 
-ENV QE_APP_FOLDER=/tmp/quantum-espresso
-WORKDIR ${QE_APP_FOLDER}
-
-COPY --chown=${NB_UID}:${NB_GID} src/ ${QE_APP_FOLDER}/src
-COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml *yaml README.md ${QE_APP_FOLDER}
+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 *yaml README.md ${QE_APP_SRC}
 
 # 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
@@ -67,33 +69,35 @@ RUN --mount=from=qe_conda_env,source=${QE_DIR},target=${QE_DIR} \
 # - Remove all content of home folder
 FROM ghcr.io/aiidalab/full-stack:${FULL_STACK_VER}
 ARG QE_DIR
+ARG QE_APP_SRC
 ARG UV_CACHE_DIR
 USER ${NB_USER}
 
-ENV QE_APP_FOLDER=/opt/conda/quantum-espresso
-WORKDIR "${QE_APP_FOLDER}"
-COPY --chown=${NB_UID}:${NB_GID} . ${QE_APP_FOLDER}
-# Remove all untracked files and directories.
-RUN git clean -dffx || true
-
+WORKDIR /tmp
 # 3.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 \
-    uv pip install --strict --system --compile-bytecode --cache-dir=${UV_CACHE_DIR} . && \
-    rm -rf build/ src/aiidalab_qe.egg-info/
+    --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}
-COPY --from=home_build /opt/conda/home.tar /opt/conda/home.tar
 
 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}"
diff --git a/before-notebook.d/00_untar_home.sh b/before-notebook.d/00_untar_home.sh
index 62aa23b1b..d911474e8 100644
--- a/before-notebook.d/00_untar_home.sh
+++ b/before-notebook.d/00_untar_home.sh
@@ -2,7 +2,6 @@
 set -eux
 
 home="/home/${NB_USER}"
-HOME_TAR="/opt/conda/home.tar"
 
 # Untar home archive file to restore home directory if it is empty
 if [[ $(ls -A ${home} | wc -l) = "0" ]]; then

From 2fd78f8e8c1fbfb561737a24899f99c8eef8553b Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 23 Jul 2024 17:23:08 +0100
Subject: [PATCH 20/23] Tweaks

---
 Dockerfile | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 2e7ad19d5..2785c947e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -32,7 +32,7 @@ ARG QE_APP_SRC
 
 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 *yaml README.md ${QE_APP_SRC}
+COPY --chown=${NB_UID}:${NB_GID} setup.cfg pyproject.toml LICENSE README.md ${QE_APP_SRC}
 
 # 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
@@ -62,11 +62,11 @@ RUN --mount=from=qe_conda_env,source=${QE_DIR},target=${QE_DIR} \
     cd /home/${NB_USER} && tar -cf /opt/conda/home.tar .
 
 # STAGE 3 - Final stage
-# - Copy QE env environment
-# - Copy home folder archive
-# - Copy the whole repo content into the container
 # - 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
@@ -74,7 +74,7 @@ ARG UV_CACHE_DIR
 USER ${NB_USER}
 
 WORKDIR /tmp
-# 3.Install python dependencies
+# 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 \

From 456c4a2e58206019216300817711a57ce6c596c9 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Tue, 23 Jul 2024 17:23:41 +0100
Subject: [PATCH 21/23] Min cache

---
 .github/workflows/docker-build-test-upload.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docker-build-test-upload.yml b/.github/workflows/docker-build-test-upload.yml
index 521f640ad..f8f0e3b55 100644
--- a/.github/workflows/docker-build-test-upload.yml
+++ b/.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 }}
 

From 21a1917030c0a1cc689db4d8ae961778d41181d5 Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Fri, 19 Jul 2024 11:57:15 +0100
Subject: [PATCH 22/23] Fix integration tests on PRs from forks

---
 .github/workflows/docker-build-test-upload.yml | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/docker-build-test-upload.yml b/.github/workflows/docker-build-test-upload.yml
index f8f0e3b55..d9d691b1b 100644
--- a/.github/workflows/docker-build-test-upload.yml
+++ b/.github/workflows/docker-build-test-upload.yml
@@ -80,7 +80,18 @@ jobs:
             - name: Run pytest for Chrome
               run: pytest -sv --driver Chrome tests_integration/
               env:
-                  QE_IMAGE: ${{ env.IMAGE }}@${{ steps.build-upload.outputs.digest }}
+                  # We'd like to identify the image by its unique digest, i.e ghcr.io/aiidalab/qe@sha256:<digest>
+                  # but that sadly does not work when the image is loaded to Docker locally and not published on ghcr.io
+                  # as is the case for PRs from forks. Hence this super-ugly ternary expression...
+                  # For forks, we take the image as ghcr.io/aiidalab/qe:pr-XXX
+                  # which is stored in the steps.meta.outputs.tags variable
+                  QE_IMAGE: >-
+                      ${{
+                          github.event_name == 'pull_request' &&
+                          github.event.pull_request.head.repo.fork &&
+                          steps.meta.outputs.tags ||
+                          format('{0}@{1}', env.IMAGE, steps.build-upload.outputs.imageid)
+                      }}
 
             - name: Upload screenshots as artifacts
               if: always()

From bf57f634d1fa6e82a327feea5a2033448618c45e Mon Sep 17 00:00:00 2001
From: Daniel Hollas <daniel.hollas@bristol.ac.uk>
Date: Thu, 25 Jul 2024 15:48:32 +0100
Subject: [PATCH 23/23] Revert "Fix integration tests on PRs from forks"

This reverts commit 21a1917030c0a1cc689db4d8ae961778d41181d5.
---
 .github/workflows/docker-build-test-upload.yml | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/.github/workflows/docker-build-test-upload.yml b/.github/workflows/docker-build-test-upload.yml
index d9d691b1b..f8f0e3b55 100644
--- a/.github/workflows/docker-build-test-upload.yml
+++ b/.github/workflows/docker-build-test-upload.yml
@@ -80,18 +80,7 @@ jobs:
             - name: Run pytest for Chrome
               run: pytest -sv --driver Chrome tests_integration/
               env:
-                  # We'd like to identify the image by its unique digest, i.e ghcr.io/aiidalab/qe@sha256:<digest>
-                  # but that sadly does not work when the image is loaded to Docker locally and not published on ghcr.io
-                  # as is the case for PRs from forks. Hence this super-ugly ternary expression...
-                  # For forks, we take the image as ghcr.io/aiidalab/qe:pr-XXX
-                  # which is stored in the steps.meta.outputs.tags variable
-                  QE_IMAGE: >-
-                      ${{
-                          github.event_name == 'pull_request' &&
-                          github.event.pull_request.head.repo.fork &&
-                          steps.meta.outputs.tags ||
-                          format('{0}@{1}', env.IMAGE, steps.build-upload.outputs.imageid)
-                      }}
+                  QE_IMAGE: ${{ env.IMAGE }}@${{ steps.build-upload.outputs.digest }}
 
             - name: Upload screenshots as artifacts
               if: always()