From 210ffff565de0d47dc6906e887bcc2113c9d6700 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Wed, 24 Apr 2024 16:48:25 -0700 Subject: [PATCH] make diagnostics optional --- .buildkite/benchmarks/pipeline.yml | 4 +- config/benchmark_configs/amip_diagedmf.yml | 1 + .../benchmark_configs/gpu_amip_diagedmf.yml | 1 + experiments/AMIP/cli_options.jl | 4 ++ experiments/AMIP/coupler_driver.jl | 65 +++++++++++-------- 5 files changed, 47 insertions(+), 28 deletions(-) diff --git a/.buildkite/benchmarks/pipeline.yml b/.buildkite/benchmarks/pipeline.yml index 8f25f9cef9..a8ad834dd4 100644 --- a/.buildkite/benchmarks/pipeline.yml +++ b/.buildkite/benchmarks/pipeline.yml @@ -76,7 +76,7 @@ steps: agents: slurm_gpus_per_task: 1 slurm_cpus_per_task: 4 - slurm_ntasks: 2 + slurm_ntasks: 4 slurm_mem: 16GB - label: "GPU AMIP with diagnostic EDMF" @@ -86,7 +86,7 @@ steps: agents: slurm_gpus_per_task: 1 slurm_cpus_per_task: 4 - slurm_ntasks: 2 + slurm_ntasks: 4 slurm_mem: 16GB - group: "Generate output table" diff --git a/config/benchmark_configs/amip_diagedmf.yml b/config/benchmark_configs/amip_diagedmf.yml index ac3670b15a..fc890ffe80 100644 --- a/config/benchmark_configs/amip_diagedmf.yml +++ b/config/benchmark_configs/amip_diagedmf.yml @@ -14,3 +14,4 @@ run_name: "amip_diagedmf" start_date: "19790301" t_end: "12hours" turb_flux_partition: "CombinedStateFluxes" +use_coupler_diagnostics: false diff --git a/config/benchmark_configs/gpu_amip_diagedmf.yml b/config/benchmark_configs/gpu_amip_diagedmf.yml index a3b528f579..b798e82967 100644 --- a/config/benchmark_configs/gpu_amip_diagedmf.yml +++ b/config/benchmark_configs/gpu_amip_diagedmf.yml @@ -14,3 +14,4 @@ run_name: "gpu_amip_diagedmf" start_date: "19790301" t_end: "12hours" turb_flux_partition: "CombinedStateFluxes" +use_coupler_diagnostics: false diff --git a/experiments/AMIP/cli_options.jl b/experiments/AMIP/cli_options.jl index 9a1ac6369f..7406c5f1fd 100644 --- a/experiments/AMIP/cli_options.jl +++ b/experiments/AMIP/cli_options.jl @@ -78,6 +78,10 @@ function argparse_settings() help = "Device type to use [`auto` (default) `CPUSingleThreaded`, `CPUMultiThreaded`, `CUDADevice`]" arg_type = String default = "auto" + "--use_coupler_diagnostics" + help = "Boolean flag indicating whether to compute and output coupler diagnostics [`true` (default), `false`]" + arg_type = Bool + default = true # ClimaAtmos specific "--surface_setup" help = "Triggers ClimaAtmos into the coupled mode [`PrescribedSurface` (default)]" # retained here for standalone Atmos benchmarks diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl index 0adb5acb12..c5cd201f60 100644 --- a/experiments/AMIP/coupler_driver.jl +++ b/experiments/AMIP/coupler_driver.jl @@ -129,6 +129,7 @@ restart_dir = config_dict["restart_dir"] restart_t = Int(config_dict["restart_t"]) evolving_ocean = config_dict["evolving_ocean"] dt_rad = config_dict["dt_rad"] +use_coupler_diagnostics = config_dict["use_coupler_diagnostics"] #= ## Setup Communication Context @@ -470,27 +471,39 @@ User can write custom diagnostics in the `user_diagnostics.jl`. Note, this will be replaced by the diagnostics framework currently in ClimaAtmos, once it is abstracted into a more general package, so we can use it to save fields from surface models. =# +if use_coupler_diagnostics + monthly_3d_diags = Diagnostics.init_diagnostics( + (:T, :u, :q_tot, :q_liq_ice), + atmos_sim.domain.center_space; + save = TimeManager.Monthly(), + operations = (; accumulate = Diagnostics.TimeMean([Int(0)])), + output_dir = COUPLER_OUTPUT_DIR, + name_tag = "monthly_mean_3d_", + ) -monthly_3d_diags = Diagnostics.init_diagnostics( - (:T, :u, :q_tot, :q_liq_ice), - atmos_sim.domain.center_space; - save = TimeManager.Monthly(), - operations = (; accumulate = Diagnostics.TimeMean([Int(0)])), - output_dir = dir_paths.output, - name_tag = "monthly_mean_3d_", -) + monthly_3d_diags = Diagnostics.init_diagnostics( + (:T, :u, :q_tot, :q_liq_ice), + atmos_sim.domain.center_space; + save = TimeManager.Monthly(), + operations = (; accumulate = Diagnostics.TimeMean([Int(0)])), + output_dir = dir_paths.output, + name_tag = "monthly_mean_3d_", + ) -monthly_2d_diags = Diagnostics.init_diagnostics( - (:precipitation_rate, :toa_fluxes, :T_sfc, :tubulent_energy_fluxes), - boundary_space; - save = TimeManager.Monthly(), - operations = (; accumulate = Diagnostics.TimeMean([Int(0)])), - output_dir = dir_paths.output, - name_tag = "monthly_mean_2d_", -) + monthly_2d_diags = Diagnostics.init_diagnostics( + (:precipitation_rate, :toa_fluxes, :T_sfc, :tubulent_energy_fluxes), + boundary_space; + save = TimeManager.Monthly(), + operations = (; accumulate = Diagnostics.TimeMean([Int(0)])), + output_dir = dir_paths.output, + name_tag = "monthly_mean_2d_", + ) -diagnostics = (monthly_3d_diags, monthly_2d_diags) -Utilities.show_memory_usage(comms_ctx) + diagnostics = (monthly_3d_diags, monthly_2d_diags) + Utilities.show_memory_usage(comms_ctx) +else + diagnostics = () +end #= ## Initialize Conservation Checks @@ -713,13 +726,14 @@ function solve_coupler!(cs) CO2_current = BCReader.interpolate_midmonth_to_daily(cs.dates.date[1], cs.mode.CO2_info) Interfacer.update_field!(atmos_sim, Val(:co2), CO2_current) - ## calculate and accumulate diagnostics at each timestep - ClimaComms.barrier(comms_ctx) - Diagnostics.accumulate_diagnostics!(cs) - - ## save and reset monthly averages - Diagnostics.save_diagnostics(cs) + ## calculate and accumulate diagnostics at each timestep, if we're using diagnostics in this run + if !isempty(cs.diagnostics) + ClimaComms.barrier(comms_ctx) + Diagnostics.accumulate_diagnostics!(cs) + ## save and reset monthly averages + Diagnostics.save_diagnostics(cs) + end end ## compute global energy @@ -842,9 +856,8 @@ if ClimaComms.iamroot(comms_ctx) end ## plotting AMIP results - if cs.mode.name == "amip" + if cs.mode.name == "amip" && !isempty(cs.diagnostics) ## plot data that correspond to the model's last save_hdf5 call (i.e., last month) - @info "AMIP plots" ## ClimaESM