Clean up process_props/exchange/reductions API - #1380
Conversation
…singletons Replace exchange=/global_reductions= parameters across migrated production APIs (InterpolationFieldsFactory, MetricsFieldsFactory, GridGeometry, Icon4pyDriver, create_static_field_factories) with required process_props, deriving exchange and reductions internally via create_exchange/create_reduction. Add module-level exchange and reductions caches keyed by (id(process_props), id(decomp_info)) with weakref stale-entry detection for both process_props and decomp_info, plus clear_exchange_cache() / clear_reductions_cache() for test cleanup. An autouse pytest fixture calls both after each test. Remove the single_node_exchange / single_node_reductions module-level singletons and the create_single_node_exchange / create_single_reduction_exchange singledispatch registrations. The base create_exchange / create_reduction now falls back to SingleNodeExchange() / SingleNodeReductions() for any ProcessProperties where is_single_rank() is True, else raises NotImplementedError. MPI registrations in mpi_decomposition.py are preserved with their comm_size > 1 check intact. Make compute_nflat_gradp's min_reduction a required argument (DD5). Switch grid_wrapper.py single-node branch to create_exchange(process_props). Remove grid_manager._single_process_props; GridManager.__call__ now requires process_props (no default). Update all test call sites: migrated factory tests pass process_props=SingleNodeProcessProperties(); non-migrated leaf consumer tests pass exchange=SingleNodeExchange().
…e-reduction-cleanup
|
cscs-ci run default |
There was a problem hiding this comment.
Pull request overview
This PR streamlines distributed-runtime wiring by making higher-level constructors accept process_props as the single source of truth, and deriving/caching halo-exchange and global-reduction helpers internally to reduce redundant parameter threading.
Changes:
- Updated core factories/constructors (geometry/interpolation/metrics/driver) to take
process_propsand buildexchange/reductionsinternally. - Introduced per-(process_props, decomposition_info) caching for exchange/reductions, and added pytest hooks to clear caches between tests.
- Updated tests/fixtures to stop using removed
single_node_exchange/single_node_reductionssingletons and instead constructSingleNodeProcessProperties()/SingleNodeExchange()where appropriate.
Reviewed changes
Copilot reviewed 32 out of 32 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| model/testing/src/icon4py/model/testing/pytest_hooks.py | Clears new decomposition caches after each test for isolation. |
| model/testing/src/icon4py/model/testing/grid_utils.py | Thread process_props into GridManager construction (defaulting to single-node). |
| model/testing/src/icon4py/model/testing/fixtures/benchmark.py | Updates benchmark fixtures to new process_props-based factory APIs. |
| model/standalone_driver/tests/standalone_driver/topography/test_jablonowski_williamson_topography.py | Replaces removed singleton exchange usage in test. |
| model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py | Driver now derives exchange/reductions from process_props. |
| model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py | Static-field factory creation now only requires process_props. |
| model/driver/tests/driver/integration_tests/test_icon4py.py | Replaces removed singleton exchange usage in integration test. |
| model/common/tests/common/states/unit_tests/test_factory.py | Replaces removed singleton exchange usage in unit tests. |
| model/common/tests/common/metrics/unit_tests/test_metrics_factory.py | Updates metrics factory tests to pass process_props instead of exchange. |
| model/common/tests/common/metrics/unit_tests/test_metric_fields.py | Replaces removed singleton exchange usage in unit test. |
| model/common/tests/common/metrics/unit_tests/test_compute_zdiff_gradp.py | Replaces removed singleton exchange usage in unit test. |
| model/common/tests/common/metrics/unit_tests/test_compute_weight_factors.py | Replaces removed singleton exchange usage in unit test. |
| model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py | Replaces removed singleton exchange usage in unit tests. |
| model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py | Updates interpolation factory tests to pass process_props. |
| model/common/tests/common/grid/unit_tests/test_vertical.py | Replaces removed singleton exchange usage in unit test. |
| model/common/tests/common/grid/unit_tests/test_topography.py | Replaces removed singleton exchange usage in unit test. |
| model/common/tests/common/grid/unit_tests/test_grid_manager.py | Updates GridManager call sites to pass process_props. |
| model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py | Updates MPI grid tests to new process_props-based APIs. |
| model/common/tests/common/fixtures.py | Fixtures now rely on factories deriving exchange/reductions from process_props. |
| model/common/src/icon4py/model/common/metrics/metrics_factory.py | Metrics factory derives exchange/reductions from process_props. |
| model/common/src/icon4py/model/common/metrics/metric_fields.py | Makes reduction dependency explicit via required min_reduction argument. |
| model/common/src/icon4py/model/common/interpolation/interpolation_factory.py | Interpolation factory derives exchange from process_props. |
| model/common/src/icon4py/model/common/grid/grid_manager.py | Makes process_props required when constructing the grid decomposition. |
| model/common/src/icon4py/model/common/grid/geometry.py | Geometry derives exchange/reductions from process_props. |
| model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py | Adjusts factory registrations to new create_exchange/create_reduction registration surface. |
| model/common/src/icon4py/model/common/decomposition/definitions.py | Implements cached create_exchange/create_reduction and removes single-node singletons. |
| model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py | Replaces removed singleton exchange usage in dycore tests. |
| model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py | Replaces removed singleton exchange usage in benchmark test. |
| model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py | Replaces removed singleton exchange usage in diffusion tests. |
| model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py | Replaces removed singleton exchange usage in diffusion benchmark test. |
| model/atmosphere/advection/tests/advection/integration_tests/test_advection.py | Replaces removed singleton exchange usage in advection tests. |
| bindings/src/icon4py/bindings/grid_wrapper.py | Uses create_exchange for single-node path instead of constructing exchange directly. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| key = (id(process_props), id(decomp_info)) | ||
| if (entry := _exchange_cache.get(key)) is not None: | ||
| if entry.process_props_ref() is None or ( | ||
| entry.decomp_info_ref is not None and entry.decomp_info_ref() is None | ||
| ): | ||
| del _exchange_cache[key] | ||
| else: | ||
| return entry.exchange | ||
| exchange = _create_exchange_impl(process_props, decomp_info) | ||
| _exchange_cache[key] = _ExchangeCacheEntry( | ||
| exchange=exchange, | ||
| process_props_ref=weakref.ref(process_props), | ||
| decomp_info_ref=weakref.ref(decomp_info) if decomp_info is not None else None, | ||
| ) |
There was a problem hiding this comment.
The cache intentionally stores strong references; the weak-reference check only detects stale keys when the same (process_props, decomp_info) pair is reused. Production runs create one ProcessProperties/DecompositionInfo pair for the whole run, and the test autouse fixture clears these caches between tests. Failing fast for missing decomp_info in distributed runs is a sensible follow-up but out of scope here.
🤖 Posted by an agent on behalf of @msimberg
| key = (id(process_props), id(decomposition_info)) | ||
| if (entry := _reduction_cache.get(key)) is not None: | ||
| if entry.process_props_ref() is None or ( | ||
| entry.decomp_info_ref is not None and entry.decomp_info_ref() is None | ||
| ): | ||
| del _reduction_cache[key] | ||
| else: | ||
| return entry.reductions | ||
| reductions = _create_reduction_impl(process_props, decomposition_info) | ||
| _reduction_cache[key] = _ReductionCacheEntry( | ||
| reductions=reductions, | ||
| process_props_ref=weakref.ref(process_props), | ||
| decomp_info_ref=weakref.ref(decomposition_info) if decomposition_info is not None else None, | ||
| ) |
There was a problem hiding this comment.
The reductions cache uses strong references for the same reason as the exchange cache, and stale entries are cleaned the same way. The single-pair-per-run lifetime and the test fixture keep this manageable. Adding a distributed-run fail-fast check is a separate follow-up.
🤖 Posted by an agent on behalf of @msimberg
| experiment: definitions.Experiment, | ||
| keep_skip_values: bool, | ||
| allocator: gtx_typing.Allocator, | ||
| process_props: decomposition.ProcessProperties | None = None, |
There was a problem hiding this comment.
It's not clear to me why this should be allowed to be nullable. In general we are explicit. Is this particular use case special in any way that would warrant allowing passing None and defaulting to single node process props instead of being explicit?
There was a problem hiding this comment.
process_props is nullable only as a testing-layer convenience for single-node tests that do not need decomposition information. GridManager.__call__ still requires process_props in production. We recorded this as discussion item R21.
🤖 Posted by an agent on behalf of @msimberg
| @functools.singledispatch | ||
| def _create_reduction_impl( | ||
| process_props: ProcessProperties, | ||
| decomposition_info: DecompositionInfo | None = None, |
There was a problem hiding this comment.
In which cases can/should DecompositionInfo be None? Why is that allowed or necessary? Same below in create_reduction.
There was a problem hiding this comment.
DecompositionInfo is None only on the single-node fallback path in _create_exchange_impl/_create_reduction_impl. MPI implementations with comm_size > 1 require a real DecompositionInfo and raise if one is missing.
🤖 Posted by an agent on behalf of @msimberg
|
|
||
|
|
||
| @functools.singledispatch | ||
| def _create_exchange_impl( |
There was a problem hiding this comment.
General question about the changes in this file: why does register/dispatch/etc. setup need changing? What purpose does it serve and why was the old setup not sufficient?
There was a problem hiding this comment.
The previous setup had separate singledispatch registrations for create_single_node_exchange and create_single_reduction_exchange; those were removed and the single-node fallback moved into _create_exchange_impl/_create_reduction_impl. create_exchange/create_reduction are now caching wrappers that forward .register to the underlying implementation so MPI-specific registrations still work.
🤖 Posted by an agent on behalf of @msimberg
…rops-exchange-reduction-cleanup # Conflicts: # model/atmosphere/advection/tests/advection/integration_tests/test_advection.py # model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py
🤖 Posted by an agent on behalf of @msimberg |
|
When developing, you can test your changes on CSCS CI before merge with the You can pass options to override pipeline variables, for example:
Available options are:
For each option, See The Merging Once your PR is approved and ready for merging, add it to the merge queue. The Optional Tests To run benchmarks you can use:
For more detailed information please look at CI in the EXCLAIM universe. |
|
cscs-ci run default |
Opened by an agent.
This PR cleans up the icon4py APIs that currently accept an
exchangeobject and/or aglobal_reductionsobject alongside (or instead of)process_props. Wherever it makes sense, these constructors now take onlyprocess_propsand derive exchange and reductions internally.process_propsis the communicator-like object that already contains everything needed to build both halo-exchange and global-reduction facilities, so passing derived objects through multiple layers was redundant.The production APIs that changed are
GridGeometry,InterpolationFieldsFactory,MetricsFieldsFactory,Icon4pyDriver, andcreate_static_field_factories. Each now takes a requiredprocess_propsargument and uses the existingcreate_exchange/create_reductionfactories internally. To keep the shared GHEX communication object from being duplicated across factories that use the sameprocess_props, exchange and reductions instances are cached per(process_props, decomposition_info)with weak-reference stale-entry cleanup.The
single_node_exchangeandsingle_node_reductionsmodule-level singletons, plus theircreate_single_node_*factory registrations, are removed. Test call sites and helpers are updated to constructSingleNodeProcessProperties()orSingleNodeExchange()directly where needed.Leaf consumers that already receive a derived
exchange(includingDiffusion,SolveNonhydro, the advection constructors, the states/provider plumbing, andinitialize_granules) are intentionally left unchanged. They still takeexchangebecause they use it directly rather than constructing one.