diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml index 516c8226..f20a3763 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 : \ @@ -72,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 b25d90dc..7dada70d 100644 --- a/.github/workflows/test.yml +++ b/.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 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 0874c15d..99433bfa 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 @@ -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) 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/core/cli.py b/src/ebfm/core/cli.py index 2ecdfc19..284671f7 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,8 +530,31 @@ 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. @@ -539,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" @@ -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." ) 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 193b2998..cb3ef726 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 @@ -203,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. @@ -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 @@ -229,8 +233,10 @@ 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"]) + 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) @@ -386,7 +392,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_cli.py b/tests/core/test_cli.py index 896d0528..bb3e41f3 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,12 +383,34 @@ 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.""" + """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"), @@ -420,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 41d6f5bb..9381f61c 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -371,6 +371,41 @@ 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 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. 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 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)