Support for main args for slurm

benchmarks/src/helpers/main_config_helper.jl

@@ -4,7 +4,7 @@ using DrWatson
 using TOML
 
 export mainsim_config_path, load_main_config, listof_main_physics,
-  collect_simsfor_physics, has_config_args, config_args
+  collect_simsfor_physics, has_config_args, get_config_args, get_config_arg
 
 mainsim_config_path() = srcdir("main_config.toml")
 
@@ -66,9 +66,22 @@ function has_config_args(config_info, snd::SimNameData)
   return !isempty(tag_info)
 end
 
-function config_args(config_info, snd::SimNameData)
+function get_config_args(config_info, snd::SimNameData)
   if !(has_config_args(config_info, snd))
     error("Arguments for $(snd) were not found in the main configuration or provided architecture $(snd.arch) is invalid")
   end
   return config_info[snd.physics][snd.arch][snd.tag]
 end
+
+function get_config_arg(config_info, snd::SimNameData, arg)
+  if !has_config_args(config_info, snd)
+    return nothing
+  end
+
+  args = get_config_args(config_info, snd)
+  if !haskey(args, arg)
+    return nothing
+  end
+
+  return args[arg]
+end

benchmarks/src/helpers/main_source.jl

@@ -6,6 +6,10 @@ using MLStyle
 
 include(helpersdir("physics_config_helper.jl"))
 
+const DEFAULT_CPU_SLURMARGS = `--partition=hpg-milan`
+const DEFAULT_CUDA_SLURMARGS = `--partition=gpu --gres=gpu:a100:1`
+const DEFAULT_MAXJOBS = 6
+
 function generate_all_configs()
   for physics in listof_main_physics()
       generate_configsfor_physics(physics)
@@ -66,25 +70,41 @@ function run_single_physics(physics, physics_configs)
 end
 
 function run_single_config!(dependency_ids, sim_namedata::SimNameData)
-  config_data = TOML.parsefile(simconfig_path(sim_namedata))
 
-  concurrent_jobs = 6
+  main_config_info = load_main_config()
+  concurrent_jobs = get_concur_jobs(main_config_info, sim_namedata)
+  slurm_args = get_slurm_args(main_config_info, sim_namedata)
 
   physics = sim_namedata.physics
   arch = sim_namedata.arch
   tag = sim_namedata.tag
 
-  count = simconfig_size(config_data)
+  count = simconfig_size(load_simconfig(sim_namedata))
 
   @info "Running $count $arch tasks"
 
-  args = @match arch begin
-    "cpu" => `--partition=hpg-milan`
-    "cuda" => `--partition=gpu --gres=gpu:a100:1`
-  end
-
-  jobid = readchomp(`sbatch --array=1-$(count)%$(concurrent_jobs) --parsable $(args) $(scriptsdir("array.sh")) $physics $arch $tag`)
+  jobid = readchomp(`sbatch --array=1-$(count)%$(concurrent_jobs) --parsable $(slurm_args) $(scriptsdir("array.sh")) $physics $arch $tag`)
 
   @info "Job ID for $(tagfor_run(sim_namedata)) is: $jobid"
   push!(dependency_ids, jobid)
+  return dependency_ids
+end
+
+function get_slurm_args(main_config_info, snd::SimNameData)
+
+  slurm_args = get_config_arg(main_config_info, snd, "slurm_args")
+
+  if isnothing(slurm_args)
+    return @match snd.arch begin
+      "cpu" => DEFAULT_CPU_SLURMARGS
+      "cuda" => DEFAULT_CUDA_SLURMARGS
+    end
+  end
+
+  return Cmd(slurm_args)
+end
+
+function get_concur_jobs(main_config_info, snd::SimNameData)
+  job_arg = get_config_arg(main_config_info, snd, "concur_jobs")
+  return isnothing(job_arg) ? DEFAULT_MAXJOBS : job_arg
 end

benchmarks/src/helpers/physics_config_helper.jl

@@ -25,7 +25,7 @@ function load_save_benchmark_data(sim)
 end
 
 has_physicsconfig_args(sim) = has_config_args(load_physicsconfig(sim), sim)
-physicsconfig_args(sim) = config_args(load_physicsconfig(sim), sim)
+physicsconfig_args(sim) = get_config_args(load_physicsconfig(sim), sim)
 
 function process_simulation_config(entry)
   received_params = dict_list(entry)

benchmarks/test/main_config.jl

@@ -84,7 +84,7 @@ end
 
   hcpa = SimNameData("heat", "cpu", "a")
   @test has_config_args(main_config_info, hcpa)
-  hcpa_args = config_args(main_config_info, hcpa)
+  hcpa_args = get_config_args(main_config_info, hcpa)
   @test hcpa_args["args"] == "test"
 
   @test is_supported_arch("cpu")
@@ -94,28 +94,28 @@ end
 
   no_entries = SimNameData("heat", "cuda", "c")
   @test has_config_args(main_config_info, no_entries) == false
-  @test_throws "not found in the main config" config_args(main_config_info, no_entries)
+  @test_throws "not found in the main config" get_config_args(main_config_info, no_entries)
 
   no_tag = SimNameData("heat", "cpu", "badtag")
   @test has_config_args(main_config_info, no_tag) == false
-  @test_throws "not found in the main config" config_args(main_config_info, no_tag)
+  @test_throws "not found in the main config" get_config_args(main_config_info, no_tag)
 
   no_arch = SimNameData("not_heat", "cpu", "c")
   @test has_config_args(main_config_info, no_arch) == false
-  @test_throws "not found in the main config" config_args(main_config_info, no_arch)
+  @test_throws "not found in the main config" get_config_args(main_config_info, no_arch)
 
   bad_arch = SimNameData("not_heat", "fake", "c")
   @test has_config_args(main_config_info, bad_arch) == false
-  @test_throws "not found in the main config" config_args(main_config_info, bad_arch)
+  @test_throws "not found in the main config" get_config_args(main_config_info, bad_arch)
 
   no_physics = SimNameData("brussel", "cuda", "c")
   @test has_config_args(main_config_info, no_physics) == false
-  @test_throws "not found in the main config" config_args(main_config_info, no_physics)
+  @test_throws "not found in the main config" get_config_args(main_config_info, no_physics)
 
   empty_toml = TOML.parse("")
   empty_snd = SimNameData("", "", "")
   @test has_config_args(empty_toml, empty_snd) == false
-  @test_throws "not found in the main configuration" config_args(empty_toml, empty_snd)
+  @test_throws "not found in the main configuration" get_config_args(empty_toml, empty_snd)
 end
 
 end