Comments and multithreading

.github/workflows/ci.yml

@@ -37,10 +37,10 @@ jobs:
           arch: ${{ matrix.arch }}
       - uses: julia-actions/julia-buildpkg@v1
       - name: Surface Fluxes Perfect Model Test
-        run:
+        run: |
           julia --project=experiments/surface_fluxes_perfect_model -e 'using Pkg; Pkg.instantiate(;verbose=true)'
       - run: |
-          julia --project=experiments/surface_fluxes_perfect_model test/test_sf.jl
+          julia -t 10 --project=experiments/surface_fluxes_perfect_model test/test_sf.jl
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v3

experiments/pipeline.jl

@@ -4,7 +4,7 @@ using Pkg
 experiment_id = ARGS[1]
 Pkg.activate("experiments/$experiment_id")
 Pkg.instantiate()
-import CalibrateAtmos
+using CalibrateAtmos
 pkg_dir = pkgdir(CalibrateAtmos)
 experiment_path = "$pkg_dir/experiments/$experiment_id"
 include("$experiment_path/model_interface.jl")

experiments/surface_fluxes_perfect_model/model_interface.jl

@@ -84,8 +84,18 @@ end
 Runs the model with the given an AbstractDict object.
 """
 
-function run_forward_model(::SurfaceFluxModel, config::AbstractDict)
-    x_inputs = load_profiles(config["x_data_file"])
+function run_forward_model(
+    ::SurfaceFluxModel,
+    config::AbstractDict;
+    lk = nothing,
+)
+    x_inputs = if isnothing(lk)
+        load_profiles(config["x_data_file"])
+    else
+        lock(lk) do
+            load_profiles(config["x_data_file"])
+        end
+    end
     FT = typeof(x_inputs.profiles_int[1].T)
     obtain_ustar(FT, x_inputs, config)
 end

src/ekp_interface.jl

@@ -5,6 +5,7 @@ using Distributions
 import EnsembleKalmanProcesses as EKP
 using EnsembleKalmanProcesses.ParameterDistributions
 using EnsembleKalmanProcesses.TOMLInterface
+import ClimaComms
 
 """
     path_to_iteration(output_dir, iteration)
@@ -152,7 +153,7 @@ include(joinpath(experiment_path, "generate_truth.jl"))
 eki = CalibrateAtmos.calibrate(experiment_id)
 ```
 """
-function calibrate(experiment_id)
+function calibrate(experiment_id; device = ClimaComms.device())
     ekp_config =
         YAML.load_file(joinpath("experiments", experiment_id, "ekp_config.yml"))
     # initialize the CalibrateAtmos
@@ -165,11 +166,11 @@ function calibrate(experiment_id)
     for i in 0:(N_iter - 1)
         # run G model to produce output from N_mem ensemble members
         physical_model = get_forward_model(Val(Symbol(experiment_id)))
-
-        for m in 1:N_mem # TODO: parallelize with pmap!
+        lk = ReentrantLock()
+        ClimaComms.@threaded device for m in 1:N_mem
             # model run for each ensemble member
             model_config = get_config(physical_model, m, i, experiment_id)
-            run_forward_model(physical_model, model_config)
+            run_forward_model(physical_model, model_config; lk)
             @info "Finished model run for member $m at iteration $i"
         end
 

test/ekp.jl → test/test_emulate_sample.jl

@@ -10,13 +10,6 @@ using EnsembleKalmanProcesses.TOMLInterface
 
 import CalibrateAtmos as CAL
 
-@testset "Test `calibrate`" begin
-    dir = pwd()
-    cd(pkgdir(CAL))
-    @test_throws ErrorException CAL.calibrate("surface_fluxes_perfect_model")
-    cd(dir)
-end
-
 @testset "Emulate and Sample tests" begin
     y_obs = [261.5493]
     y_noise_cov = [0.02619;;]

test/test_model_interface.jl

@@ -29,3 +29,11 @@ end
         "observation_map not implemented for experiment Val{:test}() at iteration 1",
     ) CalibrateAtmos.observation_map(Val(:test), 1)
 end
+
+# This test depends on `surface_fluxes_perfect_model` in the `experiments/` folder 
+@testset "Test `calibrate`" begin
+    dir = pwd()
+    cd(pkgdir(CalibrateAtmos))
+    @test_throws ErrorException CAL.calibrate("surface_fluxes_perfect_model")
+    cd(dir)
+end