Test mpi versions

.github/workflows/unittests.yml

@@ -2,18 +2,20 @@ name: unittests
 on:
   pull_request:
   push:
-    branches: main
+#    branches: main
 
 jobs:
   test:
-    runs-on: ubuntu-20.04
+    runs-on: ${{ matrix.os }}
     defaults:
       run:
         shell: bash -l {0}
     strategy:
       fail-fast: false
       matrix:
+        os: [ubuntu-18.04, ubuntu-20.04]
         mpi_install: [system, conda]
+        mpi_impl: [openmpi, mpich]
         python: [3.8]
     steps:
       - uses: actions/checkout@v2
@@ -26,13 +28,14 @@ jobs:
           conda --version
           which python
           ./ci/install_mpi.sh
+          ip addr
         env:
           MPI_INSTALL: ${{ matrix.mpi_install }}
+          MPI_IMPL: ${{ matrix.mpi_impl }}
       - name: Run unit tests
         run: |
           pip install -e .[test]
-          echo "localhost slots=16">hostfile
-          pytest --cov-report=xml --cov=dicodile
+          mpiexec -np 1 --mca btl_tcp_if_include eth0 --host localhost:16 python -m pytest --cov-report=xml --cov=dicodile
       - name: Upload coverage
         uses: codecov/codecov-action@v1
         with:

ci/install_mpi.sh

@@ -3,12 +3,23 @@ set -euo pipefail
 
 case "$MPI_INSTALL" in
     "conda")
-        conda install -y openmpi mpi4py
+        conda install -y "$MPI_IMPL" mpi4py
 	;;
     "system")
         sudo apt-get update
-	sudo apt-get install -qy libopenmpi-dev openmpi-bin
+	case "$MPI_IMPL" in
+	    "openmpi")
+		sudo apt-get install -qy libopenmpi-dev openmpi-bin
+		;;
+	    "mpich")
+		sudo apt-get install -qy mpich
+		;;
+	    *)
+		false
+		;;
+	esac
 	;;
     *)
-        false;;
+	false;;
 esac
+

dicodile/_dicodile.py

@@ -227,5 +227,7 @@ def dicodile(X, D_init, reg=.1, n_iter=100, eps=1e-5, window=False,
 
     runtime = np.sum(times)
     encoder.release_workers()
+
+    encoder.shut_down_workers()
     print("[INFO:{}] Finished in {:.0f}s".format(name, runtime))
     return D_hat, z_hat, pobj, times