Attempt to add cuda support to distributed ci pipeline

Author: msimberg

ci/distributed.yml

@@ -80,7 +80,8 @@ build_distributed_cpu:
   parallel:
     matrix:
       - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, common]
-        BACKEND: [embedded, gtfn_cpu, dace_cpu]
+        # BACKEND: [embedded, gtfn_cpu, dace_cpu, dace_gpu]
+        BACKEND: [dace_cpu, dace_gpu]
   rules:
     - if: $COMPONENT == 'atmosphere/diffusion'
       variables:

ci/docker/base_mpi.Dockerfile

@@ -5,22 +5,23 @@ ENV LC_ALL C.UTF-8
 
 ARG DEBIAN_FRONTEND=noninteractive
 RUN apt-get update -qq && apt-get install -qq -y --no-install-recommends \
-    strace \
+    autoconf \
+    automake \
     build-essential \
-    tar \
-    wget \
+    ca-certificates \
     curl \
+    git && \
     libboost-dev \
     libnuma-dev \
     libopenmpi-dev \
-    ca-certificates \
+    libreadline-dev \
     libssl-dev \
-    autoconf \
-    automake \
     libtool \
+    nvidia-cuda-toolkit \
     pkg-config \
-    libreadline-dev \
-    git && \
+    strace \
+    tar \
+    wget \
     rm -rf /var/lib/apt/lists/*
 
 # Install uv: https://docs.astral.sh/uv/guides/integration/docker

ci/docker/checkout_mpi.Dockerfile

@@ -8,4 +8,4 @@ ARG PYVERSION
 ARG VENV
 ENV UV_PROJECT_ENVIRONMENT=$VENV
 ENV MPI4PY_BUILD_BACKEND="scikit-build-core"
-RUN uv sync --extra distributed --python=$PYVERSION
+RUN uv sync --extra all --python=$PYVERSION

scripts/ci-mpi-wrapper.sh

@@ -17,6 +17,8 @@ else
     exit 1
 fi
 
+export CUDA_VISIBLE_DEVICES="${rank}"
+
 log_file="${CI_PROJECT_DIR:+${CI_PROJECT_DIR}/}pytest-log-rank-${rank}.txt"
 
 if [[ "${rank}" -eq 0 ]]; then