free MPI comm after use

veros/distributed.py

@@ -284,12 +284,15 @@ def exchange_cyclic_boundaries(vs, arr):
 def _reduce(vs, arr, op, axis=None):
     if axis is None:
         comm = rs.mpi_comm
+        disconnect_comm = False
     else:
         assert axis in (0, 1)
         pi = proc_rank_to_index(rst.proc_rank)
         other_axis = 1 - axis
         comm = rs.mpi_comm.Split(pi[other_axis], rst.proc_rank)
+        disconnect_comm = True
 
+    try:
         if np.isscalar(arr):
             squeeze = True
             arr = np.array([arr])
@@ -309,6 +312,9 @@ def _reduce(vs, arr, op, axis=None):
             res = res[0]
 
         return res
+    finally:
+        if disconnect_comm:
+            comm.Disconnect()
 
 
 @dist_context_only