Enable run without mpi4py (#45)

* add file ddm/mpi.py for MockMPI and MPIWrapper (previously in Struphy, now moved here)

* fix assertions in cart.py

* fix mpi4py imports everywhere

* make mpi4py an optional dependency; run tests without mpi too

* add forgotton change

* change version to 2.5.0 for last release still on GitLab

* install Struphy from branch 467-no-mpi; perform Struphy unit tests on 1 and 2 processes

* checkout correct Struphy branch 467-no-mpi

* correct TODO to devel

* remove the unneccessary if TYPE_CHECKING thing from all files, except mpi.py

* set mpi_enabled to True if mpi4py is installed, TODO: check if it runs under mpirun

* test with updated struphy branch

* test with updated struphy branch

Author: spossann

.github/workflows/testing.yml

@@ -328,7 +328,7 @@ jobs:
       #       f = h5py.File('parallel_test.hdf5', 'w', driver='mpio', comm=MPI.COMM_WORLD)
       #       print(f)"
 
-      - name: Install project
+      - name: Install project without mpi
         run: |
           python -m pip install ".[test]" --no-cache-dir
           python -m pip freeze
@@ -362,15 +362,20 @@ jobs:
         run: |
           python -m pytest -n auto --pyargs psydac -m "not parallel and not petsc"
 
-      - name: Run MPI tests with Pytest
+      - name: Run single-process PETSc tests with Pytest
         working-directory: ./pytest
         run: |
-          python mpi_tester.py --mpirun="mpiexec -n 4 ${MPI_OPTS}" --pyargs psydac -m "parallel and not petsc"
+          python -m pytest -n auto --pyargs psydac -m "not parallel and petsc"
 
-      - name: Run single-process PETSc tests with Pytest
+      - name: Install project with mpi
+        run: |
+          python -m pip install ".[mpi]" 
+          python -m pip freeze
+
+      - name: Run MPI tests with Pytest
         working-directory: ./pytest
         run: |
-          python -m pytest -n auto --pyargs psydac -m "not parallel and petsc"
+          python mpi_tester.py --mpirun="mpiexec -n 4 ${MPI_OPTS}" --pyargs psydac -m "parallel and not petsc"
 
       - name: Run MPI PETSc tests with Pytest
         working-directory: ./pytest
@@ -381,6 +386,7 @@ jobs:
         if: always()
         run: |
           rm -rf pytest
+  
   test_struphy:
     runs-on: ${{ matrix.os }}
     strategy:
@@ -469,7 +475,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
 
-      - name: Install project
+      - name: Install project without mpi
         run: |
           python -m pip install ".[test]" --no-cache-dir
           python -m pip freeze
@@ -479,16 +485,17 @@ jobs:
           echo "PSYDAC_DIR=$GITHUB_WORKSPACE" >> $GITHUB_ENV
           echo "STRUPHY_DIR=$GITHUB_WORKSPACE/struphy" >> $GITHUB_ENV
 
-      - name: Clone struphy from GitLab #TODO: Set branch to devel
+      - name: Clone struphy from GitLab
         run: |
           git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git $STRUPHY_DIR
       
-      - name: Install struphy
+      - name: Install struphy without mpi #TODO: Set branch to devel
         working-directory: ${{ env.STRUPHY_DIR }}
         run: |
           echo "Psydac location for this branch"
           pip show psydac
           pip uninstall psydac -y
+          git checkout 467-no-mpi
           python -m pip install ".[phys]" --no-cache-dir
           echo "Psydac location after installing struphy"
           pip show psydac
@@ -520,6 +527,16 @@ jobs:
         run: |
             struphy test models --fast
 
+      - name: Install mpi4py
+        working-directory: ${{ env.STRUPHY_DIR }}
+        run: |
+            pip install -U mpi4py
+
+      - name: Run struphy unit tests with mpi
+        working-directory: ${{ env.STRUPHY_DIR }}
+        run: |
+            struphy test unit --mpi 2
+
       - name: Remove test directory
         if: always()
         run: |

mpi_tester.py

@@ -101,7 +101,8 @@ def MPITest(commsize):
     def test_stuff(comm):
         pass
     """
-    from mpi4py import MPI
+    from psydac.ddm.mpi import mpi as MPI
+        
     if not isinstance(commsize, (tuple, list)):
         commsize = (commsize,)
 
@@ -182,7 +183,7 @@ def __init__(self):
     #---------------------------------------------------------------------------
     @property
     def comm(self):
-        from mpi4py import MPI
+        from psydac.ddm.mpi import mpi as MPI
         return MPI.COMM_WORLD
 
     #---------------------------------------------------------------------------

psydac/ddm/blocking_data_exchanger.py

@@ -1,7 +1,7 @@
 # coding: utf-8
 
 import numpy as np
-from mpi4py import MPI
+from psydac.ddm.mpi import mpi as MPI
 
 from .cart import CartDecomposition, find_mpi_type
 from .basic import CartDataExchanger

psydac/ddm/cart.py

@@ -3,8 +3,9 @@
 import os
 import numpy as np
 from itertools import product
-from mpi4py    import MPI
 
+from psydac.ddm.mpi import mpi as MPI
+from psydac.ddm.mpi import MockMPI
 from psydac.ddm.partition import compute_dims, partition_procs_per_patch
 
 
@@ -31,11 +32,14 @@ def find_mpi_type( dtype ):
         MPI datatype to be used for communication.
 
     """
-    if isinstance( dtype, MPI.Datatype ):
-        mpi_type = dtype
+    if not isinstance(MPI, MockMPI):
+        if isinstance( dtype, MPI.Datatype ):
+            mpi_type = dtype
+        else:
+            nt = np.dtype( dtype )
+            mpi_type = MPI._typedict[nt.char]
     else:
-        nt = np.dtype( dtype )
-        mpi_type = MPI._typedict[nt.char]
+        mpi_type = np.dtype( dtype )
 
     return mpi_type
 
@@ -63,7 +67,8 @@ class MultiPatchDomainDecomposition:
     def __init__(self, ncells, periods, comm=None, num_threads=None):
 
         assert len( ncells ) == len( periods )
-        if comm is not None:assert isinstance( comm, MPI.Comm )
+        if not isinstance(MPI, MockMPI) and comm is not None:
+            assert isinstance( comm, MPI.Comm )
         num_threads = num_threads if num_threads else int(os.environ.get('OMP_NUM_THREADS', 1))
 
         # Store input arguments
@@ -206,7 +211,12 @@ def __init__(self, ncells, periods, comm=None, global_comm=None, num_threads=Non
         assert len( ncells ) == len( periods )
         assert all( n >=1 for n in ncells )
         assert all( isinstance( period, bool ) for period in periods )
-        if comm is not None: assert isinstance( comm, MPI.Comm )
+        if isinstance(MPI, MockMPI): 
+            comm = None
+        else:
+            if comm is not None:
+                assert isinstance( comm, MPI.Comm )
+        
 
         self._ncells       = tuple ( ncells )
         self._periods      = tuple ( periods )

psydac/ddm/interface_data_exchanger.py

@@ -1,6 +1,6 @@
 # coding: utf-8
 
-from mpi4py import MPI
+from psydac.ddm.mpi import mpi as MPI
 
 from .cart import InterfaceCartDecomposition, find_mpi_type
 

psydac/ddm/mpi.py

@@ -0,0 +1,100 @@
+from dataclasses import dataclass
+from time import time
+from typing import TYPE_CHECKING
+
+
+# Might not be needed
+class MPICommWrapper:
+    def __init__(self, use_mpi=True):
+        self.use_mpi = use_mpi
+        if use_mpi:
+            from mpi4py import MPI
+
+            self.comm = MPI.COMM_WORLD
+        else:
+            self.comm = MockComm()
+
+    def __getattr__(self, name):
+        return getattr(self.comm, name)
+
+
+class MockComm:
+    def __getattr__(self, name):
+        # Return a function that does nothing and returns None
+        def dummy(*args, **kwargs):
+            return None
+
+        return dummy
+
+    # Override some functions
+    def Get_rank(self):
+        return 0
+
+    def Get_size(self):
+        return 1
+
+    def Barrier(self):
+        return
+
+
+class MPIwrapper:
+    def __init__(self, use_mpi: bool = False):
+        self.use_mpi = use_mpi
+        if use_mpi:
+            from mpi4py import MPI
+
+            self._MPI = MPI
+            print("MPI is enabled")
+        else:
+            self._MPI = MockMPI()
+            print("MPI is NOT enabled")
+
+    @property
+    def MPI(self):
+        return self._MPI
+
+
+class MockMPI:
+    def __getattr__(self, name):
+        # Return a function that does nothing and returns None
+        def dummy(*args, **kwargs):
+            return None
+
+        return dummy
+
+    # Override some functions
+    @property
+    def COMM_WORLD(self):
+        return MockComm()
+
+    # def comm_Get_rank(self):
+    #     return 0
+
+    # def comm_Get_size(self):
+    #     return 1
+
+
+try:
+    from mpi4py import MPI
+
+    _comm = MPI.COMM_WORLD
+    rank = _comm.Get_rank()
+    size = _comm.Get_size()
+    mpi_enabled = True
+except ImportError:
+    # mpi4py not installed
+    mpi_enabled = False
+except Exception:
+    # mpi4py installed but not running under mpirun
+    mpi_enabled = False
+
+# TODO: add environment variable for mpi use
+mpi_wrapper = MPIwrapper(use_mpi=mpi_enabled)
+
+# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime.
+if TYPE_CHECKING:
+    from mpi4py import MPI
+
+    mpi = MPI
+else:
+    mpi = mpi_wrapper.MPI

psydac/ddm/nonblocking_data_exchanger.py

@@ -2,8 +2,8 @@
 
 import numpy as np
 from itertools import product
-from mpi4py import MPI
 
+from psydac.ddm.mpi import mpi as MPI
 from .cart import CartDecomposition, find_mpi_type
 from .basic import CartDataExchanger
 

psydac/ddm/tests/test_cart_1d.py

@@ -11,7 +11,7 @@
 def run_cart_1d( data_exchanger_type, verbose=False ):
 
     import numpy as np
-    from mpi4py       import MPI
+    from psydac.ddm.mpi import mpi as MPI
     from psydac.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

psydac/ddm/tests/test_cart_2d.py

@@ -9,7 +9,7 @@
 def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis=None):
 
     import numpy as np
-    from mpi4py       import MPI
+    from psydac.ddm.mpi import mpi as MPI  
     from psydac.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------

psydac/ddm/tests/test_cart_3d.py

@@ -9,7 +9,7 @@
 def run_cart_3d( data_exchanger_type, verbose=False ):
 
     import numpy as np
-    from mpi4py       import MPI
+    from psydac.ddm.mpi import mpi as MPI 
     from psydac.ddm.cart import DomainDecomposition, CartDecomposition
 
     #---------------------------------------------------------------------------