From 04d66bcced21b571aa2fe9b6c1edbc49ae6b8fa3 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 17:18:19 +0200 Subject: [PATCH 1/6] Keep MPI runs in CI from hanging setup-mpi makes Open MPI use TCP between ranks, and on GitHub's runners Open MPI's TCP transport also picks the docker0 interface, over which messages never arrive. Small messages still get through, but large ones hang: in #197 a gather on 4 ranks hung until GitHub's limit of 6 hours, while the same run passed locally, where the ranks use shared memory. The next commits run the MATLAB example on up to 4 ranks, with gathers large enough to hang. Open MPI now only uses the loopback interface in both workflows. It is set in mca-params.conf, since tox does not pass OMPI_MCA_* variables to the tests. In addition, the build-and-test job now stops after 20 minutes and its MPI steps after 5, so a hang fails quickly instead of blocking a runner for 6 hours. Taken from 5e28f4d and 47b4c8f of #197. Assisted-by: Claude Code:claude-opus-5-5 --- .github/workflows/build-and-test.yaml | 5 +++++ .github/workflows/test.yml | 3 +++ 2 files changed, 8 insertions(+) diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml index 516c8226..9745c940 100644 --- a/.github/workflows/build-and-test.yaml +++ b/.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,6 +69,7 @@ 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 : \ diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index b25d90dc..6990f15a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -71,6 +71,9 @@ jobs: # single-rank and 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 From e05b183fe41769998402fd8d295bbdc9f91555b9 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 17:25:05 +0200 Subject: [PATCH 2/6] Select the columns of a gridded grid by their flat index MATLAB and NetCDF grids now keep the flat indices of their columns in the 2-D fields as `ind` and take every per-column field from the 2-D fields with them, instead of applying the mask to each field again. A subset of the columns then only needs a subset of `ind`, which the next commits use to split a MATLAB grid over MPI ranks. The columns, their order and the model output are unchanged. Assisted-by: Claude Code:claude-opus-5-5 --- src/ebfm/core/INIT.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index 0874c15d..8c206bb0 100644 --- a/src/ebfm/core/INIT.py +++ b/src/ebfm/core/INIT.py @@ -320,16 +320,16 @@ 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 + ind = np.flatnonzero(mask_2D == 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 +351,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(), ( From e0dace179f27a22474dd2d48c89a546a6c0bd9aa Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 17:08:32 +0200 Subject: [PATCH 3/6] Write the output of all MPI ranks to a single file on rank 0 The output writer now gathers the columns of all ranks with ColumnGather and rank 0 writes them to the file a single rank would write, one variable at a time to limit its memory. Only structured grids are gathered; a partitioned mesh still skips the output file. Taken from 8b5a9eb of #197. Assisted-by: Claude Code:claude-opus-5-5 --- src/ebfm/core/LOOP_general_functions.py | 14 +++++ src/ebfm/core/LOOP_write_to_file.py | 72 +++++++++++++++++++------ src/ebfm/main.py | 4 +- tests/core/test_loop_write_to_file.py | 9 +++- 4 files changed, 79 insertions(+), 20 deletions(-) diff --git a/src/ebfm/core/LOOP_general_functions.py b/src/ebfm/core/LOOP_general_functions.py index c415366e..576953a6 100644 --- a/src/ebfm/core/LOOP_general_functions.py +++ b/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. diff --git a/src/ebfm/core/LOOP_write_to_file.py b/src/ebfm/core/LOOP_write_to_file.py index 18c110e3..7ea8c3a9 100644 --- a/src/ebfm/core/LOOP_write_to_file.py +++ b/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 diff --git a/src/ebfm/main.py b/src/ebfm/main.py index 193b2998..abffceec 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -229,7 +229,7 @@ def _main_impl(): C = INIT.init_constants() grid = INIT.init_grid(grid, io, grid_config, ebfm_comm) - # Rank 0 collects the columns of all ranks, e.g. for --dump-reference + # Rank 0 collects the columns of all ranks, e.g. to write them to the output file or for --dump-reference column_gather = ColumnGather(ebfm_comm, grid["column_id"]) OUT, IN, OUTFILE = INIT.init_initial_conditions(C, grid, io, time, column, init_with_restart_file=args.restart_init) @@ -386,7 +386,7 @@ def _main_impl(): # Write output to files (for unpartitioned grids only; see the warning emitted before the time loop) # TODO: partitioned output should be supported in a dedicated parallel-safe writer. if not grid["is_partitioned"]: - io, OUTFILE = LOOP_write_to_file.main(OUTFILE, io, OUT, grid, t, time, column) + io, OUTFILE = LOOP_write_to_file.main(OUTFILE, io, OUT, grid, t, time, column, column_gather) else: logger.warning("Skipping writing output to file for partitioned runs.") diff --git a/tests/core/test_loop_write_to_file.py b/tests/core/test_loop_write_to_file.py index 15240812..18e21894 100644 --- a/tests/core/test_loop_write_to_file.py +++ b/tests/core/test_loop_write_to_file.py @@ -14,6 +14,8 @@ from ebfm.core import LOOP_write_to_file from ebfm.core.LOOP_write_to_file import _output_dimensions, main as write_output +from ebfm.core.comm import defaultComm +from ebfm.core.comm.gather import ColumnGather from ebfm.core.config import ColumnDiscretizationConfig NL = 4 @@ -22,6 +24,9 @@ GLACIER_INDICES = (0, 3, 5, 7, 11) GPSUM = len(GLACIER_INDICES) +# Element ids of the unstructured grid. Like the ids of an Elmer mesh, they neither start at 0 nor are contiguous. +ELEMENT_IDS = (3, 5, 12, 27, 40) + _NETCDF_OUTPUT = 2 @@ -94,6 +99,7 @@ def _structured_grid(): "lon_2D": np.zeros((NY, NX)), "lat_2D": np.zeros((NY, NX)), "ind": np.array(GLACIER_INDICES), + "column_id": np.array(GLACIER_INDICES), } @@ -101,6 +107,7 @@ def _unstructured_grid(): return { "is_unstructured": True, "is_partitioned": False, + "column_id": np.array(ELEMENT_IDS), "x": np.arange(GPSUM, dtype=np.float64), "y": np.arange(GPSUM, dtype=np.float64), "z": np.full(GPSUM, 100.0), @@ -131,7 +138,7 @@ def _write_single_step(outdir, grid, extra_out=None): io = {"outdir": str(outdir), "freqout": 1, "output_type": _NETCDF_OUTPUT} time = {"TCUR": datetime(2020, 1, 1), "tn": 1} - io, outfile = write_output({}, io, out, grid, 0, time, column) + io, outfile = write_output({}, io, out, grid, 0, time, column, ColumnGather(defaultComm, grid["column_id"])) return Path(outdir) / "model_output.nc", outfile, out From b19fa9365e0131b75bd94cf608820905f5c35f9a Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 13:01:09 +0200 Subject: [PATCH 4/6] Reject shading in runs on several MPI ranks The shading look-up table is computed by walking from every glacier cell of the 2-D grid, not only from the columns of one rank, so a rank holding part of the columns cannot build its rows of the table. Shading on several ranks is now a usage error asking for --no-shading, and _precompute_shading_matlab asserts that it is not called in such a run. No grid that supports shading runs on several ranks yet; the next commits split MATLAB grids over the ranks. Taken from 43a676e of #197. Assisted-by: Claude Code:claude-opus-5-5 --- src/ebfm/core/INIT.py | 4 ++++ src/ebfm/core/cli.py | 23 ++++++++++++++++++++++- src/ebfm/main.py | 4 ++++ tests/core/test_cli.py | 23 +++++++++++++++++++++++ tests/core/test_init.py | 4 ++++ 5 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index 8c206bb0..c33a08d6 100644 --- a/src/ebfm/core/INIT.py +++ b/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 @@ -432,6 +433,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) diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 2ecdfc19..0aa1db77 100644 --- a/src/ebfm/core/cli.py +++ b/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,6 +530,27 @@ 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. diff --git a/src/ebfm/main.py b/src/ebfm/main.py index abffceec..d08cecbf 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -29,6 +29,7 @@ validate_use_part_parallel_compat, validate_forcing_parallel_compat, validate_restart_parallel_compat, + validate_shading_parallel_compat, ) import ebfm.coupling @@ -218,6 +219,9 @@ def _main_impl(): # Ensure restart files are only read and written in runs on a single rank. validate_restart_parallel_compat(args, ebfm_comm) + # Ensure shading routine is only used in runs on a single rank. + validate_shading_parallel_compat(grid_config, ebfm_comm) + logger.debug("Successfully completed consistency checks.") # Model setup & initialization diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 896d0528..7624e457 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -19,6 +19,7 @@ validate_forcing_parallel_compat, validate_grid_parallel_compat, validate_restart_parallel_compat, + validate_shading_parallel_compat, validate_shading_coupling_compat, validate_use_part_parallel_compat, ) @@ -382,6 +383,28 @@ def test_shading_with_coupling_rejected(self): self.assertEqual(ctx.exception.code, 2) +class TestValidateShadingParallelCompat(unittest.TestCase): + """validate_shading_parallel_compat raises a CLI error when shading is on in a run on several ranks.""" + + def setUp(self): + # Ensure _parser is set by running a valid parse first. + parse_cli_args(_MATLAB) + + def _validate(self, use_shading: bool, comm_size: int): + validate_shading_parallel_compat(MagicMock(use_shading=use_shading), MagicMock(size=comm_size)) + + def test_shading_on_single_rank_ok(self): + self._validate(use_shading=True, comm_size=1) # must not raise + + def test_no_shading_on_several_ranks_ok(self): + self._validate(use_shading=False, comm_size=4) # must not raise + + def test_shading_on_several_ranks_rejected(self): + with self.assertRaises(SystemExit) as ctx: + self._validate(use_shading=True, comm_size=4) + self.assertEqual(ctx.exception.code, 2) + + class TestValidateGridParallelCompat(unittest.TestCase): """validate_grid_parallel_compat raises a CLI error on several ranks for every grid but a partitioned Elmer mesh.""" diff --git a/tests/core/test_init.py b/tests/core/test_init.py index 41d6f5bb..594b20ec 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -371,6 +371,10 @@ def test_shading_enabled_precomputes_the_lookup_table(self): self.assertTrue(grid["has_shading"]) precompute.assert_called_once() + def test_lookup_table_is_rejected_in_a_parallel_run(self): + with self.assertRaises(AssertionError): + INIT._precompute_shading_matlab({"is_parallel_run": True}, mask_2D=None) + class TestReadNetcdfGrid(unittest.TestCase): """Missing orography is NaN unless the user opts in to filling it with zero elevation. From a4fab57199fd31889458248aa2155824deceb15e Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 17:09:31 +0200 Subject: [PATCH 5/6] Split a MATLAB grid over the MPI ranks On several ranks, each rank now computes a contiguous block of the columns of a MATLAB grid, the first block on rank 0. The 2-D fields stay whole on every rank: the slope needs the neighbours of each cell, and the output writer places the columns of all ranks into them. The domain-wide latitude mean, lapse rate and random seed of #203 make the results equal those of a single rank. Taken from 46f47a2 of #197. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + README.md | 17 ++++++++++++++++- src/ebfm/core/INIT.py | 6 +++++- src/ebfm/core/cli.py | 13 +++++++++---- src/ebfm/core/config/grid.py | 14 ++++++++++++-- src/ebfm/main.py | 4 +++- tests/core/test_cli.py | 6 +++--- tests/core/test_init.py | 31 +++++++++++++++++++++++++++++++ 8 files changed, 80 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5bf9239..c597e6bf 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index fdf881e5..7f5f80fd 100644 --- a/README.md +++ b/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 diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index c33a08d6..99433bfa 100644 --- a/src/ebfm/core/INIT.py +++ b/src/ebfm/core/INIT.py @@ -321,8 +321,12 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm): grid["lon_2D"] = input_data["lon"] grid["lat_2D"] = input_data["lat"] - # Columns as flat indices into the 2-D fields + # 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 diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 0aa1db77..284671f7 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -552,7 +552,9 @@ def validate_shading_parallel_compat(grid_config, comm) -> None: 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. @@ -560,11 +562,14 @@ def validate_grid_parallel_compat(grid_config, comm) -> None: @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" @@ -573,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." ) diff --git a/src/ebfm/core/config/grid.py b/src/ebfm/core/config/grid.py index 6f68b78f..f12f2ac7 100644 --- a/src/ebfm/core/config/grid.py +++ b/src/ebfm/core/config/grid.py @@ -26,7 +26,9 @@ class GridConfig: dem_file: Path | None # Path to the DEM file (only relevant for CUSTOM grid type) is_partitioned: bool # Whether the grid is partitioned is_unstructured: bool # Whether the grid is unstructured - partition_id: int | None # Partition ID (only relevant if is_partitioned is True) + partition_id: int | None # Partition ID (only relevant if is_partitioned is True or num_partitions > 1) + # Number of blocks EBFM splits the columns of a MATLAB grid into, one per MPI rank; 1 for a grid that is not split + num_partitions: int elmer_mesh_crs_epsg: int # EPSG code of Elmer mesh coordinates use_shading: bool # Whether to use shading for the grid fill_missing_elevation: bool # Whether to set missing elevations to 0 m (only relevant for NETCDF grid type) @@ -46,11 +48,13 @@ class GridConfig: GridInputType.MATLAB, } - def __init__(self, args: Namespace): + def __init__(self, args: Namespace, num_ranks: int = 1, rank: int = 0): """ Initialize grid configuration from command line arguments. @param[in] args command line arguments + @param[in] num_ranks number of MPI ranks, over which a MATLAB grid is split + @param[in] rank MPI rank of this process, which computes the block rank + 1 of a split grid """ selected_primary_grids = [ grid_type @@ -68,6 +72,7 @@ def __init__(self, args: Namespace): self.is_partitioned = args.is_partitioned_elmer_mesh self.dem_file = None self.partition_id = None + self.num_partitions = 1 assert ( not self.is_partitioned or elmer_mesh is not None @@ -113,6 +118,11 @@ def __init__(self, args: Namespace): ) raise Exception("Invalid grid configuration.") + if self.grid_type is GridInputType.MATLAB and num_ranks > 1: + logger.info(f"Splitting the grid over {num_ranks} MPI ranks...") + self.partition_id = rank + 1 + self.num_partitions = num_ranks + self.use_shading = self._check_shading(args) self.fill_missing_elevation = args.fill_missing_elevation diff --git a/src/ebfm/main.py b/src/ebfm/main.py index d08cecbf..cb3ef726 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -204,7 +204,7 @@ def _main_impl(): logger.debug("Reading configuration and checking for consistency.") # TODO consider introducing an ebfm_adapter_config.yaml to be parsed alternatively/additionally to command line args - grid_config = GridConfig(args) + grid_config = GridConfig(args, ebfm_comm.size, ebfm_comm.rank) # Ensure shading routine is only used in uncoupled runs. # See https://github.com/EBFMorg/EBFM/issues/11 for details. @@ -235,6 +235,8 @@ def _main_impl(): # Rank 0 collects the columns of all ranks, e.g. to write them to the output file or for --dump-reference column_gather = ColumnGather(ebfm_comm, grid["column_id"]) + if ebfm_comm.size > 1 and not grid["is_partitioned"]: + logger.info(f"Rank 0 writes the output of all {ebfm_comm.size} ranks.") OUT, IN, OUTFILE = INIT.init_initial_conditions(C, grid, io, time, column, init_with_restart_file=args.restart_init) diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 7624e457..bb3e41f3 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -406,11 +406,11 @@ def test_shading_on_several_ranks_rejected(self): class TestValidateGridParallelCompat(unittest.TestCase): - """validate_grid_parallel_compat raises a CLI error on several ranks for every grid but a partitioned Elmer mesh.""" + """validate_grid_parallel_compat rejects every grid but MATLAB and a partitioned Elmer mesh on several ranks.""" # name: (grid_type, is_partitioned, grid option named by the error on several ranks, or None if accepted) _CASES = { - "--matlab-mesh": (GridInputType.MATLAB, False, "--matlab-mesh"), + "--matlab-mesh": (GridInputType.MATLAB, False, None), "--netcdf-mesh": (GridInputType.NETCDF, False, "--netcdf-mesh"), "--elmer-mesh": (GridInputType.ELMER, False, "--elmer-mesh"), "--elmer-mesh --netcdf-dem-mesh": (GridInputType.CUSTOM, False, "--elmer-mesh"), @@ -443,7 +443,7 @@ def test_several_ranks(self): self.assertEqual(ctx.exception.code, 2) error = stderr.getvalue().splitlines()[-1] self.assertIn(f"error: {rejected_grid} ", error) - self.assertIn("Only a partitioned Elmer mesh (--is-partitioned-elmer-mesh)", error) + self.assertIn("Only --matlab-mesh and a partitioned Elmer mesh", error) @unittest.skipUnless(mpi_available, "requires MPI (pip install 'ebfm[mpi]')") diff --git a/tests/core/test_init.py b/tests/core/test_init.py index 594b20ec..9381f61c 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -376,6 +376,37 @@ def test_lookup_table_is_rejected_in_a_parallel_run(self): INIT._precompute_shading_matlab({"is_parallel_run": True}, mask_2D=None) +class TestMatlabGridSplit(unittest.TestCase): + """Split over several MPI ranks, a MATLAB grid gives each rank a contiguous block of the serial columns.""" + + N_RANKS = 3 + + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.tmpdir.cleanup() + + def _init_grid(self, num_ranks, rank): + args = parse_cli_args(["--matlab-mesh", str(_MATLAB_MESH), "--no-shading"]) + grid_config = GridConfig(args, num_ranks, rank) + grid, io, _, _ = init_config(TimeConfig(args), grid_config, Path(self.tmpdir.name), False) + return init_grid(grid, io, grid_config, defaultComm) + + def test_blocks_of_all_ranks_make_up_the_serial_grid(self): + serial = self._init_grid(num_ranks=1, rank=0) + blocks = [self._init_grid(self.N_RANKS, rank) for rank in range(self.N_RANKS)] + + self.assertEqual( + [len(block["mask"]) for block in blocks], + [len(b) for b in np.array_split(serial["mask"], self.N_RANKS)], + ) + for key in ("column_id", "x", "y", "z", "lat", "lon", "slope_beta", "slope_gamma"): + with self.subTest(key=key): + np.testing.assert_array_equal(np.concatenate([block[key] for block in blocks]), serial[key]) + np.testing.assert_array_equal(np.concatenate([block["ind"][0] for block in blocks]), serial["ind"][0]) + + class TestReadNetcdfGrid(unittest.TestCase): """Missing orography is NaN unless the user opts in to filling it with zero elevation. From 25d7e270052e7f5779799afd8b0827cb7d8c004c Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 13:04:26 +0200 Subject: [PATCH 6/6] Test that a parallel MATLAB run gives the results of a serial run The MATLAB example runs for one day on a single rank and on 2, 3 and 4 ranks under mpirun; each parallel run must give a byte-identical output file and an equal --dump-reference snapshot. CI also runs the MATLAB example on 4 ranks. Taken from db2b582 and e26d3e8 of #197. Assisted-by: Claude Code:claude-opus-5-5 --- .github/workflows/build-and-test.yaml | 4 ++ .github/workflows/test.yml | 2 +- tests/core/test_parallel_runs.py | 83 +++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 1 deletion(-) create mode 100644 tests/core/test_parallel_runs.py diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml index 9745c940..f20a3763 100644 --- a/.github/workflows/build-and-test.yaml +++ b/.github/workflows/build-and-test.yaml @@ -77,6 +77,10 @@ jobs: # 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 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 6990f15a..7dada70d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -68,7 +68,7 @@ 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 diff --git a/tests/core/test_parallel_runs.py b/tests/core/test_parallel_runs.py new file mode 100644 index 00000000..9d717906 --- /dev/null +++ b/tests/core/test_parallel_runs.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""A run of the MATLAB example on several MPI ranks gives exactly the results of a run on a single rank. + +EBFM runs once on a single rank and once per rank count under mpirun, and the runs must give byte-identical output +files and equal --dump-reference snapshots. +""" + +import filecmp +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +MATLAB_MESH = REPO_ROOT / "examples" / "dem_and_mask.mat" +TIMEOUT = 300 + +requires_mpirun = pytest.mark.skipif( + shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, + reason="mpirun or mpi4py not available", +) + + +# One day: the output file gets its first record after 8 time steps of 3 hours +_TIME = ["--random-seed", "1", "--start-time", "1979-07-01T00:00:00Z", "--end-time", "1979-07-02T00:00:00Z"] +ARGS = ["--matlab-mesh", str(MATLAB_MESH), "--no-shading", *_TIME] + + +def _run(args: list[str], n_ranks: int, workdir: Path) -> Path: + """Run EBFM on n_ranks ranks in workdir, with --dump-reference; return workdir.""" + workdir.mkdir(parents=True) + # Once mpi4py has initialised MPI in this process, the process environment holds the launcher variables of an + # MPI singleton, and an mpirun inheriting them exits 1 without running anything. + env = {k: v for k, v in os.environ.items() if not k.startswith(("OMPI_", "PMIX_", "PMI_", "HYDRA_"))} + cmd = [sys.executable, "-m", "ebfm.main", *args, "--dump-reference", "reference.npz", "--log-file", "ebfm.log"] + if n_ranks > 1: + cmd = ["mpirun", "-np", str(n_ranks), *cmd] + + try: + result = subprocess.run(cmd, cwd=workdir, env=env, capture_output=True, text=True, timeout=TIMEOUT) + except subprocess.TimeoutExpired: + pytest.fail(f"{' '.join(cmd)} timed out after {TIMEOUT} s. Last log lines per rank:\n{_log_tails(workdir)}") + assert result.returncode == 0, f"{' '.join(cmd)} failed:\n{result.stdout}\n{result.stderr}" + return workdir + + +def _log_tails(workdir: Path, n_lines: int = 15) -> str: + """The last lines of the log file of every rank in workdir.""" + tails = [] + for log_file in sorted(workdir.glob("ebfm.log*")): + lines = log_file.read_text(errors="replace").splitlines()[-n_lines:] + tails.append(f"--- {log_file.name}\n" + "\n".join(lines)) + return "\n".join(tails) or "(no log files)" + + +def _assert_same_results(serial: Path, parallel: Path) -> None: + assert filecmp.cmp(serial / "Output" / "model_output.nc", parallel / "Output" / "model_output.nc", shallow=False) + + with np.load(serial / "reference.npz") as expected, np.load(parallel / "reference.npz") as actual: + assert sorted(actual.files) == sorted(expected.files) + for key in expected.files: + np.testing.assert_array_equal(actual[key], expected[key], err_msg=key) + + +@pytest.fixture(scope="module") +def serial(tmp_path_factory) -> Path: + return _run(ARGS, 1, tmp_path_factory.mktemp("serial") / "run") + + +@requires_mpirun +@pytest.mark.parametrize("n_ranks", [2, 3, 4]) +def test_parallel_run_equals_serial_run(tmp_path, serial, n_ranks): + parallel = _run(ARGS, n_ranks, tmp_path / f"np{n_ranks}") + + _assert_same_results(serial, parallel)