Merge pull request #3461 from CliMA/ne/slurm_workers

Add slurm workers for calibration end-to-end test

.buildkite/pipeline.yml

@@ -586,8 +586,12 @@ steps:
         command: julia --project=calibration/test calibration/test/interface.jl
       - label: "end to end test"
         command: julia --project=calibration/test calibration/test/e2e_test.jl
+        agents:
+          slurm_ntasks: 10
+          slurm_cpus_per_task: 1
+          slurm_mem: 96GB
+          slurm_time: "00:30:00"
         artifact_paths: "calibration_end_to_end_test/*"
-        soft_fail: true
 
   - group: "Diagnostic EDMFX"
     steps:

.gitignore

@@ -50,6 +50,7 @@ Artifacts.toml
 *.g
 *.lock
 *.loc
+*.out
 
 # misc
 .DS_Store

calibration/model_interface.jl

@@ -10,39 +10,15 @@ import YAML
 import ClimaComms
 ClimaComms.@import_required_backends
 using ClimaUtilities.ClimaArtifacts
-import ClimaCalibrate:
-    set_up_forward_model,
-    run_forward_model,
-    path_to_ensemble_member,
-    ExperimentConfig
-
-"""
-    set_up_forward_model(member, iteration, experiment_dir::AbstractString)
-    set_up_forward_model(member, iteration, ::ExperimentConfig; experiment_dir)
-    set_up_forward_model(member, iteration, config_dict::AbstractDict)
-
-Return an AtmosConfig object for the given member and iteration.
-
-Turns off default diagnostics and sets the TOML parameter file to the member's path.
-This assumes that the  config dictionary has an `output_dir` key.
-"""
-function set_up_forward_model(
-    member,
-    iteration,
-    ::ExperimentConfig;
-    experiment_dir = dirname(Base.active_project()),
-)
-    # Assume experiment_dir is project dir
-    config_dict = YAML.load_file(joinpath(experiment_dir, "model_config.yml"))
-    set_up_forward_model(member, iteration, config_dict::AbstractDict)
-end
-
-function set_up_forward_model(member, iteration, experiment_dir::AbstractString)
-    config_dict = YAML.load_file(joinpath(experiment_dir, "model_config.yml"))
-    set_up_forward_model(member, iteration, config_dict::AbstractDict)
-end
-
-function set_up_forward_model(member, iteration, config_dict::AbstractDict)
+import ClimaCalibrate: forward_model, path_to_ensemble_member
+import ClimaCalibrate as CAL
+
+function CAL.forward_model(iteration, member, config_dict = nothing)
+    experiment_dir = dirname(Base.active_project())
+    if isnothing(config_dict)
+        config_dict =
+            YAML.load_file(joinpath(experiment_dir, "model_config.yml"))
+    end
     output_dir = config_dict["output_dir"]
     member_path = path_to_ensemble_member(output_dir, iteration, member)
     config_dict["output_dir"] = member_path
@@ -55,21 +31,16 @@ function set_up_forward_model(member, iteration, config_dict::AbstractDict)
 
     # Turn off default diagnostics
     config_dict["output_default_diagnostics"] = false
-    return CA.AtmosConfig(config_dict)
-end
-
-"""
-    run_forward_model(atmos_config::CA.AtmosConfig)
-
-Run the atmosphere model with the given an AtmosConfig object.
-Currently only has basic error handling.
-"""
-function run_forward_model(atmos_config::CA.AtmosConfig)
+    atmos_config = CA.AtmosConfig(
+        config_dict;
+        comms_ctx = ClimaComms.SingletonCommsContext(),
+    )
     simulation = CA.get_simulation(atmos_config)
     sol_res = CA.solve_atmos!(simulation)
     if sol_res.ret_code == :simulation_crashed
         error(
             "The ClimaAtmos simulation has crashed. See the stack trace for details.",
         )
     end
+    return simulation
 end

calibration/test/Project.toml

@@ -5,11 +5,12 @@ ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 ClimaCalibrate = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
+ClusterManagers = "34f1f09b-3a8b-5176-ab39-66d58a4d544e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 
 [compat]
-ClimaCalibrate = "0.0.2 - 0.0.4"
+ClimaCalibrate = "0.0.6"