MPI Parallelization MVP for Matlab Grid

.github/workflows/build-and-test.yaml

@@ -11,6 +11,7 @@ jobs:
   build-and-test:
     name: Build and test (${{ matrix.variant }})
     runs-on: ubuntu-latest
+    timeout-minutes: 20
     strategy:
       matrix:
         include:
@@ -38,6 +39,9 @@ jobs:
           # several ranks report the same rank number -- which is exactly what the handshake step
           # below checks.
           mpi: openmpi
+      - name: Keep Open MPI's TCP traffic on the loopback interface
+        if: ${{ matrix.run_mpi }}
+        run: echo "btl_tcp_if_include=lo" >> ~/.openmpi/mca-params.conf
       - name: Create virtual environment and install EBFM
         run: |
           python3 -m venv .venv
@@ -65,13 +69,18 @@ jobs:
         run: .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --couple-to-icon-atmo --couple-to-icon-land --fake-coupling --no-shading
       - name: Run example with MPI handshake
         if: ${{ matrix.run_mpi }}
+        timeout-minutes: 5
         run: |
           # Run dummy script first; dummy on rank 0, EBFM on rank 1
           mpirun -np 1 .venv/bin/python tests/core/comm/mpi_handshake_dummy.py first : \
                  -np 1 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --end-time "1-Jan-1979 03:00"
           # Run dummy script second; EBFM on rank 0, dummy on rank 1
           mpirun -np 1 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --end-time "1-Jan-1979 03:00" : \
                  -np 1 .venv/bin/python tests/core/comm/mpi_handshake_dummy.py second
+      - name: Run example on several MPI ranks
+        if: ${{ matrix.run_mpi }}
+        timeout-minutes: 5
+        run: mpirun -np 4 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading
 
   # LFS checkout is expensive (downloads all tracked binary files).
   # Run this stage only after build-and-test has passed to avoid wasting

.github/workflows/test.yml

@@ -68,9 +68,12 @@ jobs:
         # Not the default (mpich): the MPICH package on Ubuntu 24.04 is broken
         # (https://bugs.launchpad.net/ubuntu/+source/mpich/+bug/2072338) and setup-mpi only patches
         # around it by copying newer libmpich12/libucx0 over the apt install. The tests are
-        # single-rank and implementation-agnostic, so use the one that works out of the box.
+        # implementation-agnostic, so use the one that works out of the box.
         mpi: openmpi
 
+    - name: Keep Open MPI's TCP traffic on the loopback interface
+      run: echo "btl_tcp_if_include=lo" >> ~/.openmpi/mca-params.conf
+
     - name: Install tox
       run: |
         python -m pip install --upgrade pip

CHANGELOG.md

@@ -26,6 +26,7 @@ SPDX-License-Identifier: CC-BY-4.0
     * `--dump-reference` now holds the columns of all ranks, ordered by Elmer element id. Every rank previously wrote its own partition to the same file.
     * `--print-diagnostics` now reports the column count and `smb` statistics of the whole domain, logged by rank 0. Each rank previously reported its own partition.
 * In a run on several MPI ranks, only rank 0 logs to the console at `--log-level-console`; the other ranks only log warnings and errors there, instead of repeating every message. Log files (`--log-file`, one per rank) are unchanged. The current model time is logged instead of printed. https://github.com/EBFMorg/EBFM/pull/204
+* A MATLAB grid (`--matlab-mesh`) can now run on several MPI ranks, e.g. `mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading`: each rank computes a contiguous block of the columns, and rank 0 writes the output file, which is identical to that of a run on a single rank. Shading is not supported on several ranks and stops the run with an error. https://github.com/EBFMorg/EBFM/pull/202
 
 # v0.8.0
 

README.md

@@ -225,7 +225,7 @@ On several MPI ranks, `--use-part` is not allowed: rank i computes the partition
 mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --elmer-mesh-crs-epsg 3413
 ```
 
-A partitioned mesh is the only grid EBFM runs on several MPI ranks.
+Besides a partitioned mesh, only a MATLAB grid runs on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)).
 
 Note: The partitioned mesh is not included in this repository.
 
@@ -286,6 +286,21 @@ ebfm --matlab-mesh examples/dem_and_mask.mat --restart-dir examples/restart --re
 
 Note that it is necessary to provide a `--start-time` consistent with the given restart file.
 
+### Running on several MPI ranks
+
+EBFM splits a MATLAB grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`):
+
+```sh
+mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading
+```
+
+* Each rank computes a contiguous block of the columns, the first block on rank 0.
+* Rank 0 writes the output file and `--dump-reference` for all ranks. They are
+  identical to those of a run on a single rank.
+* Not supported on several ranks, and rejected with an error: shading,
+  `--forcing-dir`, `--restart-dir`/`--restart-init`, and grids other than a
+  MATLAB grid or a partitioned Elmer mesh.
+
 ### Performance and Profiling Runs
 
 Multi-threaded Numba kernels (`--with-numba`), GPU offloading (`--with-gpu`), timing, output comparison and

src/ebfm/core/INIT.py

@@ -198,6 +198,7 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm):
         config (GridConfig): Grid configuration object containing settings for grid initialization.
         comm: EBFM's MPI communicator, whose ranks hold the columns of the domain.
     """
+    grid["is_parallel_run"] = comm.size > 1
     grid["is_partitioned"] = config.is_partitioned
     grid["is_unstructured"] = config.is_unstructured
     grid["has_shading"] = config.use_shading
@@ -320,16 +321,20 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm):
             grid["lon_2D"] = input_data["lon"]
             grid["lat_2D"] = input_data["lat"]
 
-        # Create 1-D mask
-        grid["mask"] = mask_2D[mask_2D == 1]
+        # Columns as flat indices into the 2-D fields; in a run on several MPI ranks a contiguous block of them,
+        # the first block on rank 0. The 2-D fields stay whole: the output writer places all columns into them.
+        ind = np.flatnonzero(mask_2D == 1)
+        if config.num_partitions > 1:
+            assert 1 <= config.partition_id <= config.num_partitions, f"Invalid partition id {config.partition_id}."
+            ind = np.array_split(ind, config.num_partitions)[config.partition_id - 1]
+        grid["ind"] = (ind,)
+        grid["column_id"] = ind
 
         # Store 1-D (vectorized) grid information
-        mask_flat = mask_2D.flatten()
-        grid["x"] = grid["x_2D"].flatten()[mask_flat == 1]
-        grid["y"] = grid["y_2D"].flatten()[mask_flat == 1]
-        grid["z"] = grid["z_2D"].flatten()[mask_flat == 1]
-        grid["ind"] = np.where(mask_flat == 1)
-        grid["column_id"] = grid["ind"][0]
+        grid["mask"] = mask_2D.flat[ind]
+        grid["x"] = grid["x_2D"].flat[ind]
+        grid["y"] = grid["y_2D"].flat[ind]
+        grid["z"] = grid["z_2D"].flat[ind]
 
         # ---------------------------------------------------------------------
         # Grid slope and aspect
@@ -351,12 +356,12 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm):
         grid["slope_beta_2D"] = np.arctan(grid["slope"])
 
         # Convert 2-D fields to 1-D vectors
-        grid["slope"] = grid["slope"].flatten()[mask_flat == 1]
-        grid["slope_x"] = grid["slope_x"].flatten()[mask_flat == 1]
-        grid["slope_y"] = grid["slope_y"].flatten()[mask_flat == 1]
-        grid["aspect"] = grid["aspect"].flatten()[mask_flat == 1]
-        grid["lat"] = grid["lat_2D"].flatten()[mask_flat == 1]
-        grid["lon"] = grid["lon_2D"].flatten()[mask_flat == 1]
+        grid["slope"] = grid["slope"].flat[ind]
+        grid["slope_x"] = grid["slope_x"].flat[ind]
+        grid["slope_y"] = grid["slope_y"].flat[ind]
+        grid["aspect"] = grid["aspect"].flat[ind]
+        grid["lat"] = grid["lat_2D"].flat[ind]
+        grid["lon"] = grid["lon_2D"].flat[ind]
 
         # four individual asserts to avoid overhead when turning on optimization
         assert np.isfinite(grid["slope"]).all(), (
@@ -432,6 +437,9 @@ def _precompute_shading_matlab(grid: GridDict, mask_2D) -> None:
                         `az_array` and `maxgridangle`.
     @param[in] mask_2D 2-D glacier mask, 1 for glacier cells.
     """
+    # The walk starts from every glacier cell of the 2-D grid, not only from the columns of this rank.
+    assert not grid["is_parallel_run"], "The shading look-up table cannot be computed in a run on several MPI ranks."
+
     logger.warning("Using EBFM's internal shading method. Should only be used on regular grids!")
 
     grid["shading_method"] = ShadingMethod.LUT  # shading based on look-up table (lut)

src/ebfm/core/LOOP_general_functions.py

@@ -21,6 +21,20 @@ def is_first_time_step(t: int) -> bool:
     return t == 0
 
 
+def is_output_time_step(t: int, io: dict) -> bool:
+    """
+    Check if the current time step is a time step where output has to be written
+
+    Parameters:
+        t (int): Current time step (0-based index).
+        io (dict): Dictionary with I/O settings (e.g. bootfilein, bootfileout, homedir).
+
+    Returns:
+        bool: True if the current time step is output time step, False otherwise.
+    """
+    return (t + 1) % io["freqout"] == 0
+
+
 def is_final_time_step(t: int, time: dict) -> bool:
     """
     Check if the current time step is the final time step.

src/ebfm/core/LOOP_write_to_file.py

@@ -3,14 +3,36 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import os
+
 import numpy as np
 from netCDF4 import Dataset, date2num
 
-from .LOOP_general_functions import is_first_time_step, is_final_time_step
+from .LOOP_general_functions import is_first_time_step, is_output_time_step, is_final_time_step
+from .comm.gather import ColumnGather
 
 _FILL_VALUE = -9999.0
 
 
+def _gather_grid(gather: ColumnGather, grid):
+    """
+    Grid of the columns of all ranks, on the root rank, as far as the output file needs it.
+
+    @returns the grid on the root rank, None on all other ranks
+    """
+    if gather.comm.size == 1:
+        return grid
+
+    assert not grid["is_unstructured"], "Gathering the output of an unstructured grid is not supported."
+
+    # Every rank holds the full 2-D fields. The column ids of a structured grid are the flat indices of the columns
+    # in them, so the sorted ids of all ranks are the positions of all columns.
+    output_grid = {**grid, "ind": gather.column_ids}
+
+    if not gather.is_root:
+        return None
+    return output_grid
+
+
 def _write_unstructured_grid(nc_file, grid):
     """Write unstructured mesh coordinates and topology to an open NetCDF file."""
     n_cells = grid["x"].shape[0]
@@ -140,7 +162,13 @@ def _write_output_variable(nc_file, varname, var_1D, time_index, grid, column):
         nc_file[varname][time_index, :, :] = var_2D
 
 
-def main(OUTFILE, io, OUT, grid, t, time, column):
+def main(OUTFILE, io, OUT, grid, t, time, column, gather: ColumnGather):
+    """
+    Accumulate the output variables of this time step and write them to file every io["freqout"] steps.
+
+    Collective over gather.comm: each rank accumulates its own columns, the root rank gathers the columns of
+    all ranks and writes them to a single file.
+    """
     # Specify variables to be written
     if is_first_time_step(t):
         OUTFILE["varsout"] = [
@@ -235,7 +263,7 @@ def save_binary_files():
         """
 
         # Save output to binary files at the first time step
-        if is_first_time_step(t):
+        if is_first_time_step(t) and gather.is_root:
             if not os.path.exists(io["outdir"]):
                 os.makedirs(io["outdir"])
 
@@ -246,14 +274,18 @@ def save_binary_files():
                 io["fid"][varname] = open(filepath, "wb")
 
         # Write variables to binary files when `freqout` is met
-        if (t + 1) % io["freqout"] == 0:  # TODO: correct?
+        if is_output_time_step(t, io):  # TODO: correct?
+            for entry in OUTFILE["varsout"]:
+                varname = entry[0]
+                OUTFILE[varname] = gather.gather(OUTFILE["TEMP"][varname])
+
+        if is_output_time_step(t, io) and gather.is_root:
             for entry in OUTFILE["varsout"]:
                 varname = entry[0]
-                OUTFILE[varname] = OUTFILE["TEMP"][varname]
                 io["fid"][varname].write(OUTFILE[varname].astype("float32").tobytes())
 
         # Close all file streams and save run metadata at the final time step
-        if is_final_time_step(t, time):
+        if is_final_time_step(t, time) and gather.is_root:
             for file in io["fid"].values():
                 file.close()
 
@@ -275,8 +307,12 @@ def save_netCDF_file():
         time_units = "days since 1970-01-01 00:00:00"
         time_calendar = "gregorian"
 
-        # Initialize NetCDF file at the first time step
+        # The root rank defines the file from the grid of all ranks
         if is_first_time_step(t):
+            io["output_grid"] = _gather_grid(gather, grid)
+
+        # Initialize NetCDF file at the first time step
+        if is_first_time_step(t) and gather.is_root:
             if not os.path.exists(io["outdir"]):
                 os.makedirs(io["outdir"])
 
@@ -285,10 +321,10 @@ def save_netCDF_file():
             io["nc_file"] = Dataset(nc_filepath, "w", format="NETCDF4")
             io["nc_file"].createDimension("time", None)  # Unlimited time dimension
 
-            if grid["is_unstructured"]:
-                _write_unstructured_grid(io["nc_file"], grid)
+            if io["output_grid"]["is_unstructured"]:
+                _write_unstructured_grid(io["nc_file"], io["output_grid"])
             else:
-                _write_structured_grid(io["nc_file"], grid)
+                _write_structured_grid(io["nc_file"], io["output_grid"])
 
             io["nc_file"].createDimension("nl", column.nl)  # Vertical layers for `sub` variables
 
@@ -298,7 +334,7 @@ def save_netCDF_file():
                 var_units = entry[1]
                 var_desc = entry[3]
 
-                dimensions, chunksizes = _output_dimensions(varname, grid, column)
+                dimensions, chunksizes = _output_dimensions(varname, io["output_grid"], column)
 
                 nc_var = io["nc_file"].createVariable(
                     varname=varname,
@@ -321,23 +357,25 @@ def save_netCDF_file():
             nc_time.description = "Time at which data is recorded, in days since 1970-01-01 00:00:00"
 
         # Write data to NetCDF at the specified frequency (e.g., for average/summed values)
-        if (t + 1) % io["freqout"] == 0:  # TODO: correct?
+        if is_output_time_step(t, io):  # TODO: correct?
             time_index = t // io["freqout"]  # Determine the time slice index
 
             # Calculate time in "days since 1970-01-01"
             time_days_since_1970 = date2num(time["TCUR"], units=time_units, calendar=time_calendar)
 
             # Write time variable
-            io["nc_file"]["time"][time_index] = time_days_since_1970
+            if gather.is_root:
+                io["nc_file"]["time"][time_index] = time_days_since_1970
 
-            # Write variables to NetCDF
+            # Write variables to NetCDF, gathering one variable at a time to limit the memory needed on the root rank
             for entry in OUTFILE["varsout"]:
                 varname = entry[0]
-                var_1D = OUTFILE["TEMP"][varname]
-                _write_output_variable(io["nc_file"], varname, var_1D, time_index, grid, column)
+                var_1D = gather.gather(OUTFILE["TEMP"][varname])
+                if gather.is_root:
+                    _write_output_variable(io["nc_file"], varname, var_1D, time_index, io["output_grid"], column)
 
         # Close the NetCDF file at the final time step
-        if is_final_time_step(t, time):
+        if is_final_time_step(t, time) and gather.is_root:
             io["nc_file"].close()
 
         return True

src/ebfm/core/cli.py

@@ -250,7 +250,7 @@ def parse_cli_args(args: list[str] | None = None) -> Namespace:
         "--shading",
         default=argparse.SUPPRESS,
         action=argparse.BooleanOptionalAction,
-        help="Enable/disable shading. " + shading_default_info,
+        help="Enable/disable shading. " + shading_default_info + " Not supported for runs on several MPI ranks.",
     )
 
     example_restart_file_name = INIT.create_restart_file_name(CliDefaults.START_TIME.value)
@@ -530,20 +530,46 @@ def validate_shading_coupling_compat(grid_config, coupling_config) -> None:
         )
 
 
+def validate_shading_parallel_compat(grid_config, comm) -> None:
+    """Check that shading is not enabled in a run on several MPI ranks.
+
+    Must be called on every rank of comm.
+
+    @param[in] grid_config resolved GridConfig
+    @param[in] comm        EBFM's MPI communicator
+    """
+    if comm.size == 1:
+        # Shading works on a single rank
+        return
+
+    if grid_config.use_shading:
+        # The shading look-up table is computed for the columns of the whole grid, not for those of a single rank
+        _cli_error(
+            f"Shading routine not implemented for runs on several MPI ranks ({comm.size} ranks). "
+            "Please deactivate shading via --no-shading or run on a single rank. "
+            "See https://github.com/EBFMorg/EBFM/issues/183"
+        )
+
+
 def validate_grid_parallel_compat(grid_config, comm) -> None:
-    """Check that a run on several MPI ranks uses a partitioned Elmer mesh, of which each rank computes one partition.
+    """Check that a run on several MPI ranks uses a grid that runs there.
+
+    EBFM splits a MATLAB grid over the ranks; of a partitioned Elmer mesh, each rank computes one partition.
 
     Must be called on every rank of comm.
 
     @param[in] grid_config resolved GridConfig
     @param[in] comm        EBFM's MPI communicator
     """
     if comm.size == 1:
+        # Every grid runs on a single rank
         return
 
     if grid_config.grid_type is GridInputType.MATLAB:
-        rejected_grid = "--matlab-mesh"
-    elif grid_config.grid_type is GridInputType.NETCDF:
+        # EBFM splits the columns over the ranks
+        return
+
+    if grid_config.grid_type is GridInputType.NETCDF:
         rejected_grid = "--netcdf-mesh"
     elif not grid_config.is_partitioned:
         rejected_grid = "--elmer-mesh without --is-partitioned-elmer-mesh"
@@ -552,7 +578,7 @@ def validate_grid_parallel_compat(grid_config, comm) -> None:
 
     _cli_error(
         f"{rejected_grid} is not supported for runs on several MPI ranks ({comm.size} ranks). "
-        "Only a partitioned Elmer mesh (--is-partitioned-elmer-mesh) runs on several ranks, one partition per rank. "
+        "Only --matlab-mesh and a partitioned Elmer mesh (--is-partitioned-elmer-mesh) run on several ranks. "
         "Please run on a single rank."
     )