Unify logging for slurm pipeline (#68)

Project.toml

@@ -6,7 +6,6 @@ version = "0.1.0"
 [deps]
 CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
-ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
@@ -20,7 +19,7 @@ YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 [compat]
 CalibrateEmulateSample = "0.5"
 ClimaComms = "0.5"
-ClimaCore = "0.11, 0.12"
+ClimaParams = "0.10"
 Distributions = "0.25"
 EnsembleKalmanProcesses = "1"
 JLD2 = "0.4"

experiments/initialize.sbatch

@@ -6,12 +6,12 @@
 
 experiment_id=$1
 
-echo "Initializing calibration for experiment $experiment_id"
+echo "Initializing calibration for experiment: $experiment_id"
 julia --color=no --project=experiments/$experiment_id -e 'using Pkg; Pkg.instantiate(;verbose=true)'
 
 julia --color=no --project=experiments/$experiment_id -e '
     import CalibrateAtmos
     CalibrateAtmos.initialize("'$experiment_id'")
 '
 
-echo "Calibration initialized."
+echo "Calibration initialized"

experiments/pipeline.sh

@@ -7,16 +7,21 @@ module load climacommon/2024_02_27
 experiment_id=${1?Error: no experiment ID given}
 tasks_per_model_run=${2?Error: no tasks per model run given}
 
-echo "Running experiment $experiment_id with $tasks_per_model_run tasks per model run."
-echo 'Initializing ensemble for calibration.'
+# Get ensemble size, number of iterations, and output dir from EKP config file
+ensemble_size=$(grep "ensemble_size:" experiments/$experiment_id/ekp_config.yml | awk '{print $2}')
+n_iterations=$(grep "n_iterations:" experiments/$experiment_id/ekp_config.yml | awk '{print $2}')
+output=$(grep "output_dir:" experiments/$experiment_id/ekp_config.yml | awk '{print $2}')
+
+mkdir $output
+
+echo "Running experiment $experiment_id with $tasks_per_model_run tasks per model run"
 init_id=$(sbatch --parsable \
-                 --output=init_$experiment_id.out \
+                 --output=$output/log.out \
+                 --open-mode=append \
                  --partition=expansion \
                  experiments/initialize.sbatch $experiment_id)
-
-# Get ensemble size and the number of iterations from configuration file
-ensemble_size=$(grep "ensemble_size:" experiments/$experiment_id/ekp_config.yml | awk '{print $2}')
-n_iterations=$(grep "n_iterations:" experiments/$experiment_id/ekp_config.yml | awk '{print $2}')
+echo "Initialization job_id: $init_id"
+echo ""
 
 # Loop over iterations
 dependency="afterok:$init_id"
@@ -40,7 +45,8 @@ do
     update_id=$(
         sbatch --dependency=$dependency --kill-on-invalid-dep=yes --parsable \
                --job=update-$i \
-               --output=output/$experiment_id/$format_i/update_log.out \
+               --output=$output/log.out \
+               --open-mode=append \
                --partition=expansion \
                experiments/update.sbatch $experiment_id $i)
 

experiments/sphere_held_suarez_rhoe_equilmoist/model_interface.jl

@@ -69,7 +69,11 @@ Runs the atmosphere model with the given an AtmosConfig object.
 Currently only has basic error handling.
 """
 
-function run_forward_model(::ClimaAtmosModel, atmos_config::CA.AtmosConfig)
+function run_forward_model(
+    ::ClimaAtmosModel,
+    atmos_config::CA.AtmosConfig;
+    lk = nothing,
+)
     simulation = CA.get_simulation(atmos_config)
     sol_res = CA.solve_atmos!(simulation)
     if sol_res.ret_code == :simulation_crashed

experiments/update.sbatch

@@ -3,8 +3,7 @@
 
 experiment_id=$1
 i=$2
-
-echo "Iteration $i"
+echo "Running update step after iteration $i"
 
 julia --color=no --project=experiments/$experiment_id -e '
     import YAML, JLD2