Run MPI tests in CI with CPU backends

.github/workflows/mandatory_and_optional_test_reminder.yml

@@ -28,6 +28,10 @@ jobs:
 
             * `cscs-ci run dace`
 
+            To run tests with MPI you can use:
+
+            * `cscs-ci run distributed`
+
             To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:
              * `cscs-ci run extra`
 

ci/distributed.yml

@@ -0,0 +1,103 @@
+include:
+  - remote: 'https://gitlab.com/cscs-ci/recipes/-/raw/master/templates/v2/.ci-ext.yml'
+
+stages:
+  - baseimage
+  - image
+  - build
+  - test
+  - benchmark
+
+variables:
+  PYVERSION_PREFIX: py310
+  PYVERSION: 3.10.9
+
+# Base image build step with SHA256 checksum for caching
+.build_distributed_baseimage:
+  stage: baseimage
+  before_script:
+    # include build arguments in hash since we use a parameterized Docker file
+    - DOCKER_TAG=`echo "$(cat $DOCKERFILE) $DOCKER_BUILD_ARGS" | sha256sum | head -c 16`
+    - export PERSIST_IMAGE_NAME=$CSCS_REGISTRY_PATH/public/$ARCH/base/icon4py:$DOCKER_TAG-$PYVERSION-mpi
+    - echo "BASE_IMAGE_${PYVERSION_PREFIX}=$PERSIST_IMAGE_NAME" >> build.env
+  artifacts:
+    reports:
+      dotenv: build.env
+  variables:
+    DOCKERFILE: ci/docker/base_mpi.Dockerfile
+    # change to 'always' if you want to rebuild, even if target tag exists already (if-not-exists is the default, i.e. we could also skip the variable)
+    CSCS_REBUILD_POLICY: if-not-exists
+
+build_distributed_baseimage_aarch64:
+  extends: [.container-builder-cscs-gh200, .build_distributed_baseimage]
+  variables:
+    DOCKER_BUILD_ARGS: '["ARCH=$ARCH", "PYVERSION=$PYVERSION"]'
+
+.build_distributed_template:
+  variables:
+    DOCKERFILE: ci/docker/checkout_mpi.Dockerfile
+    # Unique image name based on commit SHA,
+    DOCKER_BUILD_ARGS: '["PYVERSION=$PYVERSION", "BASE_IMAGE=${BASE_IMAGE_${PYVERSION_PREFIX}}", "VENV=${UV_PROJECT_ENVIRONMENT}"]'
+    PERSIST_IMAGE_NAME: $CSCS_REGISTRY_PATH/public/$ARCH/icon4py/icon4py-ci:$CI_COMMIT_SHA-$UV_PROJECT_ENVIRONMENT-$PYVERSION-mpi
+    USE_MPI: NO
+    SLURM_MPI_TYPE: pmix
+    PMIX_MCA_psec: native
+    PMIX_MCA_gds: "^shmem2"
+
+.build_distributed_cpu:
+  extends: [.build_distributed_template]
+  variables:
+    UV_PROJECT_ENVIRONMENT: venv_dist
+
+build_distributed_cpu:
+  stage: image
+  extends: [.container-builder-cscs-gh200, .build_distributed_cpu]
+  needs: [build_distributed_baseimage_aarch64]
+
+.test_template_distributed:
+  timeout: 8h
+  image: $CSCS_REGISTRY_PATH/public/$ARCH/icon4py/icon4py-ci:$CI_COMMIT_SHA-$UV_PROJECT_ENVIRONMENT-$PYVERSION-mpi
+  extends: [.container-runner-santis-gh200, .build_distributed_cpu]
+  needs: [build_distributed_cpu]
+  variables:
+    SLURM_JOB_NUM_NODES: 1
+    SLURM_CPU_BIND: 'verbose'
+    SLURM_NTASKS: 4
+    TEST_DATA_PATH: "/icon4py/testdata"
+    ICON4PY_ENABLE_GRID_DOWNLOAD: false
+    ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false
+    CSCS_ADDITIONAL_MOUNTS: '["/capstor/store/cscs/userlab/d126/icon4py/ci/testdata_003:$TEST_DATA_PATH"]'
+
+.test_distributed_aarch64:
+  stage: test
+  extends: [.test_template_distributed]
+  before_script:
+    - cd /icon4py
+    - echo "using virtual environment at ${UV_PROJECT_ENVIRONMENT}"
+    - source ${UV_PROJECT_ENVIRONMENT}/bin/activate
+    - echo "running with $(python --version)"
+  script:
+    - scripts/ci-mpi-wrapper.sh pytest -sv -k mpi_tests --with-mpi --backend=$BACKEND model/$COMPONENT
+  parallel:
+    matrix:
+      - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, common]
+        BACKEND: [embedded, gtfn_cpu, dace_cpu]
+  rules:
+    - if: $COMPONENT == 'atmosphere/diffusion'
+      variables:
+        SLURM_TIMELIMIT: '00:05:00'
+    - if: $COMPONENT == 'atmosphere/dycore' && $BACKEND == 'dace_cpu'
+      variables:
+        SLURM_TIMELIMIT: '00:20:00'
+    - if: $COMPONENT == 'atmosphere/dycore'
+      variables:
+        SLURM_TIMELIMIT: '00:15:00'
+    - when: on_success
+      variables:
+        SLURM_TIMELIMIT: '00:30:00'
+  artifacts:
+    paths:
+      - pytest-log-rank-*.txt
+
+test_model_distributed:
+  extends: [.test_distributed_aarch64]

ci/docker/base_mpi.Dockerfile

@@ -0,0 +1,27 @@
+FROM ubuntu:25.04
+
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
+
+ARG DEBIAN_FRONTEND=noninteractive
+RUN apt-get update -qq && apt-get install -qq -y --no-install-recommends \
+    strace \
+    build-essential \
+    tar \
+    wget \
+    curl \
+    libboost-dev \
+    libnuma-dev \
+    libopenmpi-dev\
+    ca-certificates \
+    libssl-dev \
+    autoconf \
+    automake \
+    libtool \
+    pkg-config \
+    libreadline-dev \
+    git && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install uv: https://docs.astral.sh/uv/guides/integration/docker
+COPY --from=ghcr.io/astral-sh/uv:0.9.24@sha256:816fdce3387ed2142e37d2e56e1b1b97ccc1ea87731ba199dc8a25c04e4997c5 /uv /uvx /bin/

ci/docker/checkout_mpi.Dockerfile

@@ -0,0 +1,11 @@
+ARG BASE_IMAGE
+FROM $BASE_IMAGE
+
+COPY . /icon4py
+WORKDIR /icon4py
+
+ARG PYVERSION
+ARG VENV
+ENV UV_PROJECT_ENVIRONMENT=$VENV
+ENV MPI4PY_BUILD_BACKEND="scikit-build-core"
+RUN uv sync --extra distributed --python=$PYVERSION

model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py

@@ -22,6 +22,7 @@
 
 
 @pytest.mark.mpi
+@pytest.mark.uses_concat_where
 @pytest.mark.parametrize(
     "experiment, step_date_init, step_date_exit",
     [
@@ -147,6 +148,7 @@ def test_parallel_diffusion(
     )
 
 
+@pytest.mark.skip("SKIP: orchestration is currently broken on CI")
 @pytest.mark.mpi
 @pytest.mark.parametrize(
     "experiment, step_date_init, step_date_exit",

model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py

@@ -62,6 +62,9 @@ def test_run_solve_nonhydro_single_step(
     decomposition_info: definitions.DecompositionInfo,  # : F811 fixture
     backend: gtx_typing.Backend | None,
 ) -> None:
+    if test_utils.is_embedded(backend):
+        pytest.xfail("ValueError: axes don't match array")
+
     parallel_helpers.check_comm_size(processor_props)
     print(
         f"rank={processor_props.rank}/{processor_props.comm_size}: inializing dycore for experiment 'mch_ch_r04_b09_dsl"

model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py

@@ -28,7 +28,7 @@
 from icon4py.model.common import dimension as dims
 from icon4py.model.common.decomposition import definitions, mpi_decomposition
 from icon4py.model.testing import definitions as test_defs, serialbox
-from icon4py.model.testing.parallel_helpers import check_comm_size, processor_props
+from icon4py.model.testing.parallel_helpers import check_comm_size
 
 from ...fixtures import (
     backend,
@@ -40,16 +40,17 @@
     icon_grid,
     interpolation_savepoint,
     metrics_savepoint,
+    processor_props,
     ranked_data_path,
 )
 
 
 """
 running tests with mpi:
 
-mpirun -np 2 python -m pytest -v --with-mpi tests/mpi_tests/test_parallel_setup.py
+mpirun -np 2 python -m pytest -v --with-mpi tests/mpi_tests/test_mpi_decomposition.py
 
-mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/
+mpirun -np 2 pytest -v --with-mpi -k mpi_tests/
 
 
 """
@@ -58,6 +59,7 @@
 @pytest.mark.parametrize("processor_props", [True], indirect=True)
 def test_props(processor_props: definitions.ProcessProperties) -> None:
     assert processor_props.comm
+    assert processor_props.comm_size > 1
 
 
 @pytest.mark.mpi(min_size=2)
@@ -257,7 +259,7 @@ def test_exchange_on_dummy_data(
     exchange = definitions.create_exchange(processor_props, decomposition_info)
     grid = grid_savepoint.construct_icon_grid()
 
-    number = processor_props.rank + 10.0
+    number = processor_props.rank + 10
     input_field = data_alloc.constant_field(
         grid,
         number,

model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py

@@ -99,6 +99,7 @@ def test_distributed_geometry_attrs_for_inverse(
     grid_name: str,
     lb_domain: h_grid.Domain,
 ) -> None:
+    pytest.xfail()
     parallel_helpers.check_comm_size(processor_props)
     parallel_helpers.log_process_properties(processor_props)
     parallel_helpers.log_local_field_size(decomposition_info)

model/common/tests/common/grid/mpi_tests/test_parallel_icon.py

@@ -14,6 +14,7 @@
 
 import icon4py.model.common.dimension as dims
 import icon4py.model.common.grid.horizontal as h_grid
+from icon4py.model.common.decomposition import definitions as decomp_defs
 from icon4py.model.testing import definitions as test_defs, parallel_helpers
 
 from ...fixtures import (
@@ -31,12 +32,13 @@
 if TYPE_CHECKING:
     import gt4py.next as gtx
 
-    from icon4py.model.common.decomposition import definitions as decomp_defs
     from icon4py.model.common.grid import base as base_grid
 
 
 try:
     import mpi4py  # type: ignore[import-not-found] # F401:  import mpi4py to check for optional mpi dependency
+
+    from icon4py.model.common.decomposition import mpi_decomposition
 except ImportError:
     pytest.skip("Skipping parallel on single node installation", allow_module_level=True)
 

model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py

@@ -131,6 +131,9 @@ def test_distributed_interpolation_grg(
     decomposition_info: decomposition.DecompositionInfo,
     interpolation_factory_from_savepoint: interpolation_factory.InterpolationFieldsFactory,
 ) -> None:
+    if test_utils.is_dace(backend):
+        pytest.xfail("Segmentation fault with dace backend")
+
     parallel_helpers.check_comm_size(processor_props)
     intp_factory = interpolation_factory_from_savepoint
     field_ref = interpolation_savepoint.geofac_grg()
@@ -204,6 +207,7 @@ def test_distributed_interpolation_rbf(
     intrp_name: str,
     atol: int,
 ) -> None:
+    pytest.xfail()
     parallel_helpers.check_comm_size(processor_props)
     parallel_helpers.log_process_properties(processor_props)
     parallel_helpers.log_local_field_size(decomposition_info)

model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py

@@ -42,6 +42,7 @@
 
 @pytest.mark.datatest
 @pytest.mark.mpi
+@pytest.mark.uses_concat_where
 @pytest.mark.parametrize("processor_props", [True], indirect=True)
 @pytest.mark.parametrize(
     "attrs_name, metrics_name",
@@ -68,6 +69,9 @@ def test_distributed_metrics_attrs(
     metrics_name: str,
     experiment: test_defs.Experiment,
 ) -> None:
+    if attrs_name == attrs.COEFF_GRADEKIN:
+        pytest.xfail()
+
     parallel_helpers.check_comm_size(processor_props)
     parallel_helpers.log_process_properties(processor_props)
     parallel_helpers.log_local_field_size(decomposition_info)
@@ -80,6 +84,7 @@ def test_distributed_metrics_attrs(
 
 @pytest.mark.datatest
 @pytest.mark.mpi
+@pytest.mark.uses_concat_where
 @pytest.mark.parametrize("processor_props", [True], indirect=True)
 @pytest.mark.parametrize(
     "attrs_name, metrics_name",
@@ -151,6 +156,8 @@ def test_distributed_metrics_attrs_no_halo_regional(
     metrics_name: str,
     experiment: test_defs.Experiment,
 ) -> None:
+    if test_utils.is_embedded(backend):
+        pytest.xfail("ValueError: axes don't match array")
     if experiment == test_defs.Experiments.EXCLAIM_APE:
         pytest.skip(f"Fields not computed for {experiment}")
     parallel_helpers.check_comm_size(processor_props)

model/testing/src/icon4py/model/testing/fixtures/datatest.py

@@ -164,6 +164,11 @@ def download_ser_data(
     if "not datatest" in request.config.getoption("-k", ""):
         return
 
+    with_mpi = request.config.getoption("with_mpi", False)
+    if with_mpi and experiment == definitions.Experiments.GAUSS3D:
+        # TODO(msimberg): Fix? Need serialized data.
+        pytest.skip("GAUSS3D experiment does not support MPI tests")
+
     _download_ser_data(processor_props.comm_size, ranked_data_path, experiment)
 
 

model/testing/src/icon4py/model/testing/parallel_helpers.py

@@ -5,6 +5,7 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+
 import logging
 from collections.abc import Iterable
 

scripts/ci-mpi-wrapper.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+
+# Log all output to separate logfiles, stored as artifacts in gitlab. Output to
+# stdout only from rank 0.
+
+set -euo pipefail
+
+# Check a few different possibilities for the rank.
+if [[ ! -z "${PMI_RANK:-}" ]]; then
+    rank="${PMI_RANK}"
+elif [[ ! -z "${OMPI_COMM_WORLD_RANK:-}" ]]; then
+    rank="${OMPI_COMM_WORLD_RANK}"
+elif [[ ! -z "${SLURM_PROCID:-}" ]]; then
+    rank="${SLURM_PROCID}"
+else
+    echo "Could not determine MPI rank. Set PMI_RANK, OMPI_COMM_WORLD_RANK, or SLURM_PROCID."
+    exit 1
+fi
+
+log_file="${CI_PROJECT_DIR:+${CI_PROJECT_DIR}/}pytest-log-rank-${rank}.txt"
+
+if [[ "${rank}" -eq 0  ]]; then
+    echo "Starting pytest on rank ${rank}, logging to stdout and ${log_file}"
+    $@ |& tee "${log_file}"
+else
+    echo "Starting pytest on rank ${rank}, logging to ${log_file}"
+    $@ >& "${log_file}"
+fi