Add driver bencher benchmarks - #1450
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds continuous benchmarking coverage for the distributed standalone driver (MPI), integrating new pytest-benchmark-based timings into the existing Nox/Bencher benchmarking workflows and CSCS CI pipelines.
Changes:
- Adds MPI driver benchmark tests measuring initialization, timeloop, and total runtime, with new pytest CLI options to parameterize experiment/steps/rounds.
- Introduces shared benchmark helpers (rank resolution, upload-rank gating, grid override validation) and unit tests for them.
- Wires new intra-node and inter-node MPI driver benchmark jobs into the existing CSCS Bencher baseline/feature pipelines and excludes
benchmark_onlytests from default test runs.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
noxfile.py |
Adds MPI driver benchmark session plus Bencher upload sessions; updates MPI test selection to exclude benchmark_only. |
model/testing/src/icon4py/model/testing/benchmark.py |
Adds shared benchmark utilities (rank resolution, upload gating, grid override validation). |
model/testing/src/icon4py/model/testing/pytest_hooks.py |
Registers new --driver-benchmark-* pytest options used by the new benchmarks. |
model/testing/tests/testing/unit_tests/test_benchmark.py |
Adds unit tests validating the new benchmark helper behaviors. |
model/driver/tests/driver/mpi_tests/test_benchmark_driver.py |
Adds three MPI pytest-benchmark tests for driver init, timeloop, and total run timing. |
.cscs-ci/benchmark_bencher.yml |
Adds Bencher feature-branch MPI driver benchmark jobs (intra/inter-node). |
.cscs-ci/benchmark_bencher_baseline.yml |
Adds Bencher baseline MPI driver benchmark jobs (intra/inter-node). |
Suppressed comments (1)
noxfile.py:257
- The benchmark JSON is generated under
model/driver/(seebenchmark_driver_mpisessionchdir). This bencher upload command refers to the file without that directory, so it likely won’t find the JSON to upload.
--github-actions {os.environ['GD_COMMENT_TOKEN']} \
--ci-number {os.environ['PR_ID']} \
--ci-id run-{bencher_testbed.replace(':', '_')}-{int(datetime.now().strftime('%Y%m%d%H%M%S%f'))} \
--file pytest_benchmark_results_{session.python}_{rank}.json".split(),
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
msimberg
left a comment
There was a problem hiding this comment.
High level comments, no comments on implementation itself yet.
…ariants [C1/C7][C2]
…op top-level icon4py import [C1/C7][C4][C5]
Move the prognostic/tracer/nonhydro/diagnostic state creation, initial condition application, DriverStates assembly and consistency validation that followed initialize_driver into a reusable helper. Both run_driver and the upcoming benchmarks will call it, eliminating the duplication between the driver entry point and the benchmark tests.
…iver benchmark Drop the configurable --driver-benchmark-* options from the MPI benchmark module and replace them with module constants: JW experiment, 100 steps, 5 rounds and 2 warmup rounds. Explicitly rescale dtime to 50 s so the run is stable on the R02B06_GLOBAL override. Remove the limited-area grid presets and the validate_grid_override gate now that dtime rescaling is explicit. Use the new initialize_driver_states helper from the driver library for the timeloop benchmark setup.
Add a non-MPI driver benchmark file under integration_tests with init, timeloop and total measurements. It uses the shared initialize_driver_states helper, hardcodes JW on R02B06_GLOBAL with 100 steps and a 50 s dtime, and runs on the default single-node process_props fixture.
…mark options Drop the icon4py.model.testing.benchmark helpers and their unit tests. The rank-gate helpers move back into noxfile.py as inline private functions used only by the MPI upload sessions. Remove the four --driver-benchmark-* pytest options from pytest_hooks.py now that the benchmarks hardcode their parameters.
…parameters Drop the configurable --driver-benchmark-* arguments from the CI driver nox job now that the benchmark module hardcodes JW, 100 steps, 5 rounds and 2 warmup rounds. Rename the MPI driver jobs to remove the 'inter' segment and add a BACKEND matrix so the backend appears in the CI job name. Add matching serial single-rank driver bencher jobs that run the new integration-tests benchmark without MPI.
… and simplify testbeds - Move experiment selection to a test-level pytest.mark.parametrize(..., ids=[e.name]) in driver benchmark tests. - Drop the experiment segment from the bencher testbed strings built in noxfile.py. - Extend _name_from_fullname in pytest_hooks.py to shorten function-based benchmark names as well as StencilTest class names.
…chmark trial Add rules: - when: never to the four existing serial benchmark jobs so the driver-benchmark trial can run without overlapping serial baselines. Revert this commit before merging.
|
cscs-ci run benchmark_bencher |
|
cscs-ci run benchmark-bencher (correction: previous comment used underscore; the registered pipeline name is hyphen) |
|
cscs-ci run benchmark-bencher (re-trigger after merging main; shared partition removed on main) |
|
| Project | icon4py |
| Branch | bencher-distributed-driver |
| Testbed | ci-runner:santis-gh200:dace_gpu:R02B06_GLOBAL |
⚠️ WARNING: No Threshold found!Without a Threshold, no Alerts will ever be generated.
Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the--ci-only-thresholdsflag.
Click to view all benchmark results
| Benchmark | Latency | seconds (s) |
|---|---|---|
| test_benchmark_driver_single_rank_init[False-exclaim_nh35_tri_jws] | 📈 view plot | 59.14 s |
| test_benchmark_driver_single_rank_timeloop[False-exclaim_nh35_tri_jws] | 📈 view plot | 8.21 s |
| test_benchmark_driver_single_rank_total[False-exclaim_nh35_tri_jws] | 📈 view plot | 69.38 s |
…ata)
Only mpitask{1,2,4} JW archives exist on rgw.cscs.ch; an 8-rank run
403s at download_experiment. JW namelists are byte-identical across
ranks and the benchmark uses an analytic IC (does not read ser_data),
so 4-rank intra-node uses the existing mpitask4 archive. Reverts the
C2 inter-node decision for now; follow-up (Option E) will separate the
experiment description from serialized data.
[skip check-versions due to pre-existing .nox scan issue]
| This wraps the creation of the prognostic/tracer/nonhydro/diagnostic states, the | ||
| application of the initial condition, the assembly of ``DriverStates``, and the | ||
| consistency check that precedes the time loop. It is shared between ``run_driver`` | ||
| and the driver benchmarks. |
There was a problem hiding this comment.
This comment shouldn't say who uses it.
| def initialize_driver_states( | ||
| icon4py_driver: Icon4pyDriver, | ||
| allocator: gtx.typing.Allocator, | ||
| ) -> driver_states.DriverStates: |
There was a problem hiding this comment.
Even though it's only two parameters these might still be worth passing by keyword.
| @@ -0,0 +1,250 @@ | |||
| # ICON4Py - ICON inspired code in Python and GT4Py | |||
There was a problem hiding this comment.
This file should drop the single_rank suffix.
| indirect=True, | ||
| ids=[e.name for e in BENCHMARK_EXPERIMENTS], | ||
| ) | ||
| def test_benchmark_driver_single_rank_init( |
There was a problem hiding this comment.
This also needs to drop single_rank
| "driver_benchmark_experiment", | ||
| BENCHMARK_EXPERIMENTS, | ||
| indirect=True, | ||
| ids=[e.name for e in BENCHMARK_EXPERIMENTS], |
There was a problem hiding this comment.
This is a bit silly. It should be lambda e: e.name.
| if mpi_decomposition.mpi4py is None: | ||
| pytest.skip( | ||
| "Skipping parallel driver benchmark tests on single-node installation", | ||
| allow_module_level=True, | ||
| ) |
| _name_from_fullname_pattern = re.compile( | ||
| # StencilTest fullnames (``<path>::<class_name>::test_stencil[<variant>]``) are | ||
| # shortened to ``<class_name>[<variant>]``. Function-based benchmark fullnames | ||
| # (``<path>::test_name[<params>]``) are shortened to ``<test_name>[<params>]``. | ||
| _stenciltest_name_pattern = re.compile( | ||
| r""" | ||
| ::(?P<class>[A-Za-z_]\w*) # capture class name | ||
| (?::: [A-Za-z_]\w* # skip method name | ||
| (?:\[(?P<params>.+)\])? ) # optional parameterization; allow '[' and ']' inside params | ||
| """, | ||
| ::(?P<class>[A-Za-z_]\w*) # capture class name | ||
| (?::: [A-Za-z_]\w* # skip method name | ||
| (?:\[(?P<params>.+)\])? ) # optional parameterization | ||
| """, | ||
| re.VERBOSE, | ||
| ) | ||
| _function_name_pattern = re.compile( | ||
| r""" | ||
| ::(?P<func>[A-Za-z_]\w*) # trailing function name | ||
| (?:\[(?P<params>.+)\])? # optional parameterization | ||
| $ | ||
| """, |
There was a problem hiding this comment.
Double-check that this makes sense.
|
@edopao @iomaganaris @yiluchen1066 @OngChia this is not ready for merging, but I'd appreciate an early quick review from any of you. Two things I'd especially appreciate feedback on is 1. the bencher setup with testbed, test naming, etc. if it makes sense to you with more bencher experience and 2. whether you think the current test seems sane. I've used JW with R02B06 as a decent compromise between between size and time. I'm running it on one GPU and on 4 GPUs, all intra-node. For 8 ranks I bumped into the problem in #1456 which I don't want to resolve for this work yet, so staying intra-node for now. There's a small speedup, but not perfect scaling, so I'm hoping this would be a decent size to see improvements to both scaling (halo exchange overheads) and general performance improvements. The benchmarks run now init, one timestep, and the full simulation with init+timestepping as three separate benchmarks, with the full simulation taking ~1 minute (with 100 timesteps). I've yet to enable other backends than dace_gpu. |
Second sweep point: 64 threads gave only a marginal improvement over 72 (single-rank init 79.4s -> 77.9s); try 32.
|
cscs-ci run benchmark-bencher OMP sweep run 2: OMP_NUM_THREADS=32 |
|
| Project | icon4py |
| Branch | bencher-distributed-driver |
| Testbed | ci-runner:santis-gh200:dace_gpu:R02B06_GLOBAL:1N4R:nccl |
⚠️ WARNING: No Threshold found!Without a Threshold, no Alerts will ever be generated.
Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the--ci-only-thresholdsflag.
Click to view all benchmark results
| Benchmark | Latency | seconds (s) |
|---|---|---|
| test_benchmark_driver_init[exclaim_nh35_tri_jws-True] | 📈 view plot | 47.86 s |
| test_benchmark_driver_timeloop[exclaim_nh35_tri_jws-True] | 📈 view plot | 9.05 s |
| test_benchmark_driver_total[exclaim_nh35_tri_jws-True] | 📈 view plot | 59.49 s |
|
| Project | icon4py |
| Branch | bencher-distributed-driver |
| Testbed | ci-runner:santis-gh200:dace_gpu:R02B06_GLOBAL |
⚠️ WARNING: No Threshold found!Without a Threshold, no Alerts will ever be generated.
Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the--ci-only-thresholdsflag.
Click to view all benchmark results
| Benchmark | Latency | seconds (s) |
|---|---|---|
| test_benchmark_driver_single_rank_init[False-exclaim_nh35_tri_jws] | 📈 view plot | 76.84 s |
| test_benchmark_driver_single_rank_timeloop[False-exclaim_nh35_tri_jws] | 📈 view plot | 8.30 s |
| test_benchmark_driver_single_rank_total[False-exclaim_nh35_tri_jws] | 📈 view plot | 87.73 s |
Third sweep point: 32/64/72 are within 3% of each other; the old 60.8s single-rank init at OMP=8 was measured with MPS enabled, so re-measure OMP=8 with MPS disabled to isolate the cause.
|
cscs-ci run benchmark-bencher OMP sweep run 3: OMP_NUM_THREADS=8 (MPS disabled) |
|
| Project | icon4py |
| Branch | bencher-distributed-driver |
| Testbed | ci-runner:santis-gh200:dace_gpu:R02B06_GLOBAL:1N4R:nccl |
⚠️ WARNING: No Threshold found!Without a Threshold, no Alerts will ever be generated.
Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the--ci-only-thresholdsflag.
Click to view all benchmark results
| Benchmark | Latency | seconds (s) |
|---|---|---|
| test_benchmark_driver_init[exclaim_nh35_tri_jws-True] | 📈 view plot | 47.13 s |
| test_benchmark_driver_timeloop[exclaim_nh35_tri_jws-True] | 📈 view plot | 9.12 s |
| test_benchmark_driver_total[exclaim_nh35_tri_jws-True] | 📈 view plot | 60.10 s |
|
cscs-ci run benchmark-bencher re-run: serial job died at container startup (bwrap uid map permission denied, runner-side flake); MPI passed. Re-measuring OMP_NUM_THREADS=8. |
|
| Project | icon4py |
| Branch | bencher-distributed-driver |
| Testbed | ci-runner:santis-gh200:dace_gpu:R02B06_GLOBAL:1N4R:nccl |
⚠️ WARNING: No Threshold found!Without a Threshold, no Alerts will ever be generated.
Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the--ci-only-thresholdsflag.
Click to view all benchmark results
| Benchmark | Latency | seconds (s) |
|---|---|---|
| test_benchmark_driver_init[exclaim_nh35_tri_jws-True] | 📈 view plot | 47.55 s |
| test_benchmark_driver_timeloop[exclaim_nh35_tri_jws-True] | 📈 view plot | 8.87 s |
| test_benchmark_driver_total[exclaim_nh35_tri_jws-True] | 📈 view plot | 62.97 s |
|
| Project | icon4py |
| Branch | bencher-distributed-driver |
| Testbed | ci-runner:santis-gh200:dace_gpu:R02B06_GLOBAL |
⚠️ WARNING: No Threshold found!Without a Threshold, no Alerts will ever be generated.
Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the--ci-only-thresholdsflag.
Click to view all benchmark results
| Benchmark | Latency | seconds (s) |
|---|---|---|
| test_benchmark_driver_single_rank_init[False-exclaim_nh35_tri_jws] | 📈 view plot | 77.21 s |
| test_benchmark_driver_single_rank_timeloop[False-exclaim_nh35_tri_jws] | 📈 view plot | 8.14 s |
| test_benchmark_driver_single_rank_total[False-exclaim_nh35_tri_jws] | 📈 view plot | 87.54 s |
OMP sweep on the serial job (MPS disabled, dedicated GPU): init is 76.8-79.4s for 8/32/64/72 threads, i.e. flat. The old 60.8s at OMP=8 (iid=388) predates the CPUS_PER_TASK/GPUS_PER_TASK pinning and is not reproducible with the current placement; the remaining init difference is therefore not an OMP effect. Keep 72 (full core count per rank).
| - when: never # TEMPORARY: disabled for driver bencher trial; restore before merge | ||
|
|
||
| benchmark_bencher_driver_mpi_feature_aarch64: | ||
| extends: [.bencher_driver_mpi_feature_tests, .retry_on_transient_failure] |
There was a problem hiding this comment.
Not directly related to this PR. I was not aware of .retry_on_transient_failure, interesting. I see it is used in multiple places (in other fiels as well) and I wonder if a refactoring PR could move it down to some base config such as .test_runner_mpi and .test_runner_serial.
There was a problem hiding this comment.
Yeah, that might be a good idea!
FYI, I introduced it to restart jobs that take long, e.g. because of JIT compilation taking longer than expected if filesystems are unhappy. It does come with false positives, i.e. if there's an actual hang caused by our bugs (or e.g. GHEX) the job would get restarted as well, but I figured it's a decent tradeoff since infrastructure-related hangs seem to be more common.
| SLURM_JOB_NUM_NODES: 1 | ||
| ICON4PY_TEST_MPI_SUBCOMM_SIZE: $SLURM_NTASKS | ||
| rules: | ||
| - if: $ICON4PY_DISABLE_DRIVER_BENCHMARK == "1" || $ICON4PY_DISABLE_DRIVER_BENCHMARK == "true" |
There was a problem hiding this comment.
Is the variable ICON4PY_DISABLE_DRIVER_BENCHMARK only temporary? If we want a CI job that is triggered manually I think we can use when: manual, and there will be a "button" in the CI-Ext page to trigger it. I did this a long time ago, I don't know if it still works like that.
There was a problem hiding this comment.
Yes, temporary! For testing I only wanted the driver benchmark to be run and I haven't made up my mind if it makes sense to always run the driver benchmark on benchmark-bencher pipelines. With a warm cache it's not so expensive, but it's also not for free. Opinions? If we actually want to make it optional, we probably need a similar setup as we have now for the regular tests, i.e. something where it's possible to choose what runs in the cscs-ci run comment.
| dt_utils.download_experiment(experiment, process_props) | ||
| experiment_path = dt_utils.get_path_for_experiment(experiment, process_props) | ||
| config = driver_config.read_experiment_config_from_fortran(experiment_path) | ||
| return config.with_overrides( |
There was a problem hiding this comment.
I guess this is where I should plugin the External workspace memory for dace benchmarks on AMD GPUs.
Review feedback (PR C2SM#1450, 2026-09-03/04 round): - Drop the single_rank suffix from the serial benchmark test module and its test functions; ids=lambda e: e.name in both benchmark files. - Remove the redundant module-level mpi4py skip guard from the MPI benchmark test (convention from C2SM#1453). - Pass initialize_driver_states arguments by keyword; docstring no longer names its callers. - Trim CI comments to the non-obvious facts (MPS for shared GPUs, allow-task-sharing for NCCL peer access); drop the venv activation line from the bencher before_script overrides (.venv/bin is on PATH in the CI image). - Drop ICON4PY_DISABLE_DRIVER_BENCHMARK; the driver bencher jobs are no longer disableable. - BACKEND matrix extended to all four backends; GRID moved into the parallel matrix (job names now carry the grid for GT4Py cache separation). Not changed: NCCL env vars are already fully present in .test_runner_base (from C2SM#1356, sourced from alps-extended-images); NCCL_NCHANNELS_PER_NET_PEER differs (1 here vs 4 upstream) and is left as is, flagged for discussion.
|
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. |
| stage: benchmark | ||
| script: | ||
| - .cscs-ci/scripts/ci-mpi-wrapper.sh nox -s "${NOX_SESSION}" -- --backend=$BACKEND --grid=$GRID -k "${TEST_SELECTION}" | ||
| # start-cuda-mps.sh is for shared GPUs; these jobs have one GPU per task. |
There was a problem hiding this comment.
| # start-cuda-mps.sh is for shared GPUs; these jobs have one GPU per task. | |
| # MPS is explicitly not enabled here. The job uses one GPU per task. |
| - source .cscs-ci/scripts/gt4py-cache.sh | ||
| extends: [.benchmark_driver_base_variables] | ||
| variables: | ||
| # allow-task-sharing lets NCCL open peer GPUs for GPU-GPU communication. |
There was a problem hiding this comment.
| # allow-task-sharing lets NCCL open peer GPUs for GPU-GPU communication. |
| stage: benchmark | ||
| script: | ||
| - nox -s "${NOX_SESSION}" -- --backend=$BACKEND --grid=$GRID -k "${TEST_SELECTION}" | ||
| # start-cuda-mps.sh is for shared GPUs; this job has one GPU per task. |
There was a problem hiding this comment.
| # start-cuda-mps.sh is for shared GPUs; this job has one GPU per task. | |
| # MPS is explicitly not enabled here. The job uses one GPU per task. |
This PR extends the bencher continuous-benchmarking pipeline with end-to-end standalone driver benchmarks. The new jobs time the full driver for the JW global experiment across initialization, the time integration loop, and the total run. Both distributed MPI and single-rank variants are added.
The distributed benchmark lives in
model/driver/tests/driver/mpi_tests/test_benchmark_driver.pyand is selected through thebenchmark_driver_mpinox session. It runs with fixed pytest-benchmarkpedanticrounds,iterations=1, andwarmup_roundsso compilation stays out of the timed path and every rank executes the same number of rounds. Barriers bracket each timed callable to keep ranks synchronized, and only rank 0 uploads results. The current CI job uses 4 ranks on 1 node with 1 GPU per rank because onlympitask{1,2,4}JW archives are available.A single-rank variant is added in
model/driver/tests/driver/integration_tests/test_benchmark_driver.pyand run through thebenchmark_drivernox session. Both variants expose the same three benchmarks astest_benchmark_driver_init,test_benchmark_driver_timeloop, andtest_benchmark_driver_total, and both hardcode the JW experiment, 100 timesteps, a 50 s timestep, 5 rounds, and 2 warmup rounds. The earlier--driver-benchmark-*pytest options and theicon4py.model.testing.benchmarkhelper module have been removed.run_drivernow shares a newinitialize_driver_stateshelper with the benchmarks so the timeloop measurement builds fresh driver states per round in the same way the full run does.The bencher testbed for the MPI benchmark records the runner, system tag, backend, grid, node/rank layout, and GHEX transport backend so MPI and NCCL transport runs do not merge under one baseline. The single-rank testbed matches the existing serial stencil/granule shape. CI jobs follow the
test_model/test_model_mpiconventions, pin OpenMP threads, and disable MPS because each rank owns a dedicated GPU. The CI runs JW on the R02B06 global grid for all four backends (dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu). The four pre-existing serial bencher jobs for stencils and granules are temporarily disabled withrules: when: neverso this branch can run an isolated driver-benchmark trial; they must be restored before merge.OMP_NUM_THREADSwas swept for the init-heavy driver benchmark on the target GH200 nodes (8, 32, 64, 72 with MPS off, one dedicated GPU per task, andSLURM_CPUS_PER_TASK=72). Serial init stayed flat within about 3% across the sweep, and MPI init was likewise flat. The remaining init differences are attributable to CPU/GPU placement rather than the thread count, so the jobs keepOMP_NUM_THREADS: 72(full cores per rank, pinned for determinism). NCCL and libfabric environment variables are inherited from.test_runner_basein.cscs-ci/base.yml, which is sourced from alps-extended-images; the driver bencher jobs add no NCCL overrides.NCCL_MAX_CTAS: 16,NCCL_NET: "AWS Libfabric", andNCCL_NCHANNELS_PER_NET_PEER: 1are kept deliberately.