Support for Multiple Devices and Parallelism via MPI

.github/workflows/mpi_tests.yml

@@ -0,0 +1,141 @@
+name: MPI tests
+
+on:
+  push:
+    branches:
+      - master
+      - dev
+  pull_request:
+    branches:
+      - master
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  mpi_tests:
+    runs-on: ubuntu-latest
+    env:
+      GH_TOKEN: ${{ github.token }}
+
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Filter changes
+        id: changes
+        uses: dorny/paths-filter@v3
+        with:
+          filters: |
+            has_changes:
+              - 'desc/**'
+              - 'tests/test_mpi*'
+              - 'requirements.txt'
+              - 'devtools/dev-requirements.txt'
+              - 'setup.cfg'
+              - '.github/workflows/mpi_tests.yml'
+
+      - name: Check for relevant changes
+        id: check_changes
+        run: echo "has_changes=${{ !contains(github.event.pull_request.labels.*.name, 'only-docs-comments') && steps.changes.outputs.has_changes }}" >> $GITHUB_ENV
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Check full Python version
+        run: |
+          python --version
+          python_version=$(python --version 2>&1 | cut -d' ' -f2)
+          echo "Python version: $python_version"
+          echo "version=$python_version" >> $GITHUB_ENV
+
+      - name: Install MPI (OpenMPI + compiler)
+        if: env.has_changes == 'true'
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libopenmpi-dev openmpi-bin
+
+      - name: Set up virtual environment
+        run: |
+          python -m venv .venv-${{ env.version }}
+          source .venv-${{ env.version }}/bin/activate
+          python -m pip install --upgrade pip
+          pip install -r devtools/dev-requirements.txt
+          pip install matplotlib==3.9.2
+          # install mpi4py after OpenMPI is available
+          pip install mpi4py
+
+      - name: Action Details
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
+          which python
+          python --version
+          mpirun --version
+          pip list
+
+      - name: Test with pytest (MPI setup)
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
+          python -m pytest -v -m mpi_setup\
+            --durations=0 \
+            --cov-report xml:cov.xml \
+            --cov-config=setup.cfg \
+            --cov=desc/ \
+            --db ./prof.db
+
+      - name: Test with pytest (MPI run)
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
+          # ensure each MPI rank writes to a different coverage file
+          mpirun -n 3 --oversubscribe \
+            bash -c 'COVERAGE_FILE=.coverage.$OMPI_COMM_WORLD_RANK \
+                     python -m pytest -v -m mpi_run \
+                       --durations=0 \
+                       --cov=desc/ \
+                       --cov-config=setup.cfg \
+                       --cov-append \
+                       --cov-report='
+
+      - name: Run MPI tutorials
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
+          mpirun -n 2 --oversubscribe python docs/notebooks/tutorials/mpi-tutorials/mpi-eq-solve.py
+          mpirun -n 2 --oversubscribe python docs/notebooks/tutorials/mpi-tutorials/mpi-proximal.py
+
+      - name: Combine coverage files
+        if: always() && env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
+          coverage combine
+          coverage xml -o cov.xml
+
+      - name: save coverage file
+        if: always() && env.has_changes == 'true'
+        uses: actions/upload-artifact@v4
+        with:
+          name: mpi_test_artifact-${{ matrix.python-version }}
+          path: |
+            ./cov.xml
+            ./mpl_results.html
+            ./prof.db
+
+      - name: Upload coverage
+        if: env.has_changes == 'true'
+        id : codecov
+        uses: codecov/codecov-action@v5
+        with:
+          name: codecov-umbrella
+          files: ./cov.xml
+          fail_ci_if_error: true
+          verbose: true

.github/workflows/notebook_tests.yml

@@ -16,13 +16,12 @@ concurrency:
 
 jobs:
   notebook_tests:
-
     runs-on: ubuntu-latest
     env:
       GH_TOKEN: ${{ github.token }}
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ["3.10"]
         group: [1, 2, 3]
 
     steps:
@@ -45,7 +44,6 @@ jobs:
         id: check_changes
         run: echo "has_changes=${{ !contains(github.event.pull_request.labels.*.name, 'only-docs-comments') && steps.changes.outputs.has_changes }}" >> $GITHUB_ENV
 
-
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v6
         with:
@@ -94,6 +92,7 @@ jobs:
           pytest -v --nbmake "./docs/notebooks" \
             --nbmake-timeout=2000 \
             --ignore=./docs/notebooks/zernike_eval.ipynb \
+            --ignore=./docs/notebooks/tutorials/multi_device.ipynb \
             --splits 3 \
             --group ${{ matrix.group }} \
             --splitting-algorithm least_duration

CHANGELOG.md

@@ -20,6 +20,11 @@ New Features
 - Adds ``desc.objectives.DeflationOperator``, a new objective class which can be used to apply deflation techniques to equilibrium and optimization problems to find multiple local minima or multiple solutions from a single initial point, either by wrapping an existing ``desc.objectives._Objective`` object or by including as an additional penalty or constraint. Also adds a tutorial showing this functionality.
 - Sub-objectives of an `ObjectiveFunction` can now have different `use_jit` values than the `ObjectiveFunction`. These objectives have to be built before building the `ObjectiveFunction`.
 - Adds ``num_neighbors`` parameter to ``CoilSetMinDistance`` that limits the pairwise distance computation to the nearest neighbors per coil, reducing memory useage for large coilsets.
+- Adds initial support for multi-device optimization with MPI. This allows to compute derivatives and costs on multiple devices (GPUs/CPUs), and to split memory usage during these operations across devices. See the [documentation](https://desc-docs.readthedocs.io/en/stable/notebooks/tutorials/multi_device.html) for details. Couple important notes:
+    - MPI is not a default dependency of DESC, so, to use MPI functionality, the users should verify their MPI installation themselves.
+    - Using MPI is recommended only for the cases where you get out-of-memory error. If your problem fits to single GPU memory, it's unlikely that MPI will give speed improvement.
+    - MPI is not implemented for matrix decompositions (i.e. QR/SVD/Cholesky) which default optimizer ``lsq-exact`` uses. For the cases where Jacobian doesn't fit to GPU memory, matrix decompositions will be performed on CPU and will be slow. Feel free to open a PR, if you have knowledge on parallel QR/SVD or Cholesky.
+    - CUDA-aware MPI is not supported yet.
 - Method to plot frequency spectrum of inverse stream map in field line coordinates ``Bounce2D.plot_angle_spectrum``.
 - Method to compute bounce integrals in batches is now added to the public API ``Bounce2D.batch``.
 - Initiated deprecation of ``Bounce2D.compute_fieldline_length`` in favor of ``eq.compute("V_psi")``.
@@ -128,7 +133,6 @@ New Features
   - Able to set weights to zero, excluding certain coils from the objective.
 
 
-
 Bug Fixes
 
 - Fixes straight field line equilibrium conversion, see #1880