Merge pull request #152 from underworldcode/bugfix/swarm-write-timestep-parallel-151

Fix swarm.write_timestep parallel deadlock (#151)

Author: lmoresi

src/underworld3/swarm.py

@@ -1802,39 +1802,74 @@ def save(
         if filename.endswith(".h5") == False:
             raise RuntimeError("The filename must end with .h5")
 
+        local_data = self.data[:]
+        local_n = local_data.shape[0]
+        n_components = self.num_components
+
         if h5py.h5.get_config().mpi == True and not force_sequential:
+            # BUGFIX(#151): the previous parallel path called
+            #   h5f.create_dataset("data", data=self.data[:])
+            # collectively, but each rank passed its own local-sized array.
+            # In parallel HDF5 every rank must specify the *same* dataset
+            # shape on a collective create_dataset; passing different shapes
+            # leaves HDF5's internal metadata inconsistent so the collective
+            # close never synchronises, producing a silent hang.
+            #
+            # Fix: allgather the per-rank sizes, create the dataset at the
+            # global shape, then each rank writes its own slice.
+            sizes = comm.allgather(local_n)
+            total_n = sum(sizes)
+            offset = sum(sizes[: comm.rank])
+
             with h5py.File(f"{filename[:-3]}.h5", "w", driver="mpio", comm=comm) as h5f:
                 if compression == True:
-                    h5f.create_dataset("data", data=self.data[:], compression=compressionType)
+                    dset = h5f.create_dataset(
+                        "data",
+                        shape=(total_n, n_components),
+                        dtype=local_data.dtype,
+                        chunks=True,
+                        compression=compressionType,
+                    )
                 else:
-                    h5f.create_dataset("data", data=self.data[:])
+                    dset = h5f.create_dataset(
+                        "data",
+                        shape=(total_n, n_components),
+                        dtype=local_data.dtype,
+                    )
+                if local_n > 0:
+                    dset[offset : offset + local_n] = local_data
         else:
+            # Sequential fallback: rank 0 creates the file and writes its slab,
+            # then each higher rank appends in turn. Indentation here matters —
+            # the barrier/loop must be outside the rank-0 branch so all ranks
+            # synchronise (the previous version nested them inside, leaving
+            # higher ranks' data unwritten and rank 0 deadlocked at the
+            # barrier with no peers).
             if comm.rank == 0:
                 with h5py.File(f"{filename[:-3]}.h5", "w") as h5f:
                     if compression == True:
                         h5f.create_dataset(
                             "data",
-                            data=self.data[:],
+                            data=local_data,
                             chunks=True,
-                            maxshape=(None, self.data.shape[1]),
+                            maxshape=(None, n_components),
                             compression=compressionType,
                         )
                     else:
                         h5f.create_dataset(
                             "data",
-                            data=self.data[:],
+                            data=local_data,
                             chunks=True,
-                            maxshape=(None, self.data.shape[1]),
+                            maxshape=(None, n_components),
                         )
-                comm.barrier()
-                for proc in range(1, comm.size):
-                    if comm.rank == proc:
-                        if self.local_size > 0:
-                            with h5py.File(f"{filename[:-3]}.h5", "a") as h5f:
-                                incoming_size = h5f["data"].shape[0]
-                                h5f["data"].resize((h5f["data"].shape[0] + self.local_size), axis=0)
-                                h5f["data"][incoming_size:] = self.data[:, ...]
-                    comm.barrier()
+
+            comm.barrier()
+            for proc in range(1, comm.size):
+                if comm.rank == proc and local_n > 0:
+                    with h5py.File(f"{filename[:-3]}.h5", "a") as h5f:
+                        incoming_size = h5f["data"].shape[0]
+                        h5f["data"].resize((incoming_size + local_n), axis=0)
+                        h5f["data"][incoming_size:] = local_data
                 comm.barrier()
 
         ## Add swarm variable unit metadata to the file
@@ -3626,31 +3661,49 @@ def save(
             warnings.warn("Compression may slow down write times", stacklevel=2)
 
         if h5py.h5.get_config().mpi == True and not force_sequential:
-            # It seems to be a bad idea to mix mpi barriers with the access
-            # context manager so the copy-free version of this seems to hang
-            # when there are many active cores. This is probably why the parallel
-            # h5py write hangs
-
+            # BUGFIX(#151): the previous parallel path called
+            #   h5f.create_dataset("coordinates", data=points_data_copy)
+            # collectively, but each rank passed its own local-sized array.
+            # In parallel HDF5 every rank must specify the *same* dataset
+            # shape on a collective create_dataset; passing different shapes
+            # leaves HDF5's internal metadata inconsistent so the collective
+            # close never synchronises, producing a silent hang.
+            #
+            # Fix: allgather the per-rank sizes, create the dataset at the
+            # global shape, then each rank writes its own slice.
             points_data_copy = self._particle_coordinates.data[:].copy()
+            local_n = points_data_copy.shape[0]
+            cdim = points_data_copy.shape[1]
+            sizes = comm.allgather(local_n)
+            total_n = sum(sizes)
+            offset = sum(sizes[: comm.rank])
 
             with h5py.File(f"{filename[:-3]}.h5", "w", driver="mpio", comm=comm) as h5f:
                 if compression == True:
-                    h5f.create_dataset(
+                    dset = h5f.create_dataset(
                         "coordinates",
-                        data=points_data_copy,
+                        shape=(total_n, cdim),
+                        dtype=points_data_copy.dtype,
+                        chunks=True,
                         compression=compressionType,
                     )
                 else:
-                    h5f.create_dataset("coordinates", data=points_data_copy)
+                    dset = h5f.create_dataset(
+                        "coordinates",
+                        shape=(total_n, cdim),
+                        dtype=points_data_copy.dtype,
+                    )
+                if local_n > 0:
+                    dset[offset : offset + local_n] = points_data_copy
 
             del points_data_copy
 
         else:
-            # It seems to be a bad idea to mix mpi barriers with the access
-            # context manager so the copy-free version of this seems to hang
-            # when there are many active cores
+            # Sequential fallback: rank 0 creates the file and writes its slab,
+            # then each higher rank appends in turn.
 
             points_data_copy = self.points[:].copy()
+            local_n = points_data_copy.shape[0]
 
             if comm.rank == 0:
                 with h5py.File(f"{filename[:-3]}.h5", "w") as h5f:
@@ -3672,17 +3725,16 @@ def save(
 
             comm.barrier()
             for i in range(1, comm.size):
-                if comm.rank == i:
+                if comm.rank == i and local_n > 0:
+                    # BUGFIX(#151): the previous version referenced an undefined
+                    # ``data_copy`` here; passive swarms with a zero-particle
+                    # rank would have raised NameError. Use the local
+                    # ``points_data_copy`` we already have.
                     with h5py.File(f"{filename[:-3]}.h5", "a") as h5f:
-                        h5f["coordinates"].resize(
-                            (h5f["coordinates"].shape[0] + points_data_copy.shape[0]),
-                            axis=0,
-                        )
-                        # passive swarm, zero local particles is not unusual
-                        if data_copy.shape[0] > 0:
-                            h5f["coordinates"][-points_data_copy.shape[0] :] = points_data_copy[:]
+                        existing_size = h5f["coordinates"].shape[0]
+                        h5f["coordinates"].resize((existing_size + local_n), axis=0)
+                        h5f["coordinates"][existing_size:] = points_data_copy
                 comm.barrier()
-            comm.barrier()
 
             del points_data_copy
 

tests/parallel/test_0790_swarm_write_timestep_mpi.py

@@ -0,0 +1,85 @@
+"""MPI regression test for swarm.write_timestep deadlock (issue #151).
+
+The bug: ``Swarm.save()`` and ``SwarmVariable.save()`` parallel paths called
+``h5f.create_dataset(name, data=local_data)`` collectively, but each rank
+passed its own local-sized array. In parallel HDF5 every rank must specify
+the *same* dataset shape on a collective ``create_dataset``; passing
+different shapes leaves HDF5's internal metadata inconsistent so the
+collective close never synchronises, producing a silent hang.
+
+The fix uses ``allgather`` of per-rank sizes, creates the dataset at the
+global shape, then has each rank write its slice.
+
+This test runs bknight1's reproducer from #151 and fails fast via the
+pytest timeout rather than blocking indefinitely.
+"""
+
+import os
+import pytest
+
+import underworld3 as uw
+from petsc4py import PETSc
+
+
+pytestmark = [
+    pytest.mark.level_2,
+    pytest.mark.mpi(min_size=2),
+    pytest.mark.timeout(60),
+]
+
+
+@pytest.mark.mpi(min_size=2)
+def test_swarm_write_timestep_parallel(tmp_path_factory):
+    """swarm.write_timestep() must complete under MPI without deadlock.
+
+    Pre-fix at np>=2 with unequal per-rank particle counts (which is the
+    common case), the parallel HDF5 collective close hung indefinitely.
+    """
+    if uw.mpi.rank == 0:
+        out_dir = tmp_path_factory.mktemp("swarm_io_151")
+    else:
+        out_dir = None
+    out_dir = uw.mpi.comm.bcast(out_dir, root=0)
+    out_dir = str(out_dir)
+
+    mesh = uw.meshing.UnstructuredSimplexBox(
+        minCoords=(-1.0, 0.0),
+        maxCoords=(1.0, 1.0),
+        cellSize=0.2,
+        regular=False,
+        qdegree=3,
+    )
+
+    swarm = uw.swarm.Swarm(mesh=mesh)
+    material = swarm.add_variable(name="material", size=1, dtype=PETSc.IntType)
+    swarm.populate(fill_param=4)
+
+    # Sanity: this is the failing case — local sizes should NOT all be equal.
+    sizes = uw.mpi.comm.allgather(swarm.local_size)
+    assert sum(sizes) > 0, "swarm.populate produced zero global particles"
+
+    # Pre-fix: this hangs forever. The pytest timeout (60s) catches that
+    # rather than letting the run block indefinitely.
+    swarm.write_timestep(
+        filename="swarm",
+        swarmname="swarm",
+        index=0,
+        outputPath=out_dir,
+        swarmVars=[material],
+    )
+
+    # Verify file content
+    expected_h5 = os.path.join(out_dir, "swarm.swarm.00000.h5")
+    expected_var = os.path.join(out_dir, "swarm.swarm.material.00000.h5")
+    expected_xdmf = os.path.join(out_dir, "swarm.swarm.00000.xdmf")
+    if uw.mpi.rank == 0:
+        assert os.path.exists(expected_h5), expected_h5
+        assert os.path.exists(expected_var), expected_var
+        assert os.path.exists(expected_xdmf), expected_xdmf
+
+        import h5py
+        with h5py.File(expected_h5, "r") as f:
+            n_global = f["coordinates"].shape[0]
+        assert n_global == sum(sizes), (
+            f"saved coords shape {n_global} != sum of local sizes {sum(sizes)}"
+        )