Merge branch 'main' into distributed_standalone_driver

* main:
  Scheduled Halo Exchange (#980)
  Add missing metrics fields to `test_parallel_grid_manager.py` test (#1114)
  Muphys: Lowering with single precision (#1101)
  Add single-rank lsq pseudoinv factory test (#1099)
  Cleanup Diffusion config (#1060)
  Fortran bindings: fix numpy allocation and cleanups (#1112)
  fix: fix gt4py metrics extractor in the StencilTest benchmarking (#1111)
  py2fgen: don't recompile if unchanged (#1110)
  CI for standalone_driver (#1070)
  Update mpi4py and pymetis groups to make them optional (#1100)
  Bump mshick/add-pr-comment from 2 to 3 (#1109)
  Use inout fields for full_muphys as well (#1108)
  Update GPU configuration for graupel (#1104)
  Move the mask of _q_t_update outside in graupel (#1093)
  Update gt4py to v1.1.7 (#1105)
  cleanup for ugly if condition of single node default in lsq coeffs (#1103)

Author: jcanton

.github/workflows/mandatory_and_optional_test_reminder.yml

@@ -7,7 +7,7 @@ jobs:
     permissions:
       pull-requests: write
     steps:
-      - uses: mshick/add-pr-comment@v2
+      - uses: mshick/add-pr-comment@v3
         with:
           refresh-message-position: true
           message: |2

.pre-commit-config.yaml

@@ -103,7 +103,7 @@ repos:
       - <<: *uv-managed-hook
         id: mypy
         name: mypy static type checker
-        entry: uv run --group dev --frozen --isolated mypy
+        entry: uv run --group typing-distributed --frozen --isolated mypy
         types_or: [python, pyi]
         pass_filenames: false
 

ci/default.yml

@@ -45,6 +45,9 @@ test_tools_datatests_aarch64:
       when: never  # run only in daily CI, to save compute resources
     - if: $BACKEND == 'dace_gpu' && $COMPONENT == 'driver'
       when: never  # run only in daily CI, to save compute resources
+    - if: $COMPONENT == 'standalone_driver'
+      variables:
+        SLURM_TIMELIMIT: '00:50:00'
     - if: $BACKEND == 'dace_gpu' || $BACKEND == 'gtfn_gpu'
       variables:
         NUM_PROCESSES: 8

model/atmosphere/advection/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # workspace members
   "icon4py-common>=0.0.6",
   # external dependencies
-  "gt4py==1.1.6",
+  "gt4py==1.1.7",
   'packaging>=20.0'
 ]
 description = "ICON advection."

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py

@@ -207,7 +207,9 @@ def run(
         log.debug("advection run - start")
 
         log.debug("communication of prep_adv cell field: mass_flx_ic - start")
-        self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic)
+        self._exchange.exchange(
+            dims.CellDim, prep_adv.mass_flx_ic, stream=decomposition.DEFAULT_STREAM
+        )
         log.debug("communication of prep_adv cell field: mass_flx_ic - end")
 
         log.debug("running stencil copy_cell_kdim_field - start")
@@ -313,7 +315,11 @@ def run(
         log.debug("advection run - start")
 
         log.debug("communication of prep_adv cell field: mass_flx_ic - start")
-        self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic)
+        self._exchange.exchange(
+            dims.CellDim,
+            prep_adv.mass_flx_ic,
+            stream=decomposition.DEFAULT_STREAM,
+        )
         log.debug("communication of prep_adv cell field: mass_flx_ic - end")
 
         # reintegrate density for conservation of mass
@@ -396,7 +402,11 @@ def run(
 
         # exchange updated tracer values, originally happens only if iforcing /= inwp
         log.debug("communication of advection cell field: p_tracer_new - start")
-        self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new)
+        self._exchange.exchange(
+            dims.CellDim,
+            p_tracer_new,
+            stream=decomposition.DEFAULT_STREAM,
+        )
         log.debug("communication of advection cell field: p_tracer_new - end")
 
         # finalize step

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py

@@ -174,7 +174,7 @@ def apply_flux_limiter(
         )
 
         log.debug("communication of advection cell field: r_m - start")
-        self._exchange.exchange_and_wait(dims.CellDim, self._r_m)
+        self._exchange.exchange(dims.CellDim, self._r_m, stream=decomposition.DEFAULT_STREAM)
         log.debug("communication of advection cell field: r_m - end")
 
         # limit outward fluxes

model/atmosphere/advection/tests/advection/integration_tests/test_advection.py

@@ -47,6 +47,7 @@
     construct_least_squares_state,
     construct_metric_state,
     construct_prep_adv,
+    dummy_exchange,
     log_serialized,
     verify_advection_fields,
 )
@@ -161,6 +162,7 @@ def test_advection_run_single_step(
         ),
         min_rlcell_int=icon_grid.end_index(h_grid.domain(dims.CellDim)(h_grid.Zone.LOCAL)),
         geometry_type=icon_grid.geometry_type,
+        exchange=dummy_exchange,
     )
 
     least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend)
@@ -268,6 +270,7 @@ def test_compute_lsq_coeffs(
         start_idx,
         min_rlcell_int,
         icon_grid.geometry_type,
+        exchange=dummy_exchange,
     )
 
     assert test_helpers.dallclose(

model/atmosphere/advection/tests/advection/utils.py

@@ -197,3 +197,7 @@ def verify_advection_fields(
         p_tracer_new_ref.asnumpy()[p_tracer_new_range, :],
         atol=1e-16,
     )
+
+
+def dummy_exchange(*field: data_alloc.NDArray) -> None:
+    return None

model/atmosphere/diffusion/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
   # workspace members
   "icon4py-common>=0.0.6",
   # external dependencies
-  "gt4py==1.1.6",
+  "gt4py==1.1.7",
   'packaging>=20.0'
 ]
 description = "ICON diffusion."

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -153,8 +153,6 @@ def __init__(
         smagorinski_scaling_factor: float = 0.015,
         n_substeps: int = 5,
         zdiffu_t: bool = True,
-        thslp_zdiffu: float = 0.025,
-        thhgtd_zdiffu: float = 200.0,
         velocity_boundary_diffusion_denom: float = 200.0,
         temperature_boundary_diffusion_denom: float = 135.0,
         _nudge_max_coeff: float | None = None,  # default is set in __init__
@@ -212,11 +210,6 @@ def __init__(
         #: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90
         self.apply_zdiffusion_t: bool = zdiffu_t
 
-        #:slope threshold (temperature diffusion): is used to build up an index list for application of truly horizontal diffusion in mo_vertical_grid.f90
-        self.thslp_zdiffu = thslp_zdiffu
-        #: threshold [m] for height difference between adjacent grid points, defaults to 200m (temperature diffusion)
-        self.thhgtd_zdiffu = thhgtd_zdiffu
-
         # from other namelists:
         # from parent namelist mo_nonhydrostatic_nml
 
@@ -375,9 +368,6 @@ def _determine_smagorinski_factor(self, config: DiffusionConfig):
                 smagorinski_height = None
             case _:
                 raise NotImplementedError("Only implemented for diffusion type 4 and 5")
-                smagorinski_factor = None
-                smagorinski_height = None
-                pass
         return smagorinski_factor, smagorinski_height
 
 
@@ -428,7 +418,7 @@ def __init__(
         self._cell_params = cell_params
 
         self.halo_exchange_wait = decomposition.create_halo_exchange_wait(
-            self._exchange
+            self._exchange,
         )  # wait on a communication handle
         self.rd_o_cvd: float = constants.GAS_CONSTANT_DRY_AIR / (
             constants.CPD - constants.GAS_CONSTANT_DRY_AIR
@@ -772,11 +762,12 @@ def _sync_cell_fields(self, prognostic_state):
         IF ( linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN
         """
         log.debug("communication of prognostic cell fields: theta, w, exner - start")
-        self._exchange.exchange_and_wait(
+        self._exchange.exchange(
             dims.CellDim,
             prognostic_state.w,
             prognostic_state.theta_v,
             prognostic_state.exner,
+            stream=decomposition.DEFAULT_STREAM,
         )
         log.debug("communication of prognostic cell fields: theta, w, exner - done")
 
@@ -813,12 +804,17 @@ def _do_diffusion_step(
         log.debug("rbf interpolation 1: end")
 
         # 2.  HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
+        # TODO(phimuell, muellch): Is asynchronous mode okay here.
+        # NOTE: We do not specify a stream here but rely on the default argument.
+        #   We do this to ensure that the orchestrator works, but it is not aware
+        #   of the streams.
         log.debug("communication rbf extrapolation of vn - start")
         self._exchange(
             self.u_vert,
             self.v_vert,
             dim=dims.VertexDim,
-            wait=True,
+            full_exchange=True,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
         )
         log.debug("communication rbf extrapolation of vn - end")
 
@@ -861,12 +857,14 @@ def _do_diffusion_step(
         log.debug("2nd rbf interpolation: end")
 
         # 6.  HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields)
+        # TODO(phimuell, muellch): Is asynchronous mode okay here.
         log.debug("communication rbf extrapolation of z_nable2_e - start")
         self._exchange(
             self.u_vert,
             self.v_vert,
             dim=dims.VertexDim,
-            wait=True,
+            full_exchange=True,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
         )
         log.debug("communication rbf extrapolation of z_nable2_e - end")
 
@@ -882,7 +880,12 @@ def _do_diffusion_step(
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end")
 
         log.debug("communication of prognistic.vn : start")
-        handle_edge_comm = self._exchange(prognostic_state.vn, dim=dims.EdgeDim, wait=False)
+        handle_edge_comm = self._exchange(
+            prognostic_state.vn,
+            dim=dims.EdgeDim,
+            full_exchange=False,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
+        )
 
         log.debug(
             "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"
@@ -928,7 +931,8 @@ def _do_diffusion_step(
             log.debug("running stencil 13 to 16 apply_diffusion_to_theta_and_exner: end")
 
         self.halo_exchange_wait(
-            handle_edge_comm
+            handle_edge_comm,
+            # stream=decomposition.DEFAULT_STREAM,  # noqa: ERA001  # See NOTE above.
         )  # need to do this here, since we currently only use 1 communication object.
         log.debug("communication of prognogistic.vn - end")