Distributed standalone driver

model/common/src/icon4py/model/common/decomposition/definitions.py

@@ -12,10 +12,11 @@
 import dataclasses
 import functools
 import logging
+import warnings
 from collections.abc import Sequence
 from enum import Enum
 from types import ModuleType
-from typing import Any, Literal, Protocol, TypeAlias, overload, runtime_checkable
+from typing import Any, ClassVar, Literal, Protocol, TypeAlias, overload, runtime_checkable
 
 import dace  # type: ignore[import-untyped]
 import gt4py.next as gtx
@@ -408,12 +409,25 @@ def __str__(self) -> str:
 
 @dataclasses.dataclass
 class SingleNodeExchange(ExchangeRuntime):
+    _warning_emitted: ClassVar[bool] = False
+
+    @classmethod
+    def _warn_if_used(cls, *, stacklevel: int = 3) -> None:
+        if not cls._warning_emitted:
+            warnings.warn(
+                "***** SingleNodeExchange is in use; HALO EXCHANGE IS RUNNING IN SINGLE-NODE *****",
+                RuntimeWarning,
+                stacklevel=stacklevel,
+            )
+            cls._warning_emitted = True
+
     def start(
         self,
         dim: gtx.Dimension,
         *fields: gtx.Field | data_alloc.NDArray,
         stream: StreamLike = DEFAULT_STREAM,
     ) -> ExchangeResult:
+        self._warn_if_used()
         return SingleNodeResult()
 
     def my_rank(self) -> int:

model/common/src/icon4py/model/common/grid/grid_manager.py

@@ -106,8 +106,8 @@ def __call__(
         allocator: gtx_typing.Allocator | None,
         keep_skip_values: bool,
         decomposer: decomp.Decomposer = _single_node_decomposer,
-        run_properties=_single_process_props,
-    ):
+        run_properties: decomposition.ProcessProperties = _single_process_props,
+    ) -> None:
         if not run_properties.is_single_rank() and isinstance(
             decomposer, decomp.SingleNodeDecomposer
         ):

model/common/src/icon4py/model/common/grid/utils.py

@@ -5,21 +5,20 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-from types import ModuleType
 
 import numpy as np
 
 from icon4py.model.common.grid import gridfile
 
 
-def revert_repeated_index_to_invalid(offset: np.ndarray, array_ns: ModuleType):
+def revert_repeated_index_to_invalid(offset: np.ndarray):
     num_elements = offset.shape[0]
     for i in range(num_elements):
         # convert repeated indices back into -1
-        for val in array_ns.flip(offset[i, :]):
-            if array_ns.count_nonzero(val == offset[i, :]) > 1:
-                unique_values, counts = array_ns.unique(offset[i, :], return_counts=True)
+        for val in np.flip(offset[i, :]):
+            if np.count_nonzero(val == offset[i, :]) > 1:
+                unique_values, counts = np.unique(offset[i, :], return_counts=True)
                 rep_values = unique_values[counts > 1]
-                rep_indices = array_ns.where(array_ns.isin(offset[i, :], rep_values))[0]
+                rep_indices = np.where(np.isin(offset[i, :], rep_values))[0]
                 offset[i, rep_indices[1:]] = gridfile.GridFile.INVALID_INDEX
     return offset

model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py

@@ -325,6 +325,7 @@ def test_exchange_on_dummy_data(
 
 @pytest.mark.mpi
 @pytest.mark.datatest
+@pytest.mark.embedded_only
 @pytest.mark.parametrize("processor_props", [False], indirect=True)
 def test_halo_exchange_for_sparse_field(
     interpolation_savepoint: serialbox.InterpolationSavepoint,

model/common/tests/common/decomposition/unit_tests/test_definitions.py

@@ -6,13 +6,17 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import sys
+import warnings
+
 import gt4py.next as gtx
 import numpy as np
 import pytest
 from gt4py.next import common as gtx_common
 
 import icon4py.model.common.dimension as dims
 import icon4py.model.common.utils.data_allocation as data_alloc
+from icon4py.model.common import dimension as dims
 from icon4py.model.common.decomposition import definitions
 from icon4py.model.common.decomposition.definitions import (
     DecompositionInfo,
@@ -94,3 +98,57 @@ def test_decomposition_info_is_distributed(flag, expected) -> None:
         np.ones((mesh.num_cells,)) * flag,
     )
     assert decomp.is_distributed() == expected
+
+
+def test_single_node_exchange_warns_on_first_use(monkeypatch):
+    monkeypatch.setattr(SingleNodeExchange, "_warning_emitted", False)
+
+    exchange = SingleNodeExchange()
+
+    with pytest.warns(RuntimeWarning, match="SingleNodeExchange"):
+        exchange.start(dims.CellDim)
+
+
+def test_single_node_exchange_does_not_warn_on_construction_or_repeat_use(monkeypatch):
+    monkeypatch.setattr(SingleNodeExchange, "_warning_emitted", False)
+
+    with warnings.catch_warnings(record=True) as recorded_warnings:
+        warnings.simplefilter("always")
+        exchange = SingleNodeExchange()
+
+    assert len(recorded_warnings) == 0
+
+    with pytest.warns(RuntimeWarning, match="SingleNodeExchange"):
+        exchange.exchange(dims.CellDim)
+
+    with warnings.catch_warnings(record=True) as repeated_warnings:
+        warnings.simplefilter("always")
+        exchange.start(dims.CellDim)
+
+    assert len(repeated_warnings) == 0
+
+
+def _assert_warning_points_to_call_site(monkeypatch, func, expected_line):
+    monkeypatch.setattr(SingleNodeExchange, "_warning_emitted", False)
+
+    with warnings.catch_warnings(record=True) as caught_warnings:
+        warnings.simplefilter("always")
+        func()
+
+    assert len(caught_warnings) == 1
+    assert caught_warnings[0].filename == __file__
+    assert caught_warnings[0].lineno == expected_line
+
+
+def test_single_node_exchange_warning_points_to_call_site(monkeypatch):
+    exchange = SingleNodeExchange()
+
+    exchange_line = sys._getframe().f_lineno + 1
+    _assert_warning_points_to_call_site(
+        monkeypatch, lambda: exchange.start(dims.CellDim), exchange_line
+    )
+
+    wait_line = sys._getframe().f_lineno + 1
+    _assert_warning_points_to_call_site(
+        monkeypatch, lambda: exchange.exchange(dims.CellDim), wait_line
+    )

model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py

@@ -5,9 +5,7 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-import functools
 import logging
-import operator
 
 import numpy as np
 import pytest
@@ -33,7 +31,7 @@
 from icon4py.model.common.metrics import metrics_attributes, metrics_factory
 from icon4py.model.common.states import utils as state_utils
 from icon4py.model.common.utils import data_allocation as data_alloc
-from icon4py.model.testing import definitions as test_defs, grid_utils, test_utils
+from icon4py.model.testing import definitions as test_defs, grid_utils, parallel_helpers, test_utils
 from icon4py.model.testing.fixtures.datatest import (
     backend,
     experiment,
@@ -57,7 +55,7 @@ def test_grid_manager_validate_decomposer(
     processor_props: decomp_defs.ProcessProperties,
     experiment: test_defs.Experiment,
 ) -> None:
-    if experiment == test_defs.Experiments.MCH_CH_R04B09:
+    if experiment.grid.params.limited_area:
         pytest.xfail("Limited-area grids not yet supported")
 
     file = grid_utils.resolve_full_grid_file_name(experiment.grid)
@@ -85,104 +83,6 @@ def _get_neighbor_tables(grid: base.Grid) -> dict:
     }
 
 
-def gather_field(field: np.ndarray, props: decomp_defs.ProcessProperties) -> tuple:
-    constant_dims = tuple(field.shape[1:])
-    _log.info(f"gather_field on rank={props.rank} - gathering field of local shape {field.shape}")
-    # Because of sparse indexing the field may have a non-contigous layout,
-    # which Gatherv doesn't support. Make sure the field is contiguous.
-    field = np.ascontiguousarray(field)
-    constant_length = functools.reduce(operator.mul, constant_dims, 1)
-    local_sizes = np.array(props.comm.gather(field.size, root=0))
-    if props.rank == 0:
-        recv_buffer = np.empty(np.sum(local_sizes), dtype=field.dtype)
-        _log.info(
-            f"gather_field on rank = {props.rank} - setup receive buffer with size {sum(local_sizes)} on rank 0"
-        )
-    else:
-        recv_buffer = None
-
-    props.comm.Gatherv(sendbuf=field, recvbuf=(recv_buffer, local_sizes), root=0)
-    if props.rank == 0:
-        local_first_dim = tuple(sz // constant_length for sz in local_sizes)
-        _log.info(
-            f" gather_field on rank = 0: computed local dims {local_first_dim} - constant dims {constant_dims}"
-        )
-        gathered_field = recv_buffer.reshape((-1, *constant_dims))  # type: ignore [union-attr]
-    else:
-        gathered_field = None
-        local_first_dim = field.shape
-    return local_first_dim, gathered_field
-
-
-def check_local_global_field(
-    decomposition_info: decomp_defs.DecompositionInfo,
-    processor_props: decomp_defs.ProcessProperties,  # F811 # fixture
-    dim: gtx.Dimension,
-    global_reference_field: np.ndarray,
-    local_field: np.ndarray,
-    check_halos: bool,
-) -> None:
-    if dim == dims.KDim:
-        test_utils.assert_dallclose(global_reference_field, local_field)
-        return
-
-    _log.info(
-        f" rank= {processor_props.rank}/{processor_props.comm_size}----exchanging field of main dim {dim}"
-    )
-    assert (
-        local_field.shape[0]
-        == decomposition_info.global_index(dim, decomp_defs.DecompositionInfo.EntryType.ALL).shape[
-            0
-        ]
-    )
-
-    # Compare halo against global reference field
-    if check_halos:
-        test_utils.assert_dallclose(
-            global_reference_field[
-                decomposition_info.global_index(dim, decomp_defs.DecompositionInfo.EntryType.HALO)
-            ],
-            local_field[
-                decomposition_info.local_index(dim, decomp_defs.DecompositionInfo.EntryType.HALO)
-            ],
-            atol=1e-9,
-            verbose=True,
-        )
-
-    # Compare owned local field, excluding halos, against global reference
-    # field, by gathering owned entries to the first rank. This ensures that in
-    # total we have the full global field distributed on all ranks.
-    owned_entries = local_field[
-        decomposition_info.local_index(dim, decomp_defs.DecompositionInfo.EntryType.OWNED)
-    ]
-    gathered_sizes, gathered_field = gather_field(owned_entries, processor_props)
-
-    global_index_sizes, gathered_global_indices = gather_field(
-        decomposition_info.global_index(dim, decomp_defs.DecompositionInfo.EntryType.OWNED),
-        processor_props,
-    )
-
-    if processor_props.rank == 0:
-        _log.info(f"rank = {processor_props.rank}: asserting gathered fields: ")
-
-        assert np.all(
-            gathered_sizes == global_index_sizes
-        ), f"gathered field sizes do not match:  {dim} {gathered_sizes} - {global_index_sizes}"
-        _log.info(
-            f"rank = {processor_props.rank}: Checking field size on dim ={dim}: --- gathered sizes {gathered_sizes} = {sum(gathered_sizes)}"
-        )
-        _log.info(
-            f"rank = {processor_props.rank}:                      --- gathered field has size {gathered_sizes}"
-        )
-        sorted_ = np.zeros(global_reference_field.shape, dtype=gtx.float64)  # type: ignore [attr-defined]
-        sorted_[gathered_global_indices] = gathered_field
-        _log.info(
-            f" rank = {processor_props.rank}: SHAPES: global reference field {global_reference_field.shape}, gathered = {gathered_field.shape}"
-        )
-
-        test_utils.assert_dallclose(sorted_, global_reference_field, atol=1e-9, verbose=True)
-
-
 # These fields can't be computed with the embedded backend for one reason or
 # another, so we declare them here for xfailing.
 embedded_broken_fields = {
@@ -316,7 +216,7 @@ def test_geometry_fields_compare_single_multi_rank(
     field = multi_rank_geometry.get(attrs_name)
     dim = field_ref.domain.dims[0]
 
-    check_local_global_field(
+    parallel_helpers.check_local_global_field(
         decomposition_info=multi_rank_grid_manager.decomposition_info,
         processor_props=processor_props,
         dim=dim,
@@ -361,7 +261,7 @@ def test_interpolation_fields_compare_single_multi_rank(
     experiment: test_defs.Experiment,
     attrs_name: str,
 ) -> None:
-    if experiment == test_defs.Experiments.MCH_CH_R04B09:
+    if experiment.grid.params.limited_area:
         pytest.xfail("Limited-area grids not yet supported")
 
     if attrs_name in embedded_broken_fields and test_utils.is_embedded(backend):
@@ -430,7 +330,7 @@ def test_interpolation_fields_compare_single_multi_rank(
     field = multi_rank_interpolation.get(attrs_name)
     dim = field_ref.domain.dims[0]
 
-    check_local_global_field(
+    parallel_helpers.check_local_global_field(
         decomposition_info=multi_rank_grid_manager.decomposition_info,
         processor_props=processor_props,
         dim=dim,
@@ -503,7 +403,7 @@ def test_metrics_fields_compare_single_multi_rank(
     experiment: test_defs.Experiment,
     attrs_name: str,
 ) -> None:
-    if experiment == test_defs.Experiments.MCH_CH_R04B09:
+    if experiment.grid.params.limited_area:
         pytest.xfail("Limited-area grids not yet supported")
 
     if attrs_name in embedded_broken_fields and test_utils.is_embedded(backend):
@@ -660,7 +560,7 @@ def test_metrics_fields_compare_single_multi_rank(
         assert isinstance(field, state_utils.ScalarType)
         assert pytest.approx(field) == field_ref
     else:
-        check_local_global_field(
+        parallel_helpers.check_local_global_field(
             decomposition_info=multi_rank_grid_manager.decomposition_info,
             processor_props=processor_props,
             dim=field_ref.domain.dims[0],
@@ -682,7 +582,7 @@ def test_metrics_mask_prog_halo_c(
     backend: gtx_typing.Backend | None,
     experiment: test_defs.Experiment,
 ) -> None:
-    if experiment == test_defs.Experiments.MCH_CH_R04B09:
+    if experiment.grid.params.limited_area:
         pytest.xfail("Limited-area grids not yet supported")
 
     file = grid_utils.resolve_full_grid_file_name(experiment.grid)
@@ -813,7 +713,7 @@ def test_validate_skip_values_in_distributed_connectivities(
     processor_props: decomp_defs.ProcessProperties,
     experiment: test_defs.Experiment,
 ) -> None:
-    if experiment == test_defs.Experiments.MCH_CH_R04B09:
+    if experiment.grid.params.limited_area:
         pytest.xfail("Limited-area grids not yet supported")
 
     file = grid_utils.resolve_full_grid_file_name(experiment.grid)