Domain decomposition and halo construction

.github/validation_grid_utils.py

@@ -28,6 +28,7 @@
 
 
 VALIDATION_GRIDS = (
+    definitions.Grids.R01B01_GLOBAL,
     definitions.Grids.R02B04_GLOBAL,
     definitions.Grids.MCH_CH_R04B09_DSL,
     definitions.Grids.MCH_OPR_R04B07_DOMAIN01,

.github/workflows/icon4py-test-model.yml

@@ -8,6 +8,14 @@ on:
     branches:
       - main
     types: [opened, reopened, synchronize]
+
+env:
+  ICON4PY_TEST_DATA_PATH: "${{ github.workspace }}/testdata"
+  NUM_PROCESSES: auto
+  ICON4PY_ENABLE_GRID_DOWNLOAD: false
+  ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false
+  PYTEST_ADDOPTS: "--durations=0"
+
 jobs:
   test-model:
     runs-on: ubuntu-latest
@@ -47,25 +55,17 @@ jobs:
 
       - name: download grids
         if: steps.grid-cache.outputs.cache-hit != 'true'
+        env:
+          ICON4PY_ENABLE_GRID_DOWNLOAD: true
         run: |
           .github/validation_grid_utils.py download
 
       - name: run nox session - unit tests
         shell: bash
-        env:
-          NUM_PROCESSES: auto
-          ICON4PY_ENABLE_GRID_DOWNLOAD: false
-          ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false
-          PYTEST_ADDOPTS: "--durations=0"
         run: |
           uvx nox -s "test_model-${{matrix.python-version}}(basic, ${{ matrix.component }})" -- --backend=${{ matrix.backend }}
 
       - name: run nox session - stencil tests
         shell: bash
-        env:
-          NUM_PROCESSES: auto
-          ICON4PY_ENABLE_GRID_DOWNLOAD: false
-          ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false
-          PYTEST_ADDOPTS: "--durations=0"
         run: |
           uvx nox -s "test_model-${{matrix.python-version}}(stencils, ${{ matrix.component }})" -- --backend=${{ matrix.backend }} --grid=icon_regional

ci/distributed.yml

@@ -85,16 +85,16 @@ build_distributed_cpu:
   rules:
     - if: $COMPONENT == 'atmosphere/diffusion'
       variables:
-        SLURM_TIMELIMIT: '00:05:00'
+        SLURM_TIMELIMIT: '00:10:00'
     - if: $COMPONENT == 'atmosphere/dycore' && $BACKEND == 'dace_cpu'
       variables:
-        SLURM_TIMELIMIT: '00:20:00'
+        SLURM_TIMELIMIT: '00:30:00'
     - if: $COMPONENT == 'atmosphere/dycore'
       variables:
         SLURM_TIMELIMIT: '00:15:00'
     - when: on_success
       variables:
-        SLURM_TIMELIMIT: '00:30:00'
+        SLURM_TIMELIMIT: '00:45:00'
   artifacts:
     when: always
     paths:

model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py

@@ -6,6 +6,7 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import logging
 from typing import Any
 
 import pytest
@@ -21,6 +22,12 @@
 from ..fixtures import *  # noqa: F403
 
 
+if mpi_decomposition.mpi4py is None:
+    pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True)
+
+_log = logging.getLogger(__file__)
+
+
 @pytest.mark.mpi
 @pytest.mark.uses_concat_where
 @pytest.mark.parametrize(
@@ -63,24 +70,24 @@ def test_parallel_diffusion(
         raise pytest.skip("This test is only executed for `dace` backends.")
     caplog.set_level("INFO")
     parallel_helpers.check_comm_size(processor_props)
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: initializing diffusion for experiment '{definitions.Experiments.MCH_CH_R04B09}'"
     )
-    print(
+    _log.info(
         f"local cells = {decomposition_info.global_index(dims.CellDim, decomposition.DecompositionInfo.EntryType.ALL).shape} "
         f"local edges = {decomposition_info.global_index(dims.EdgeDim, decomposition.DecompositionInfo.EntryType.ALL).shape} "
         f"local vertices = {decomposition_info.global_index(dims.VertexDim, decomposition.DecompositionInfo.EntryType.ALL).shape}"
     )
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes"
     )
 
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells} Cells, {icon_grid.num_edges} Edges, {icon_grid.num_vertices} Vertices"
     )
     config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps)
     dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime")
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes"
     )
     vertical_config = v_grid.VerticalGridConfig(
@@ -113,7 +120,7 @@ def test_parallel_diffusion(
         orchestration=orchestration,
     )
 
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
 
     diagnostic_state = diffusion_states.DiffusionDiagnosticState(
         hdef_ic=savepoint_diffusion_init.hdef_ic(),
@@ -135,15 +142,15 @@ def test_parallel_diffusion(
             prognostic_state=prognostic_state,
             dtime=dtime,
         )
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
 
     utils.verify_diffusion_fields(
         config=config,
         diagnostic_state=diagnostic_state,
         prognostic_state=prognostic_state,
         diffusion_savepoint=savepoint_diffusion_exit,
     )
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE"
     )
 
@@ -190,19 +197,19 @@ def test_parallel_diffusion_multiple_steps(
     ######################################################################
     caplog.set_level("INFO")
     parallel_helpers.check_comm_size(processor_props)
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: initializing diffusion for experiment '{definitions.Experiments.MCH_CH_R04B09}'"
     )
-    print(
+    _log.info(
         f"local cells = {decomposition_info.global_index(dims.CellDim, decomposition.DecompositionInfo.EntryType.ALL).shape} "
         f"local edges = {decomposition_info.global_index(dims.EdgeDim, decomposition.DecompositionInfo.EntryType.ALL).shape} "
         f"local vertices = {decomposition_info.global_index(dims.VertexDim, decomposition.DecompositionInfo.EntryType.ALL).shape}"
     )
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes"
     )
 
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells} Cells, {icon_grid.num_edges} Edges, {icon_grid.num_vertices} Vertices"
     )
     cell_geometry = grid_savepoint.construct_cell_geometry()
@@ -218,7 +225,7 @@ def test_parallel_diffusion_multiple_steps(
     config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps)
     diffusion_params = diffusion_.DiffusionParams(config)
     dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime")
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes"
     )
     exchange = decomposition.create_exchange(processor_props, decomposition_info)
@@ -245,7 +252,7 @@ def test_parallel_diffusion_multiple_steps(
         orchestration=False,
     )
 
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
 
     diagnostic_state_dace_non_orch = diffusion_states.DiffusionDiagnosticState(
         hdef_ic=savepoint_diffusion_init.hdef_ic(),
@@ -269,8 +276,8 @@ def test_parallel_diffusion_multiple_steps(
                 prognostic_state=prognostic_state_dace_non_orch,
                 dtime=dtime,
             )
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
-    print(
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE"
     )
 
@@ -296,7 +303,7 @@ def test_parallel_diffusion_multiple_steps(
         backend=backend,
         orchestration=True,
     )
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
 
     diagnostic_state_dace_orch = diffusion_states.DiffusionDiagnosticState(
         hdef_ic=savepoint_diffusion_init.hdef_ic(),
@@ -320,8 +327,8 @@ def test_parallel_diffusion_multiple_steps(
                 prognostic_state=prognostic_state_dace_orch,
                 dtime=dtime,
             )
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
-    print(
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ")
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE"
     )
 

model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py

@@ -32,7 +32,6 @@
 from icon4py.model.common.metrics import metrics_attributes, metrics_factory
 from icon4py.model.common.states import prognostic_state as prognostics
 from icon4py.model.common.utils import data_allocation as data_alloc
-from icon4py.model.testing import grid_utils
 from icon4py.model.testing.fixtures.benchmark import (
     geometry_field_source,
     interpolation_field_source,
@@ -62,8 +61,6 @@ def solve_nonhydro(
 
     nonhydro_params = solve_nh.NonHydrostaticParams(config)
 
-    decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator)
-
     vertical_config = v_grid.VerticalGridConfig(
         mesh.num_levels,
         lowest_layer_thickness=50,
@@ -203,11 +200,7 @@ def solve_nonhydro(
         vertical_params=vertical_grid,
         edge_geometry=edge_geometry,
         cell_geometry=cell_geometry,
-        owner_mask=gtx.as_field(
-            (dims.CellDim,),
-            decomposition_info.owner_mask(dims.CellDim),  # type: ignore[arg-type] # mypy not take the type of owner_mask
-            allocator=allocator,
-        ),
+        owner_mask=geometry_field_source.get("cell_owner_mask"),
         backend=backend_like,
     )
 

model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py

@@ -7,6 +7,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 
+import logging
+
 import numpy as np
 import pytest
 from gt4py.next import typing as gtx_typing
@@ -22,6 +24,9 @@
 from ..fixtures import *  # noqa: F403
 
 
+_log = logging.getLogger(__name__)
+
+
 @pytest.mark.parametrize("processor_props", [True], indirect=True)
 @pytest.mark.datatest
 @pytest.mark.parametrize(
@@ -67,25 +72,25 @@ def test_run_solve_nonhydro_single_step(
         pytest.xfail("ValueError: axes don't match array")
 
     parallel_helpers.check_comm_size(processor_props)
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: inializing dycore for experiment 'mch_ch_r04_b09_dsl"
     )
-    print(
+    _log.info(
         f"local cells = {decomposition_info.global_index(dims.CellDim, definitions.DecompositionInfo.EntryType.ALL).shape} "
         f"local edges = {decomposition_info.global_index(dims.EdgeDim, definitions.DecompositionInfo.EntryType.ALL).shape} "
         f"local vertices = {decomposition_info.global_index(dims.VertexDim, definitions.DecompositionInfo.EntryType.ALL).shape}"
     )
     owned_cells = decomposition_info.owner_mask(dims.CellDim)
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes"
     )
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}"
     )
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo edges {np.count_nonzero(np.invert(decomposition_info.owner_mask(dims.EdgeDim)))}"
     )
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}"
     )
 
@@ -137,7 +142,7 @@ def test_run_solve_nonhydro_single_step(
         exchange=exchange,
     )
 
-    print(
+    _log.info(
         f"rank={processor_props.rank}/{processor_props.comm_size}:  entering : solve_nonhydro.time_step"
     )
 
@@ -153,7 +158,7 @@ def test_run_solve_nonhydro_single_step(
         at_first_substep=(substep_init == 1),
         at_last_substep=(substep_init == ndyn_substeps),
     )
-    print(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ")
+    _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ")
 
     expected_theta_v = savepoint_nonhydro_step_final.theta_v_new().asnumpy()
     calculated_theta_v = prognostic_states.next.theta_v.asnumpy()

model/common/pyproject.toml

@@ -53,6 +53,7 @@ io = [
   "netcdf4>=1.6.1",
   "numpy>=1.23.3",
   "scikit-learn>=1.4.0",
+  "pymetis>2022.1",
   # TODO(halungge): there are failing tests starting from uxarray==2024.4.0: when a data file does not have
   # fields of a given dimension (eg 'edge') then something in uxarray goes wrong with the dimension
   # mapping. It is not yet clear whether this is a uxarray bug or on our side.