ARM64 support (#372)

- The rabbitmq for ARM64 is installed from binary tar and starts the service from the same folder as it is installed from conda-forge in x86_64.

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
     build:
 
         runs-on: ubuntu-latest
-        timeout-minutes: 30
+        timeout-minutes: 60
 
         outputs:
             build_vars: ${{ steps.build_vars.outputs.vars }}
@@ -68,8 +68,11 @@ jobs:
                     | tee docker-bake-template-meta.json
             - name: Build images with buildx bake
               id: bake
-              uses: docker/bake-action@v2.3.0
+              uses: docker/bake-action@v3.0.1
               with:
+                  # Using provenance to disable default attestation so it will build only desired images:
+                  # https://github.com/orgs/community/discussions/45969
+                  provenance: false
                   files: |
                       docker-bake.hcl
                       build.json
@@ -90,7 +93,7 @@ jobs:
         needs: build
 
         runs-on: ubuntu-latest
-        timeout-minutes: 10
+        timeout-minutes: 20
 
         steps:
             - uses: actions/checkout@v3
@@ -107,12 +110,42 @@ jobs:
             - name: Run tests
               shell: bash -l {0}  # required to activate the conda environment
               env: ${{ fromJSON(needs.build.outputs.images) }}
-              run: pytest -v -m "not integration"
+              run: |
+                  pytest -v -m "not integration"
+
+    test-arm64:
+        needs: build
+
+        runs-on: buildjet-2vcpu-ubuntu-2204-arm
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v2
+            - name: Login to GitHub Container Registry
+              uses: docker/login-action@v2
+              with:
+                  registry: ghcr.io
+                  username: ${{ github.actor }}
+                  password: ${{ secrets.GITHUB_TOKEN }}
+            - name: Setup Python test environment
+              uses: conda-incubator/setup-miniconda@v2
+              with:
+                  activate-environment: aiidalab-docker-stack
+                  architecture: ARM64
+                  miniconda-version: "latest"
+                  environment-file: environment.yml
+                  python-version: 3.9
+                  auto-activate-base: true
+            - name: Run tests
+              shell: bash -l {0}  # required to activate the conda environment
+              env: ${{ fromJSON(needs.build.outputs.images) }}
+              run: |
+                  pytest -v -m "not integration"
 
     integration-test:
         needs: build
         runs-on: ubuntu-latest
-        timeout-minutes: 10
+        timeout-minutes: 20
 
         steps:
             - uses: actions/checkout@v3
@@ -131,6 +164,33 @@ jobs:
               env: ${{ fromJSON(needs.build.outputs.images) }}
               run: pytest -v -m "integration"
 
+    integration-test-arm64:
+        needs: build
+        runs-on: buildjet-2vcpu-ubuntu-2204-arm
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v2
+            - name: Login to GitHub Container Registry
+              uses: docker/login-action@v2
+              with:
+                  registry: ghcr.io
+                  username: ${{ github.actor }}
+                  password: ${{ secrets.GITHUB_TOKEN }}
+            - name: Setup Python test environment
+              uses: conda-incubator/setup-miniconda@v2
+              with:
+                  activate-environment: aiidalab-docker-stack
+                  architecture: ARM64
+                  miniconda-version: "latest"
+                  environment-file: environment.yml
+                  python-version: 3.9
+                  auto-activate-base: true
+            - name: Run integration tests
+              shell: bash -l {0}  # required to activate the conda environment
+              env: ${{ fromJSON(needs.build.outputs.images) }}
+              run: pytest -v -m "integration"
+
     release:
         if: >-
             github.repository == 'aiidalab/aiidalab-docker-stack'

.pre-commit-config.yaml

@@ -18,11 +18,6 @@ repos:
           - id: yamlfmt
             args: [--preserve-quotes]
 
-    - repo: https://github.com/sirosen/check-jsonschema
-      rev: 0.22.0
-      hooks:
-          - id: check-github-workflows
-
     - repo: https://github.com/psf/black
       rev: 23.3.0
       hooks:

README.md

@@ -59,6 +59,7 @@ To build the images, run `doit build` (tested with *docker buildx* version v0.8.
 The build system will attempt to detect the local architecture and automatically build images for it (tested with amd64 and arm64).
 All commands `build`, `tests`, and `up` will use the locally detected platform and use a version tag based on the state of the local git repository.
 However, you can also specify a custom platform or version with the `--platform` and `--version` parameters, example: `doit build --platform=linux/amd64 --version=my-version`.
+
 You can specify target stacks to build with `--target`, example: `doit build --target base --target full-stack`.
 
 ### Run automated tests
@@ -72,7 +73,7 @@ For manual testing, you can start the images with `doit up`, however we recommen
 
 ### Continuous integration
 
-Images are built for `linux/amd64` during continuous integration for all pull requests into the default branch and pushed to the GitHub Container Registry (ghcr.io) with tags `ghcr.io/aiidalab/*:pr-###`.
+Images are built for `linux/amd64` and `linux/arm64` during continuous integration for all pull requests into the default branch and pushed to the GitHub Container Registry (ghcr.io) with tags `ghcr.io/aiidalab/*:pr-###`.
 You can run automated or manual tests against those images by specifying the registry and version for both the `up` and `tests` commands, example: `doit up --registry=ghcr.io/ --version=pr-123`.
 Note: You may have to [log into the registry first](https://docs.github.com/en/packages/working-with-a-github-packages-registry/working-with-the-container-registry#authenticating-to-the-container-registry).
 

docker-bake.hcl

@@ -30,7 +30,7 @@ variable "REGISTRY" {
 }
 
 variable "PLATFORMS" {
-  default = ["linux/amd64"]
+  default = ["linux/amd64", "linux/arm64"]
 }
 
 variable "TARGETS" {

environment.yml

@@ -3,12 +3,12 @@ name: aiidalab-docker-stack
 channels:
     - conda-forge
 dependencies:
-    - docker-compose=1.29.2
     - doit=0.36.0
     - dunamai=1.13.0
     - pip=22.2.2
     - pytest=7.1.2
     - requests==2.28.1
     - pip:
+          - docker-compose==1.29.2
           - pytest-docker==1.0.0
           - bumpver==2022.1120

stack/base-with-services/Dockerfile

@@ -8,18 +8,51 @@ WORKDIR /opt/
 
 ARG AIIDA_VERSION
 ARG PGSQL_VERSION
+ARG TARGETARCH
 
 RUN mamba create -p /opt/conda/envs/aiida-core-services --yes \
-     aiida-core.services=${AIIDA_VERSION} \
      postgresql=${PGSQL_VERSION} \
-     rabbitmq-server=3.8.14 \
      && mamba clean --all -f -y && \
      fix-permissions "${CONDA_DIR}" && \
      fix-permissions "/home/${NB_USER}"
 
+# Install RabbitMQ in a dedicated conda environment.
+# If the architecture is arm64, we install the default version of rabbitmq provided by the generic binary,
+# # https://www.rabbitmq.com/install-generic-unix.html the version needs to be compatible with system's erlang version.
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+     mamba install -p /opt/conda/envs/aiida-core-services --yes \
+          rabbitmq-server=3.8.14  && \
+     mamba clean --all -f -y && \
+     fix-permissions "${CONDA_DIR}" && \
+     fix-permissions "/home/${NB_USER}"; \
+elif [ "$TARGETARCH" = "arm64" ]; then \
+     apt-get update && apt-get install -y --no-install-recommends  \
+          erlang && \
+     rm -rf /var/lib/apt/lists/* && \
+     apt-get clean all && \
+     export RMQ_VERSION=3.9.13 && \
+     wget -c https://github.com/rabbitmq/rabbitmq-server/releases/download/v${RMQ_VERSION}/rabbitmq-server-generic-unix-${RMQ_VERSION}.tar.xz && \
+     tar -xf rabbitmq-server-generic-unix-${RMQ_VERSION}.tar.xz && \
+     rm rabbitmq-server-generic-unix-${RMQ_VERSION}.tar.xz && \
+     mv rabbitmq_server-${RMQ_VERSION} /opt/conda/envs/aiida-core-services/ && \
+     fix-permissions "/opt/conda/envs/aiida-core-services/rabbitmq_server-${RMQ_VERSION}" && \
+     ln -sf /opt/conda/envs/aiida-core-services/rabbitmq_server-${RMQ_VERSION}/sbin/* /opt/conda/envs/aiida-core-services/bin/; \
+else \
+   echo "Unknown architecture: ${TARGETARCH}."; \
+fi
+
 # Configure AiiDA profile.
 COPY config-quick-setup.yaml .
-COPY before-notebook.d/* /usr/local/bin/before-notebook.d/
+COPY before-notebook.d/20_start-postgresql.sh /usr/local/bin/before-notebook.d/
+COPY before-notebook.d/30_start-rabbitmq-${TARGETARCH}.sh /usr/local/bin/before-notebook.d/
+
+# Supress rabbitmq version warning for arm64 since
+# it is built using latest version rabbitmq from apt install.
+# We explicitly set consumer_timeout to 100 hours in /etc/rabbitmq/rabbitmq.conf
+COPY before-notebook.d/41_suppress-rabbitmq-version-warning.sh /usr/local/bin/before-notebook.d/
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+     rm /usr/local/bin/before-notebook.d/41_suppress-rabbitmq-version-warning.sh; \
+fi
 
 USER ${NB_USER}
 

stack/base-with-services/before-notebook.d/30_start-rabbitmq-arm64.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+set -em
+
+RABBITMQ_DATA_DIR="/home/${NB_USER}/.rabbitmq"
+
+mkdir -p "${RABBITMQ_DATA_DIR}"
+fix-permissions "${RABBITMQ_DATA_DIR}"
+
+# Fix issue where the erlang cookie permissions are corrupted.
+chmod 400 "/home/${NB_USER}/.erlang.cookie" || echo "erlang cookie not created yet."
+
+# Set base directory for RabbitMQ to persist its data. This needs to be set to a folder in the system user's home
+# directory as that is the only folder that is persisted outside of the container.
+export RMQ_VERSION=3.9.13
+RMQ_ETC_DIR="/opt/conda/envs/aiida-core-services/rabbitmq_server-${RMQ_VERSION}/etc/rabbitmq"
+echo MNESIA_BASE="${RABBITMQ_DATA_DIR}" >> "${RMQ_ETC_DIR}/rabbitmq-env.conf"
+echo LOG_BASE="${RABBITMQ_DATA_DIR}/log" >> "${RMQ_ETC_DIR}/rabbitmq-env.conf"
+
+# RabbitMQ with versions >= 3.8.15 have reduced some default timeouts
+# baseimage phusion/baseimage:jammy-1.0.0 running ubuntu 22.04 will install higher version of rabbimq by apt.
+# using workaround from https://github.com/aiidateam/aiida-core/wiki/RabbitMQ-version-to-use
+# set timeout to 100 hours
+echo "consumer_timeout=3600000" >> "${RMQ_ETC_DIR}/rabbitmq.conf"
+
+# Explicitly define the node name. This is necessary because the mnesia subdirectory contains the hostname, which by
+# default is set to the value of $(hostname -s), which for docker containers, will be a random hexadecimal string. Upon
+# restart, this will be different and so the original mnesia folder with the persisted data will not be found. The
+# reason RabbitMQ is built this way is through this way it allows to run multiple nodes on a single machine each with
+# isolated mnesia directories. Since in the AiiDA setup we only need and run a single node, we can simply use localhost.
+echo NODENAME=rabbit@localhost >> "${RMQ_ETC_DIR}/rabbitmq-env.conf"
+
+mamba run -n aiida-core-services rabbitmq-server -detached

stack/base-with-services/before-notebook.d/41_suppress-rabbitmq-version-warning.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+set -em
+
+# Supress rabbitmq version warning for arm64 since
+# it is built using latest version rabbitmq from apt install.
+# We explicitly set consumer_timeout to 100 hours in /etc/rabbitmq/rabbitmq.conf
+verdi config set warnings.rabbitmq_version False

stack/full-stack/Dockerfile

@@ -1,14 +1,26 @@
 # syntax=docker/dockerfile:1
 FROM base-with-services as base
 
-FROM  lab
+FROM lab
 
 USER root
 
 COPY --from=base /opt/config-quick-setup.yaml /opt/
 COPY --from=base "${CONDA_DIR}/envs/aiida-core-services" "${CONDA_DIR}/envs/aiida-core-services"
 COPY --from=base /usr/local/bin/before-notebook.d /usr/local/bin/before-notebook.d
 
+# This is needed because we use multi-stage build.
+# the erlang package is not available after the first stage.
+# After we move base-with-services to a aiida-core repo, we can remove this.
+# Note that it is very important to having the TARGETARCH argument here, otherwise the variable is empty.
+ARG TARGETARCH
+RUN if [ "$TARGETARCH" = "arm64" ]; then \
+    # Install erlang.
+    apt-get update --yes && \
+    apt-get install --yes --no-install-recommends erlang && \
+    apt-get clean && rm -rf /var/lib/apt/lists/*; \
+fi
+
 RUN fix-permissions "${CONDA_DIR}"
 RUN fix-permissions "/home/${NB_USER}/.aiida"
 

tests/conftest.py

@@ -32,7 +32,7 @@ def notebook_service(docker_ip, docker_services):
     port = docker_services.port_for("aiidalab", 8888)
     url = f"http://{docker_ip}:{port}"
     docker_services.wait_until_responsive(
-        timeout=30.0, pause=0.1, check=lambda: is_responsive(url)
+        timeout=60.0, pause=0.1, check=lambda: is_responsive(url)
     )
     return url
 

tests/test_aiidalab.py

@@ -45,6 +45,16 @@ def test_correct_pgsql_version_installed(aiidalab_exec, pgsql_version, variant):
     assert parse(info["version"]).major == parse(pgsql_version).major
 
 
+def test_rabbitmq_can_start(aiidalab_exec, variant):
+    """Test the rabbitmq-server can start, the output should be empty if
+    the command is successful."""
+    if "lab" in variant:
+        pytest.skip()
+    output = aiidalab_exec("mamba run -n aiida-core-services rabbitmq-server -detached")
+
+    assert output == b""
+
+
 def test_correct_aiida_version_installed(aiidalab_exec, aiida_version):
     info = json.loads(
         aiidalab_exec("mamba list --json --full-name aiida-core").decode()