Repository: C2SM/icon4py
Pull request #748: Improve backend handling in icon4py_driver
Author: login: jcanton
Created: 2025-05-15T16:21:41Z | Merged: 2025-05-19T20:09:02Z

this line was not the best


Diff:
diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py
index effa0852ef..8df3cb4d93 100644
--- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py
+++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py
@@ -10,9 +10,10 @@
 import functools
 import logging
 
+from gt4py.next import backend as gtx_backend
+
 from icon4py.model.atmosphere.diffusion import diffusion
 from icon4py.model.atmosphere.dycore import solve_nonhydro as solve_nh
-from icon4py.model.common import model_backends
 from icon4py.model.common.grid import vertical as v_grid
 from icon4py.model.driver import initialization_utils as driver_init
 
@@ -24,7 +25,7 @@
 
 @dataclasses.dataclass(frozen=True)
 class Icon4pyRunConfig:
-    backend_name: str
+    backend: gtx_backend.Backend
     dtime: datetime.timedelta = datetime.timedelta(seconds=600.0)  # length of a time step
     start_date: datetime.datetime = datetime.datetime(1, 1, 1, 0, 0, 0)
     end_date: datetime.datetime = datetime.datetime(1, 1, 1, 1, 0, 0)
@@ -42,16 +43,9 @@ class Icon4pyRunConfig:
 
     restart_mode: bool = False
 
-    def __post_init__(self):
-        if self.backend_name not in model_backends.BACKENDS:
-            raise ValueError(
-                f"Invalid driver backend: {self.backend_name}. \n"
-                f"Available backends are {', '.join([f'{k}' for k in model_backends.BACKENDS.keys()])}"
-            )
-
     @functools.cached_property
     def backend(self):
-        return model_backends.BACKENDS[self.backend_name]
+        return self.backend
 
 
 @dataclasses.dataclass
@@ -63,8 +57,8 @@ class Icon4pyConfig:
 
 
 def read_config(
-    icon4py_driver_backend: str,
-    experiment_type: driver_init.ExperimentType = driver_init.ExperimentType.ANY,
+    experiment_type: driver_init.ExperimentType,
+    backend: gtx_backend.Backend,
 ) -> Icon4pyConfig:
     def _mch_ch_r04b09_vertical_config():
         return v_grid.VerticalGridConfig(
@@ -136,7 +130,7 @@ def _mch_ch_r04b09_config():
                 end_date=datetime.datetime(2021, 6, 20, 12, 0, 10),
                 n_substeps=n_substeps_reduced,
                 apply_initial_stabilization=True,
-                backend_name=icon4py_driver_backend,
+                backend=backend,
             ),
             _mch_ch_r04b09_vertical_config(),
             _mch_ch_r04b09_diffusion_config(),
@@ -149,7 +143,7 @@ def _jablownoski_Williamson_config():
             end_date=datetime.datetime(1, 1, 1, 0, 30, 0),
             apply_initial_stabilization=False,
             n_substeps=5,
-            backend_name=icon4py_driver_backend,
+            backend=backend,
         )
         jabw_vertical_config = _jabw_vertical_config()
         jabw_diffusion_config = _jabw_diffusion_config(icon_run_config.n_substeps)
@@ -184,7 +178,7 @@ def _gauss3d_config():
             end_date=datetime.datetime(1, 1, 1, 0, 0, 4),
             apply_initial_stabilization=False,
             n_substeps=5,
-            backend_name=icon4py_driver_backend,
+            backend=backend,
         )
         vertical_config = _gauss3d_vertical_config()
         diffusion_config = _gauss3d_diffusion_config(icon_run_config.n_substeps)
diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py
index 8b376b463e..76cc4f0393 100644
--- a/model/driver/src/icon4py/model/driver/icon4py_driver.py
+++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py
@@ -16,12 +16,15 @@
 import numpy as np
 from devtools import Timer
 
+from gt4py.next import backend as gtx_backend
+
 import icon4py.model.common.utils as common_utils
 from icon4py.model.atmosphere.diffusion import (
     diffusion,
     diffusion_states,
 )
 from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh
+from icon4py.model.common import model_backends
 from icon4py.model.common.decomposition import definitions as decomposition
 from icon4py.model.common.states import (
     diagnostic_state as diagnostics,
@@ -326,7 +329,7 @@ def initialize(
     grid_id: uuid.UUID,
     grid_root,
     grid_level,
-    icon4py_driver_backend: str,
+    backend: gtx_backend.Backend,
 ) -> tuple[TimeLoop, DriverStates, DriverParams]:
     """
     Initialize the driver run.
@@ -354,12 +357,12 @@ def initialize(
     """
     log.info("initialize parallel runtime")
     log.info(f"reading configuration: experiment {experiment_type}")
-    config = driver_config.read_config(icon4py_driver_backend, experiment_type)
+    config = driver_config.read_config(backend, experiment_type)
 
     decomp_info = driver_init.read_decomp_info(
         file_path,
         props,
-        config.run_config.backend,
+        backend,
         serialization_type,
         grid_id,
         grid_root,
@@ -369,7 +372,7 @@ def initialize(
     log.info(f"initializing the grid from '{file_path}'")
     icon_grid = driver_init.read_icon_grid(
         file_path,
-        backend=config.run_config.backend,
+        backend=backend,
         rank=props.rank,
         ser_type=serialization_type,
         grid_id=grid_id,
@@ -385,7 +388,7 @@ def initialize(
     ) = driver_init.read_geometry_fields(
         file_path,
         vertical_grid_config=config.vertical_grid_config,
-        backend=config.run_config.backend,
+        backend=backend,
         rank=props.rank,
         ser_type=serialization_type,
         grid_id=grid_id,
@@ -401,7 +404,7 @@ def initialize(
     ) = driver_init.read_static_fields(
         icon_grid,
         file_path,
-        config.run_config.backend,
+        backend,
         rank=props.rank,
         ser_type=serialization_type,
     )
@@ -419,14 +422,14 @@ def initialize(
         edge_geometry,
         cell_geometry,
         exchange=exchange,
-        backend=config.run_config.backend,
+        backend=backend,
     )
 
     nonhydro_params = solve_nh.NonHydrostaticParams(config.solve_nonhydro_config)
 
     solve_nonhydro_granule = solve_nh.SolveNonhydro(
         grid=icon_grid,
-        backend=config.run_config.backend,
+        backend=backend,
         config=config.solve_nonhydro_config,
         params=nonhydro_params,
         metric_state_nonhydro=solve_nonhydro_metric_state,
@@ -450,7 +453,7 @@ def initialize(
         cell_geometry,
         edge_geometry,
         file_path,
-        backend=config.run_config.backend,
+        backend=backend,
         rank=props.rank,
         experiment_type=experiment_type,
     )
@@ -562,6 +565,14 @@ def icon4py_driver(
 
     2. run time loop
     """
+
+    if icon4py_driver_backend not in model_backends.BACKENDS:
+        raise ValueError(
+            f"Invalid driver backend: {icon4py_driver_backend}. \n"
+            f"Available backends are {', '.join([f'{k}' for k in model_backends.BACKENDS.keys()])}"
+        )
+    backend = model_backends.BACKENDS[icon4py_driver_backend]
+
     parallel_props = decomposition.get_processor_properties(decomposition.get_runtype(with_mpi=mpi))
     grid_id = uuid.UUID(grid_id)
     driver_init.configure_logging(run_path, experiment_type, enable_output, parallel_props)
@@ -577,7 +588,7 @@ def icon4py_driver(
         grid_id,
         grid_root,
         grid_level,
-        icon4py_driver_backend,
+        backend,
     )
     log.info(f"Starting ICON dycore run: {time_loop.simulation_date.isoformat()}")
     log.info(
diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py
index d8688df5ee..3b72ef3032 100644
--- a/model/driver/src/icon4py/model/driver/initialization_utils.py
+++ b/model/driver/src/icon4py/model/driver/initialization_utils.py
@@ -475,17 +475,17 @@ def read_static_fields(
             bdy_halo_c=metrics_savepoint.bdy_halo_c(),
             mask_prog_halo_c=metrics_savepoint.mask_prog_halo_c(),
             rayleigh_w=metrics_savepoint.rayleigh_w(),
-            exner_exfac=metrics_savepoint.exner_exfac(),
-            exner_ref_mc=metrics_savepoint.exner_ref_mc(),
+            time_extrapolation_parameter_for_exner=metrics_savepoint.exner_exfac(),
+            reference_exner_at_cells_on_model_levels=metrics_savepoint.exner_ref_mc(),
             wgtfac_c=metrics_savepoint.wgtfac_c(),
             wgtfacq_c=metrics_savepoint.wgtfacq_c_dsl(),
             inv_ddqz_z_full=metrics_savepoint.inv_ddqz_z_full(),
-            rho_ref_mc=metrics_savepoint.rho_ref_mc(),
-            theta_ref_mc=metrics_savepoint.theta_ref_mc(),
+            reference_rho_at_cells_on_model_levels=metrics_savepoint.rho_ref_mc(),
+            reference_theta_at_cells_on_model_levels=metrics_savepoint.theta_ref_mc(),
             exner_w_explicit_weight_parameter=metrics_savepoint.vwind_expl_wgt(),
-            d_exner_dz_ref_ic=metrics_savepoint.d_exner_dz_ref_ic(),
+            ddz_of_reference_exner_at_cells_on_half_levels=metrics_savepoint.d_exner_dz_ref_ic(),
             ddqz_z_half=metrics_savepoint.ddqz_z_half(),
-            theta_ref_ic=metrics_savepoint.theta_ref_ic(),
+            reference_theta_at_cells_on_half_levels=metrics_savepoint.theta_ref_ic(),
             d2dexdz2_fac1_mc=metrics_savepoint.d2dexdz2_fac1_mc(),
             d2dexdz2_fac2_mc=metrics_savepoint.d2dexdz2_fac2_mc(),
             reference_rho_at_edges_on_model_levels=metrics_savepoint.rho_ref_me(),
diff --git a/model/driver/tests/driver_tests/test_timeloop.py b/model/driver/tests/driver_tests/test_timeloop.py
index 6b543ed30a..dfabb454e1 100644
--- a/model/driver/tests/driver_tests/test_timeloop.py
+++ b/model/driver/tests/driver_tests/test_timeloop.py
@@ -130,10 +130,9 @@ def test_run_timeloop_single_step(
     backend,
 ):
     if experiment == dt_utils.GAUSS3D_EXPERIMENT:
-        # it does not matter what backend is set here because the granules are set externally in this test
         config = icon4py_configuration.read_config(
-            icon4py_driver_backend="gtfn_cpu",
             experiment_type=experiment,
+            backend=backend,
         )
         diffusion_config = config.diffusion_config
         nonhydro_config = config.solve_nonhydro_config
@@ -148,6 +147,7 @@ def test_run_timeloop_single_step(
             timeloop_date_exit,
             timeloop_diffusion_linit_init,
             ndyn_substeps=ndyn_substeps,
+            backend=backend,
         )
 
     edge_geometry: grid_states.EdgeParams = grid_savepoint.construct_edge_geometry()
diff --git a/model/driver/tests/driver_tests/utils.py b/model/driver/tests/driver_tests/utils.py
index f684c9dd83..cb34f658fc 100644
--- a/model/driver/tests/driver_tests/utils.py
+++ b/model/driver/tests/driver_tests/utils.py
@@ -8,6 +8,8 @@
 
 from datetime import datetime, timedelta
 
+from gt4py.next import backend as gtx_backend
+
 from icon4py.model.atmosphere.diffusion import diffusion
 from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh
 from icon4py.model.driver import icon4py_configuration as driver_config
@@ -105,6 +107,7 @@ def mch_ch_r04b09_dsl_icon4pyrun_config(
     date_init: str,
     date_exit: str,
     diffusion_linit_init: bool,
+    backend: gtx_backend.Backend,
     ndyn_substeps: int,
 ) -> driver_config.Icon4pyRunConfig:
     """
@@ -120,7 +123,7 @@ def mch_ch_r04b09_dsl_icon4pyrun_config(
         n_substeps=ndyn_substeps,
         apply_initial_stabilization=True,
         restart_mode=not diffusion_linit_init,
-        backend_name="gtfn_cpu",
+        backend=backend,
     )
 
 
@@ -128,6 +131,7 @@ def exclaim_ape_icon4pyrun_config(
     date_init: str,
     date_exit: str,
     diffusion_linit_init: bool,
+    backend: gtx_backend.Backend,
     ndyn_substeps: int,
 ) -> driver_config.Icon4pyRunConfig:
     """
@@ -143,7 +147,7 @@ def exclaim_ape_icon4pyrun_config(
         n_substeps=ndyn_substeps,
         apply_initial_stabilization=False,
         restart_mode=not diffusion_linit_init,
-        backend_name="gtfn_cpu",
+        backend=backend,
     )
 
 
@@ -152,13 +156,14 @@ def construct_icon4pyrun_config(
     date_init: str,
     date_exit: str,
     diffusion_linit_init: bool,
+    backend: gtx_backend.Backend,
     ndyn_substeps: int = 5,
 ):
     if name.lower() in "mch_ch_r04b09_dsl":
         return mch_ch_r04b09_dsl_icon4pyrun_config(
-            date_init, date_exit, diffusion_linit_init, ndyn_substeps
+            date_init, date_exit, diffusion_linit_init, backend, ndyn_substeps
         )
     elif name.lower() in "exclaim_ape_r02b04":
         return exclaim_ape_icon4pyrun_config(
-            date_init, date_exit, diffusion_linit_init, ndyn_substeps
+            date_init, date_exit, diffusion_linit_init, backend, ndyn_substeps
         )


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR improves backend handling in the icon4py_driver by converting string-based backend identifiers to typed gtx_backend.Backend objects and updating related configuration components.  
- Refactors test configuration calls in driver_tests.  
- Renames and updates parameter usage for backend handling in driver and initialization modules.  
- Changes field names in configuration and static fields for clearer semantics.

### Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

| File | Description |
| ---- | ----------- |
| model/driver/tests/driver_tests/test_timeloop.py | Updates configuration call to pass the backend parameter. |
| model/driver/src/icon4py/model/driver/initialization_utils.py | Renames field keys to more descriptive names in the static fields reading function. |
| model/driver/src/icon4py/model/driver/icon4py_driver.py | Changes parameter usage from a string identifier to a typed backend with minor naming inconsistencies. |
| model/driver/src/icon4py/model/driver/icon4py_configuration.py | Refactors Icon4pyRunConfig to remove the old backend_name field and introduces a backend property. |




state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T09:05:26Z
user: login: OngChia
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T09:21:41Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR refactors backend handling in the `icon4py_driver` module by switching from string-based backend names to actual `gtx_backend.Backend` objects, adding validation for backend selection, and updating related configuration and initialization calls.

- Replace string `backend_name` parameters with typed `gtx_backend.Backend` in configuration and initialization APIs  
- Add validation and mapping from backend name to backend object in `icon4py_driver`  
- Rename several static-field parameters in `initialization_utils` for clearer naming

### Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

| File                                                           | Description                                                          |
|----------------------------------------------------------------|----------------------------------------------------------------------|
| model/driver/tests/driver_tests/test_timeloop.py               | Update test to pass `backend` fixture instead of hard-coded string  |
| model/driver/src/icon4py/model/driver/initialization_utils.py  | Rename static-field args for `MetricStateNonHydro` constructor       |
| model/driver/src/icon4py/model/driver/icon4py_driver.py        | Change initialize signature, add backend validation, propagate backend |
| model/driver/src/icon4py/model/driver/icon4py_configuration.py | Change `Icon4pyRunConfig` and `read_config` to use typed backend      |


<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**model/driver/src/icon4py/model/driver/icon4py_driver.py:569**
* [nitpick] The local variable `backend` later shadows the imported `gtx_backend` alias and may confuse readers. Consider renaming the local variable (e.g., `selected_backend`) to improve clarity.
```
    if icon4py_driver_backend not in model_backends.BACKENDS:
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T09:23:21Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T19:48:11Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T19:50:36Z
user: login: jcanton
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:17:27Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:42:50Z

Review comments:
diff_hunk: @@ -63,8 +57,8 @@ class Icon4pyConfig:
 
 
 def read_config(
-    icon4py_driver_backend: str,
-    experiment_type: driver_init.ExperimentType = driver_init.ExperimentType.ANY,
+    experiment_type: driver_init.ExperimentType,
+    backend: gtx_backend.Backend,
 ) -> Icon4pyConfig:
path: model/driver/src/icon4py/model/driver/icon4py_configuration.py
user: login: Copilot
body: [nitpick] The `read_config` signature has changed but there is no docstring or comment explaining the new parameter order. Please add or update documentation to reflect the new parameters.
```suggestion
) -> Icon4pyConfig:
    """
    Reads and constructs the configuration for an ICON4Py simulation.

    Args:
        experiment_type (driver_init.ExperimentType): The type of experiment to configure.
        backend (gtx_backend.Backend): The backend to use for the simulation.

    Returns:
        Icon4pyConfig: The complete configuration object for the simulation.
    """
```
created_at: 2025-05-19T09:23:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -130,10 +130,9 @@ def test_run_timeloop_single_step(
     backend,
 ):
     if experiment == dt_utils.GAUSS3D_EXPERIMENT:
path: model/driver/tests/driver_tests/test_timeloop.py
user: login: Copilot
body: [nitpick] The inline comment explaining the backend override was removed. Consider adding a brief note here to explain why the test passes the backend fixture explicitly.
```suggestion
    if experiment == dt_utils.GAUSS3D_EXPERIMENT:
        # Explicitly pass the backend fixture to ensure the configuration is set up
        # correctly for the specified backend during the test.
```
created_at: 2025-05-19T09:23:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -562,6 +565,14 @@ def icon4py_driver(
 
     2. run time loop
     """
+
+    if icon4py_driver_backend not in model_backends.BACKENDS:
path: model/driver/src/icon4py/model/driver/icon4py_driver.py
user: login: OngChia
body: Pytest_config.py also has a duplicated code to check if string `backend` is in `BACKENDS`. I now think that we should have a function maybe in `common/model_backends.py` to do this and call this function here and `pytest_config.py`.
created_at: 2025-05-19T19:48:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -130,10 +130,9 @@ def test_run_timeloop_single_step(
     backend,
 ):
     if experiment == dt_utils.GAUSS3D_EXPERIMENT:
-        # it does not matter what backend is set here because the granules are set externally in this test
         config = icon4py_configuration.read_config(
-            icon4py_driver_backend="gtfn_cpu",
path: model/driver/tests/driver_tests/test_timeloop.py
user: login: OngChia
body: Thanks for correcting my bad workaround here🙇
created_at: 2025-05-19T19:50:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -562,6 +565,14 @@ def icon4py_driver(
 
     2. run time loop
     """
+
+    if icon4py_driver_backend not in model_backends.BACKENDS:
path: model/driver/src/icon4py/model/driver/icon4py_driver.py
user: login: jcanton
body: shit I didn't see this comment and I already merged 😬
created_at: 2025-05-19T20:17:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -562,6 +565,14 @@ def icon4py_driver(
 
     2. run time loop
     """
+
+    if icon4py_driver_backend not in model_backends.BACKENDS:
path: model/driver/src/icon4py/model/driver/icon4py_driver.py
user: login: OngChia
body: It is okay. My bad, this comes to my mind after I approved it... you can modify it when you create another PR to use the proper icon4py configuration in `test_timeloop.py`.
created_at: 2025-05-19T20:42:50Z
author_association: CONTRIBUTOR

PR comments:
user: login: jcanton
created_at: 2025-05-15T16:25:24Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: jcanton
created_at: 2025-05-16T12:09:02Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: jcanton
created_at: 2025-05-16T13:00:40Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: jcanton
created_at: 2025-05-19T09:21:36Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-19T11:31:08Z
body: <!-- add-pr-comment:add-pr-comment -->

**Mandatory Tests**

Please make sure you run these tests via comment before you merge!

* `cscs-ci run default`

**Optional Tests**

To run benchmarks you can use:

* `cscs-ci run benchmark-bencher`

To run tests and benchmarks with the DaCe backend you can use:

* `cscs-ci run dace`

To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:
 * `cscs-ci run extra`



For more detailed information please look at [CI in the EXCLAIM universe](https://github.com/C2SM/icon-exclaim/wiki).
author_association: NONE
user: login: jcanton
created_at: 2025-05-19T11:31:32Z
body: cscs-ci run default
author_association: CONTRIBUTOR