clean up!

experiments/initialize.sbatch

@@ -6,10 +6,7 @@
 
 job_id=$1
 
-echo 'Instantiating the Julia project.'
 julia --project=experiments -e 'using Pkg; Pkg.instantiate(;verbose=true)'
-
-echo 'Initializing ensemble for calibration.'
 julia --project=experiments -e "
 import CalibrateAtmos
 using Distributions

experiments/job_array.sh

@@ -1,45 +1,44 @@
 #!/bin/bash
 
-# Set number of tasks per model run
-tasks_per_ensemble=8
-
 # Configure for cluster
 module load julia/1.9.3 cuda/12.2 ucx/1.14.1_cuda-12.2 openmpi/4.1.5_cuda-12.2 hdf5/1.12.2-ompi415
 export OMPI_MCA_opal_warn_on_missing_libcuda=0
 export JULIA_NUM_PRECOMPILE_TASKS=8
 export JULIA_CPU_TARGET='broadwell;skylake'
 
-# Get job ID from command line
+# Parse command line
 job_id=${1?Error: no job ID given}
+tasks_per_model_run=${2?Error: no tasks per model run given}
 
-echo "Running job: $job_id"
-
-init_id=$(sbatch --parsable experiments/initialize.sbatch $job_id)
+echo "Running job $job_id with $tasks_per_model_run tasks per model run."
+echo 'Initializing ensemble for calibration.'
+init_id=$(sbatch --parsable --output=/dev/null experiments/initialize.sbatch $job_id)
 
 # Get ensemble size and the number of iterations from configuration file
 ensemble_size=$(grep "ensemble_size:" experiments/$job_id/ekp_config.yml | awk '{print $2}')
 n_iterations=$(grep "n_iterations:" experiments/$job_id/ekp_config.yml | awk '{print $2}')
 
 # Loop over iterations
-dependency="after:$init_id"
+dependency="afterok:$init_id"
 for i in $(seq 0 $((n_iterations - 1)))
 do
     echo "Scheduling iteration $i"
     format_i=$(printf "iteration_%03d" "$i")
 
     ensemble_array_id=$(
-        sbatch --dependency=$dependency \
+        sbatch --dependency=$dependency --kill-on-invalid-dep=yes \
                --parsable \
+               --output=/dev/null \
                --job=model-$i \
-               --ntasks=$tasks_per_ensemble \
+               --ntasks=$tasks_per_model_run \
                --array=1-$ensemble_size \
                experiments/model_run.sbatch $job_id $i)
 
     dependency=afterany:$ensemble_array_id
     echo "Iteration $i job id: $ensemble_array_id"
 
     update_id=$(
-        sbatch --dependency=$dependency \
+        sbatch --dependency=$dependency --kill-on-invalid-dep=yes \
                --job=update-$i \
                --output=output/$job_id/$format_i/update_log.out \
                --parsable \

experiments/model_run.sbatch

@@ -13,7 +13,7 @@ member=$(printf "member_%03d" "$SLURM_ARRAY_TASK_ID")
 output=output/$job_id/$format_i/$member/model_log.out
 
 # Run the forward model
-srun --output=$output --open-mode=append julia --project=experiments -e "
+srun --output=$output julia --project=experiments -e "
     using ClimaComms
     ClimaComms.init(ClimaComms.context())
     import CalibrateAtmos