Split up perf configuration methods

.buildkite/pipeline.yml

@@ -568,7 +568,7 @@ steps:
       - label: ":computer: checkbounds"
         command: >
           julia --color=yes --check-bounds=yes --project=perf perf/benchmark.jl
-          --config_file $PERF_CONFIG_PATH/checkbounds.yml
+          checkbounds
         artifact_paths: "checkbounds/*"
         agents:
           slurm_mem: 20GB
@@ -959,14 +959,14 @@ steps:
       - label: ":computer: Benchmark: perf target (default)"
         command: >
           julia --color=yes --project=perf perf/benchmark.jl
-          --config_file $PERF_CONFIG_PATH/bm_perf_target.yml
+          bm_perf_target
         agents:
           slurm_mem: 20GB
 
       - label: ":computer: Benchmark: perf target (Threaded)"
         command: >
           julia --color=yes --threads 8 --project=perf perf/benchmark.jl
-          --config_file $PERF_CONFIG_PATH/bm_perf_target_threaded.yml
+          bm_perf_target_threaded
         agents:
           slurm_mem: 20GB
           slurm_cpus_per_task: 8
@@ -975,47 +975,47 @@ steps:
       - label: ":fire: Flame graph: perf target (default)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file $PERF_CONFIG_PATH/flame/perf_target.yml
+          default flame_perf_target
         artifact_paths: "flame_perf_target/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: perf target (with tracers)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file  $PERF_CONFIG_PATH/flame/perf_target_tracers.yml
+          flame_perf_target_tracers
         artifact_paths: "flame_perf_target_tracers/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: perf target (edmfx)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file $PERF_CONFIG_PATH/flame/perf_target_edmfx.yml
+          sphere_baroclinic_wave_rhoe_equilmoist_edmfx flame_perf_target_edmfx
         artifact_paths: "flame_perf_target_edmfx/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: perf target (diagnostic edmfx)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file  $PERF_CONFIG_PATH/flame/perf_target_diagnostic_edmfx.yml
+          diagnostic_edmfx_aquaplanet_tke flame_perf_target_diagnostic_edmfx
         artifact_paths: "flame_perf_target_diagnostic_edmfx/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: perf target (edmf)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file  $PERF_CONFIG_PATH/flame/perf_target_edmf.yml
+          sphere_baroclinic_wave_rhoe_equilmoist_edmf flame_perf_target_edmf
         artifact_paths: "flame_perf_target_edmf/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: perf target (Threaded)"
         command: >
           julia --threads 8 --color=yes --project=perf perf/flame.jl
-          --config_file  $PERF_CONFIG_PATH/flame/perf_target_threaded.yml
+          default flame_perf_target_threaded
         artifact_paths: "flame_perf_target_threaded/*"
         agents:
           slurm_cpus_per_task: 8
@@ -1024,23 +1024,23 @@ steps:
       - label: ":fire: Flame graph: perf target (Callbacks)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file  $PERF_CONFIG_PATH/flame/perf_target_callbacks.yml
+          flame_perf_target_callbacks
         artifact_paths: "flame_perf_target_callbacks/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: gravity wave"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file $PERF_CONFIG_PATH/flame/perf_gw.yml
+          flame_perf_gw
         artifact_paths: "flame_perf_gw/*"
         agents:
           slurm_mem: 20GB
 
       - label: ":fire: Flame graph: perf target (diagnostics)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
-          --config_file $PERF_CONFIG_PATH/flame/diagnostics.yml
+          flame_perf_diagnostics
         artifact_paths: "flame_perf_diagnostics/*"
         agents:
           slurm_mem: 20GB
@@ -1049,15 +1049,14 @@ steps:
       - label: ":rocket: JET n-failures (inference)"
         command: >
           julia --color=yes --project=perf perf/jet_test_nfailures.jl
-          --config_file $PERF_CONFIG_PATH/jet_n_failures.yml
+          jet_n_failures sphere_baroclinic_wave_rhoe_equilmoist_edmf
         agents:
           slurm_mem: 20GB
 
       # Latency
       - label: ":mag::rocket: Invalidations"
         command: >
           julia --color=yes --project=perf perf/invalidations.jl
-          --config_file $PERF_CONFIG_PATH/invalidations.yml
         artifact_paths: "invalidations/*"
         agents:
           slurm_mem: 20GB

config/default_configs/default_config.yml

@@ -114,7 +114,7 @@ device:
   help: "Device type to use [`auto` (default) `CPUSingleThreaded`, `CPUMultiThreaded`, `CUDADevice`]"
   value: "auto"
 target_job:
-  help: "An (optional) job to target for analyzing performance"
+  help: "Job config to use in `AtmosCoveragePerfConfig`. If the target config has argument overlapping with the perf defaults, the perf defaults take priority."
   value: ~
 reference_job_id:
   help: |-

perf/benchmark.jl

@@ -1,7 +1,19 @@
 import Random
 Random.seed!(1234)
 import ClimaAtmos as CA
-config = CA.AtmosCoveragePerfConfig()
+
+include("common.jl")
+
+target_job = ARGS[1]
+job_id = get(ARGS, 2, nothing)
+if isnothing(job_id)
+    job_id = target_job
+end
+
+config_dict =
+    target_job != "default" ? CA.config_from_target_job(target_job) :
+    CA.default_config_dict()
+config = AtmosCoveragePerfConfig(; config_dict)
 integrator = CA.get_integrator(config)
 
 (; parsed_args) = config

perf/benchmark_dump.jl

@@ -1,13 +1,3 @@
-#=
-```
-julia --project=examples perf/benchmark_dump.jl
-```
-Or, interactively,
-```
-julia --project=examples
-include(joinpath("perf", "benchmark_dump.jl"));
-```
-=#
 import Random
 Random.seed!(1234)
 import ClimaAtmos as CA
@@ -17,16 +7,16 @@ using Plots
 using PrettyTables
 import YAML
 
+# Need to generate config_dict here to override `h_elem` in the loop below
 parsed_args = CA.parse_commandline(CA.argparse_settings())
 config_dict = YAML.load_file(parsed_args["config_file"])
 output_dir = joinpath(config_dict["job_id"])
 
 steptimes = []
-
 # Iterate through varying number of horizontal elements
 for h_elem in 8:8:40
     config_dict["h_elem"] = h_elem
-    config = CA.AtmosConfig(; config_dict = config_dict)
+    config = CA.AtmosConfig(config_dict)
     integrator = CA.get_integrator(config)
     Y₀ = deepcopy(integrator.u)
 

perf/common.jl

@@ -0,0 +1,49 @@
+import YAML
+
+"""
+    AtmosCoveragePerfConfig()
+    AtmosCoveragePerfConfig(; config_dict)
+Creates a model configuration that covers many physics components.
+The configuration precedence is as follows:
+1. Configuration from the given config file/dict (highest precedence)
+2. Default perf configuration (to increase coverage)
+3. Default configuration (lowest precedence)
+
+There is an exception for the job `flame_perf_target_edmf`, which requires 
+`precip_model` and `rad` to be `nothing` and `apply_limiter` to be `false`.
+"""
+function AtmosCoveragePerfConfig(; config_dict = CA.default_config_dict())
+    perf_default_config = perf_config_dict()
+    config_dict = merge(config_dict, perf_default_config)
+    if get(config_dict, "job_id", "") == "flame_perf_target_edmf"
+        config_dict["precip_model"] = nothing
+        config_dict["rad"] = nothing
+        config_dict["apply_limiter"] = false
+        @show config_dict["job_id"]
+    end
+    @show config_dict
+    return CA.AtmosConfig(config_dict)
+end
+
+"""
+    TargetJobConfig(target_job)
+Creates a full model configuration from the given target job.
+"""
+TargetJobConfig(target_job) =
+    CA.AtmosConfig(CA.config_from_target_job(target_job))
+
+
+"""
+    perf_config_dict()
+Loads the default performance configuration from a file into a Dict.
+"""
+function perf_config_dict()
+    perf_defaults = joinpath(
+        dirname(@__FILE__),
+        "..",
+        "config",
+        "default_configs",
+        "default_perf.yml",
+    )
+    return YAML.load_file(perf_defaults)
+end

perf/flame.jl

@@ -1,21 +1,34 @@
 import Random
 Random.seed!(1234)
 import ClimaAtmos as CA
-config = CA.AtmosCoveragePerfConfig()
-integrator = CA.get_integrator(config)
 
+include("common.jl")
+
+target_job = ARGS[1]
+job_id = get(ARGS, 2, nothing)
+if isnothing(job_id)
+    job_id = target_job
+end
+
+config_dict =
+    target_job != "default" ? CA.config_from_target_job(target_job) :
+    CA.default_config_dict()
+
+# Need to set internal job_id for diagnostics saved by model
+config_dict["job_id"] = job_id
+
+config = AtmosCoveragePerfConfig(; config_dict)
+integrator = CA.get_integrator(config)
 # The callbacks flame graph is very expensive, so only do 2 steps.
 @info "running step"
 
 import SciMLBase
 SciMLBase.step!(integrator) # compile first
 CA.call_all_callbacks!(integrator) # compile callbacks
 import Profile, ProfileCanvas
-(; output_dir, job_id) = integrator.p.simulation
 output_dir = job_id
 mkpath(output_dir)
 
-
 @info "collect profile"
 Profile.clear()
 prof = Profile.@profile SciMLBase.step!(integrator)

perf/jet_test_nfailures.jl

@@ -1,11 +1,22 @@
 import Random
 Random.seed!(1234)
 import ClimaAtmos as CA
-config = CA.AtmosCoveragePerfConfig()
+
+include("common.jl")
+
+target_job = ARGS[1]
+job_id = get(ARGS, 2, nothing)
+if isnothing(job_id)
+    job_id = target_job
+end
+
+config_dict =
+    target_job != "default" ? CA.config_from_target_job(target_job) :
+    CA.default_config_dict()
+config = AtmosCoveragePerfConfig(; config_dict)
 integrator = CA.get_integrator(config)
 
 import JET
-
 import SciMLBase
 SciMLBase.step!(integrator) # Make sure no errors
 

src/solver/types.jl

@@ -389,56 +389,39 @@ struct AtmosConfig{FT, TD, PA, C}
     comms_ctx::C
 end
 
+Base.eltype(::AtmosConfig{FT}) where {FT} = FT
+
 """
-    AtmosCoveragePerfConfig()
-    AtmosCoveragePerfConfig(; s, config_dict)
-Creates a model configuration for many performance tests.
-Creates a config from the following in order of top priority to last:
-1. Configuration from the given config file/dict
-2. Default perf configuration
-3. Target job configuration
-4. Default configuration
+    AtmosConfig(config_file::String)
+
+Helper function for the AtmosConfig constructor. Reads a YAML file into a Dict
+and passes it to the AtmosConfig constructor.
 """
-function AtmosCoveragePerfConfig(;
-    s = argparse_settings(),
-    config_dict = nothing,
-)
-    parsed_args = parse_commandline(s)
-    if isnothing(config_dict)
-        config_dict = YAML.load_file(parsed_args["config_file"])
-    end
-    target_job_config = if haskey(config_dict, "target_job")
-        config_from_target_job(config_dict["target_job"])
-    else
-        Dict()
-    end
-    perf_defaults = joinpath(
-        dirname(@__FILE__),
-        "..",
-        "..",
-        "config",
-        "default_configs",
-        "default_perf.yml",
-    )
-    perf_default_config = YAML.load_file(perf_defaults)
-    config_dict = merge(target_job_config, perf_default_config, config_dict)
-    return AtmosConfig(; s, config_dict)
+function AtmosConfig(config_file::String; comms_ctx = nothing)
+    config = YAML.load_file(config_file)
+    return AtmosConfig(config; comms_ctx)
 end
 
-function AtmosConfig(;
-    s = argparse_settings(),
-    parsed_args = parse_commandline(s),
-    config_dict = nothing,
-    comms_ctx = nothing,
-)
-    config = if !isnothing(config_dict)
-        override_default_config(config_dict)
-    elseif !isnothing(parsed_args["config_file"])
-        override_default_config(parsed_args["config_file"])
-    else
-        @info "Using default configuration"
-        default_config_dict()
-    end
+"""
+    AtmosConfig(; comms_ctx = nothing)
+Helper function for the AtmosConfig constructor.
+Reads the `config_file` from the command line into a Dict
+and passes it to the AtmosConfig constructor.
+"""
+function AtmosConfig(; comms_ctx = nothing)
+    parsed_args = parse_commandline(argparse_settings())
+    return AtmosConfig(parsed_args["config_file"]; comms_ctx)
+end
+
+AtmosConfig(::Nothing; comms_ctx = nothing) = AtmosConfig(Dict(); comms_ctx)
+
+"""
+    AtmosConfig(config::Dict; comms_ctx = nothing)
+Constructs the AtmosConfig from the Dict passed in. This Dict overrides all of 
+the default configurations set in `default_config_dict()`.
+"""
+function AtmosConfig(config::Dict; comms_ctx = nothing)
+    config = override_default_config(config)
     FT = config["FLOAT_TYPE"] == "Float64" ? Float64 : Float32
     toml_dict = CP.create_toml_dict(
         FT;
@@ -465,7 +448,6 @@ function AtmosConfig(;
     PA = typeof(config)
     return AtmosConfig{FT, TD, PA, C}(toml_dict, config, comms_ctx)
 end
-Base.eltype(::AtmosConfig{FT}) where {FT} = FT
 
 # In order to specify C2F operator boundary conditions with 0 instead of FT(0),
 # we need to tell ClimaCore how to convert AxisTensor components from Int to FT.