From d0ab1adb82fb55ecee845cc1fc9b8170a4648400 Mon Sep 17 00:00:00 2001 From: "Rodrigo V. Honorato" Date: Thu, 23 Jul 2026 13:55:17 +0200 Subject: [PATCH 1/5] add sbatch headers --- src/input.rs | 149 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 149 insertions(+) diff --git a/src/input.rs b/src/input.rs index 548e847..8bce674 100644 --- a/src/input.rs +++ b/src/input.rs @@ -191,6 +191,8 @@ impl Input { anyhow::bail!("partition must be non-empty when defined"); } + self.validate_slurm_header()?; + match &self.general.execution { Execution::Local => { validate_haddock3()?; @@ -211,6 +213,31 @@ impl Input { Ok(()) } + + /// Validates `general.slurm_header` entries: keys must be recognized `sbatch` + /// long-option names (see `SBATCH_FIELDS`), and values must be scalars. + fn validate_slurm_header(&self) -> Result<()> { + let Some(slurm_header) = &self.general.slurm_header else { + return Ok(()); + }; + + for (key, value) in slurm_header { + if !SBATCH_FIELDS.contains(&key.as_str()) { + bail!( + "unknown slurm_header field '{key}': not a recognized sbatch option \ + (see docs/reference.md or https://slurm.schedmd.com/sbatch.html for valid fields)" + ); + } + + if matches!(value, Value::Sequence(_) | Value::Mapping(_)) { + bail!( + "slurm_header field '{key}' must be a string, number, or boolean, not a list/mapping" + ); + } + } + + Ok(()) + } } #[derive(Debug, Deserialize, Serialize, Clone)] @@ -227,8 +254,122 @@ pub struct General { pub preprocess: Option, pub postprocess: Option, pub gen_archive: Option, + // SBATCH optional header customization + pub slurm_header: Option>, } +/// Canonical `sbatch` long-option names accepted in `general.slurm_header`. +/// Sourced from https://slurm.schedmd.com/sbatch.html. +/// +/// Note that some options that don't make sense to be sense embedded +/// in a job script header are intentionally excluded. +pub(crate) const SBATCH_FIELDS: &[&str] = &[ + "account", + "acctg-freq", + "array", + "batch", + "bb", + "bbf", + "begin", + "chdir", + "cluster-constraint", + "clusters", + "comment", + "consolidate-segments", + "constraint", + "container", + "container-id", + "container-type", + "contiguous", + "core-spec", + "cores-per-socket", + "cpu-freq", + "cpus-per-gpu", + "cpus-per-task", + "deadline", + "delay-boot", + "dependency", + "distribution", + "error", + "exclude", + "exclusive", + "export", + "export-file", + "extra", + "extra-node-info", + "get-user-env", + "gid", + "gpu-bind", + "gpu-freq", + "gpus", + "gpus-per-node", + "gpus-per-socket", + "gpus-per-task", + "gres", + "gres-flags", + "hint", + "hold", + "ignore-pbs", + "input", + "job-name", + "kill-on-invalid-dep", + "licenses", + "mail-type", + "mail-user", + "mcs-label", + "mem", + "mem-bind", + "mem-per-cpu", + "mem-per-gpu", + "mem-update", + "mincpus", + "network", + "nice", + "no-kill", + "no-requeue", + "nodefile", + "nodelist", + "nodes", + "ntasks", + "ntasks-per-core", + "ntasks-per-gpu", + "ntasks-per-node", + "ntasks-per-socket", + "oom-kill-step", + "open-mode", + "output", + "overcommit", + "oversubscribe", + "parsable", + "partition", + "prefer", + "priority", + "profile", + "propagate", + "qos", + "reboot", + "requeue", + "reservation", + "resources", + "resv-ports", + "segment", + "signal", + "sockets-per-node", + "spread-job", + "spread-segments", + "stepmgr", + "switches", + "thread-spec", + "threads-per-core", + "time", + "time-min", + "tmp", + "tres-bind", + "uid", + "wait", + "wckey", +]; + #[derive(Debug, Deserialize, Serialize, Clone)] #[serde(rename_all = "lowercase")] pub enum Execution { @@ -343,6 +484,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -367,6 +509,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -391,6 +534,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -415,6 +559,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -440,6 +585,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -464,6 +610,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -487,6 +634,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![], }; @@ -713,6 +861,7 @@ scenarios: preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![Scenario { name: "test".to_string(), From 87cd623e1cab1f41f12d95678651ffd309864d90 Mon Sep 17 00:00:00 2001 From: "Rodrigo V. Honorato" Date: Thu, 23 Jul 2026 13:55:36 +0200 Subject: [PATCH 2/5] add tests --- src/input.rs | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/src/input.rs b/src/input.rs index 8bce674..e55749d 100644 --- a/src/input.rs +++ b/src/input.rs @@ -647,6 +647,85 @@ mod tests { ); } + fn make_input_with_slurm_header(slurm_header: IndexMap) -> Input { + Input { + general: General { + mol_suffixes: vec!["_r".to_string(), "_l".to_string()], + input_list: "test.txt".to_string(), + work_dir: PathBuf::from("/tmp"), + max_concurrent: 1, + ncores: 1, + execution: Execution::Local, + partition: None, + preprocess: None, + postprocess: None, + gen_archive: None, + slurm_header: Some(slurm_header), + }, + scenarios: vec![], + } + } + + #[test] + fn test_validate_slurm_header_unknown_field_is_error() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("partiton".to_string(), Value::String("gpu".to_string())); + let input = make_input_with_slurm_header(slurm_header); + + let result = input.validate_general(); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("partiton"), + "error should mention the offending key" + ); + } + + #[test] + fn test_validate_slurm_header_known_field_is_ok() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("nodes".to_string(), Value::Number(2.into())); + slurm_header.insert("account".to_string(), Value::String("proj".to_string())); + let input = make_input_with_slurm_header(slurm_header); + + assert!(input.validate_slurm_header().is_ok()); + } + + #[test] + fn test_validate_slurm_header_sequence_value_is_error() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert( + "nodes".to_string(), + Value::Sequence(vec![Value::Number(1.into())]), + ); + let input = make_input_with_slurm_header(slurm_header); + + let result = input.validate_general(); + assert!(result.is_err()); + } + + #[test] + fn test_validate_slurm_header_mapping_value_is_error() { + let mut slurm_header = IndexMap::new(); + let mut mapping = serde_yaml::Mapping::new(); + mapping.insert(Value::String("a".to_string()), Value::Bool(true)); + slurm_header.insert("nodes".to_string(), Value::Mapping(mapping)); + let input = make_input_with_slurm_header(slurm_header); + + let result = input.validate_general(); + assert!(result.is_err()); + } + + #[test] + fn test_validate_slurm_header_null_and_empty_values_are_ok() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("partition".to_string(), Value::Null); + slurm_header.insert("qos".to_string(), Value::String(String::new())); + slurm_header.insert("exclusive".to_string(), Value::Bool(false)); + let input = make_input_with_slurm_header(slurm_header); + + assert!(input.validate_slurm_header().is_ok()); + } + #[test] fn test_input_deserialize_unknown_top_level_field() { let yaml = r#" From 529ffc232ee72c6a0e5da9a233c5d89605fbeb5f Mon Sep 17 00:00:00 2001 From: "Rodrigo V. Honorato" Date: Thu, 23 Jul 2026 14:10:22 +0200 Subject: [PATCH 3/5] update generate_slurm_header --- src/job.rs | 215 +++++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 194 insertions(+), 21 deletions(-) diff --git a/src/job.rs b/src/job.rs index fa12cab..2380639 100644 --- a/src/job.rs +++ b/src/job.rs @@ -7,7 +7,9 @@ use std::fs; use std::io::Write; use std::path::PathBuf; +use indexmap::IndexMap; use itertools::Itertools; +use serde_yaml::Value; use crate::input::{DIGIT_SUFFIX_RE, Execution, General}; use crate::runner::local; @@ -422,19 +424,66 @@ impl Job { } fn generate_slurm_header(&self) -> String { - let mut header = "#!/bin/bash\n".to_string() - + "#SBATCH --job-name=haddock\n" - + "#SBATCH --output=haddock-%j.out\n" - + "#SBATCH --error=haddock-%j.err\n"; + // `Some(value)` renders as `#SBATCH --key=value`, `None` as a bare `#SBATCH --key` + let mut fields: IndexMap> = IndexMap::new(); - header.push_str("#SBATCH --ntasks=1\n"); - header.push_str(&format!( - "#SBATCH --cpus-per-task={}\n", - self.general.ncores - )); + // cpus-per-task is always derived from ncores and cannot be overridden. + fields.insert( + "cpus-per-task".to_string(), + Some(self.general.ncores.to_string()), + ); if let Some(partition) = &self.general.partition { - header.push_str(&format!("#SBATCH --partition={partition}\n")); + fields.insert("partition".to_string(), Some(partition.clone())); + } + + // Parse fields and assign its proper types + if let Some(slurm_header) = &self.general.slurm_header { + for (key, value) in slurm_header { + // NOTE: `cpus-per-task` is a special field, its tied to haddock's `ncores` + // so here on purpose we do not let users override it + if key == "cpus-per-task" { + log::warn!( + "slurm_header field 'cpus-per-task' is ignored: it is derived from ncores ({}) and cannot be overridden", + self.general.ncores + ); + continue; + } + + // Match values + match value { + // Not provided / explicitly unset: drop the entry and let SLURM default apply. + Value::Null | Value::Bool(false) => { + fields.shift_remove(key); + } + // Empty + Value::String(s) if s.trim().is_empty() => { + fields.shift_remove(key); + } + // Bare no-argument flag, e.g. `exclusive: true` -> `#SBATCH --exclusive`. + Value::Bool(true) => { + fields.insert(key.clone(), None); + } + Value::String(s) => { + fields.insert(key.clone(), Some(s.clone())); + } + Value::Number(n) => { + fields.insert(key.clone(), Some(n.to_string())); + } + // Sequences/mappings are rejected by `Input::validate`; ignore defensively. + Value::Sequence(_) | Value::Mapping(_) | Value::Tagged(_) => {} + } + } + } + + let mut header = "#!/bin/bash\n".to_string(); + for (key, value) in &fields { + match value { + // key=value + Some(v) => header.push_str(&format!("#SBATCH --{key}={v}\n")), + // bare + None => header.push_str(&format!("#SBATCH --{key}\n")), + } } header @@ -500,6 +549,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }; let scenario = Scenario { @@ -541,6 +591,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, scenarios: vec![ Scenario { @@ -606,6 +657,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }; let scenario = Scenario { @@ -687,6 +739,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }; @@ -750,6 +803,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }; @@ -796,17 +850,12 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }; let header = job.generate_slurm_header(); - let expected = "#!/bin/bash\n\ -#SBATCH --job-name=haddock\n\ -#SBATCH --output=haddock-%j.out\n\ -#SBATCH --error=haddock-%j.err\n\ -#SBATCH --ntasks=1\n\ -#SBATCH --cpus-per-task=4\n"; - assert_eq!(header, expected); + assert_eq!(header, "#!/bin/bash\n#SBATCH --cpus-per-task=4\n"); } #[test] @@ -841,20 +890,143 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }; + let header = job.generate_slurm_header(); + assert_eq!( + header, + "#!/bin/bash\n#SBATCH --cpus-per-task=4\n#SBATCH --partition=gpu\n" + ); + } + + fn make_slurm_job( + partition: Option, + slurm_header: Option>, + ) -> Job { + Job { + name: "test".to_string(), + status: Status::Unknown, + wd: PathBuf::from("/tmp"), + target: Target { + id: "target".to_string(), + molecules: vec![], + restraints: vec![], + toppar: vec![], + misc: vec![], + shape: None, + size: 0, + }, + scenario: Scenario { + name: "scenario".to_string(), + workflow: Workflow { + modules: IndexMap::new(), + }, + }, + general: General { + mol_suffixes: vec!["_r".to_string(), "_l".to_string()], + input_list: "test.txt".to_string(), + work_dir: PathBuf::from("/tmp"), + max_concurrent: 1, + ncores: 4, + execution: Execution::Slurm, + partition, + preprocess: None, + postprocess: None, + gen_archive: None, + slurm_header, + }, + } + } + + #[test] + fn test_generate_slurm_header_with_new_field() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("nodes".to_string(), Value::Number(1.into())); + slurm_header.insert( + "account".to_string(), + Value::String("project_XXXXXX".to_string()), + ); + + let job = make_slurm_job(None, Some(slurm_header)); + let header = job.generate_slurm_header(); + let expected = "#!/bin/bash\n\ +#SBATCH --cpus-per-task=4\n\ +#SBATCH --nodes=1\n\ +#SBATCH --account=project_XXXXXX\n"; + assert_eq!(header, expected); + } + + #[test] + fn test_generate_slurm_header_cpus_per_task_cannot_be_overridden() { + // cpus-per-task is always derived from general.ncores; a user-supplied + // value in slurm_header is ignored (a warning is logged instead). + let mut slurm_header = IndexMap::new(); + slurm_header.insert("cpus-per-task".to_string(), Value::Number(8.into())); + + let job = make_slurm_job(None, Some(slurm_header)); + let header = job.generate_slurm_header(); + assert_eq!(header, "#!/bin/bash\n#SBATCH --cpus-per-task=4\n"); + } + + #[test] + fn test_generate_slurm_header_slurm_header_overrides_partition_field() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("partition".to_string(), Value::String("gpu".to_string())); + + let job = make_slurm_job(Some("small".to_string()), Some(slurm_header)); let header = job.generate_slurm_header(); let expected = "#!/bin/bash\n\ -#SBATCH --job-name=haddock\n\ -#SBATCH --output=haddock-%j.out\n\ -#SBATCH --error=haddock-%j.err\n\ -#SBATCH --ntasks=1\n\ #SBATCH --cpus-per-task=4\n\ #SBATCH --partition=gpu\n"; assert_eq!(header, expected); } + #[test] + fn test_generate_slurm_header_null_value_is_skipped() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("partition".to_string(), Value::Null); + + let job = make_slurm_job(Some("small".to_string()), Some(slurm_header)); + let header = job.generate_slurm_header(); + assert_eq!(header, "#!/bin/bash\n#SBATCH --cpus-per-task=4\n"); + } + + #[test] + fn test_generate_slurm_header_false_value_is_skipped() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("exclusive".to_string(), Value::Bool(false)); + + let job = make_slurm_job(None, Some(slurm_header)); + let header = job.generate_slurm_header(); + assert!(!header.contains("exclusive")); + } + + #[test] + fn test_generate_slurm_header_empty_string_value_is_skipped() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("qos".to_string(), Value::String(" ".to_string())); + + let job = make_slurm_job(None, Some(slurm_header)); + let header = job.generate_slurm_header(); + assert!(!header.contains("qos")); + } + + #[test] + fn test_generate_slurm_header_true_value_renders_bare_flag() { + let mut slurm_header = IndexMap::new(); + slurm_header.insert("exclusive".to_string(), Value::Bool(true)); + + let job = make_slurm_job(None, Some(slurm_header)); + let header = job.generate_slurm_header(); + assert!( + header.contains("#SBATCH --exclusive\n"), + "expected bare --exclusive flag but got:\n{header}" + ); + assert!(!header.contains("--exclusive=")); + } + fn make_job_with_modules(modules: IndexMap) -> Job { Job { name: "test".to_string(), @@ -884,6 +1056,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, } } From 47082e5b3792b7bcfaf3f0af9f415c4f1f884f8b Mon Sep 17 00:00:00 2001 From: "Rodrigo V. Honorato" Date: Thu, 23 Jul 2026 14:10:40 +0200 Subject: [PATCH 4/5] update test --- src/queue.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/queue.rs b/src/queue.rs index 89983bc..1edaf00 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -236,6 +236,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }, Job { @@ -268,6 +269,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }, ]; @@ -325,6 +327,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }, Job { @@ -357,6 +360,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }, Job { @@ -389,6 +393,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }, }, ]; @@ -429,6 +434,7 @@ mod tests { preprocess: None, postprocess: None, gen_archive: None, + slurm_header: None, }; let job = Job { From d2befc03711c9b760e880cd9aad77a1fe5a9eab6 Mon Sep 17 00:00:00 2001 From: "Rodrigo V. Honorato" Date: Thu, 23 Jul 2026 14:22:07 +0200 Subject: [PATCH 5/5] add `slurm.yml` --- example/slurm.yml | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 example/slurm.yml diff --git a/example/slurm.yml b/example/slurm.yml new file mode 100644 index 0000000..fbcce78 --- /dev/null +++ b/example/slurm.yml @@ -0,0 +1,40 @@ +general: + max_concurrent: 4 # How many should be executed at the same time + ncores: 96 + execution: slurm + # arbitrary #SBATCH directives, added on top of/overriding the above + # if not defined, SLURM will use it's default, keep in mind + # this will change by system + slurm_header: + partition: small + nodes: 1 + account: project_XXXXXX + mem-per-cpu: 1500 + time: "24:00:00" + qos: standard + exclusive: true # renders as `#SBATCH --exclusive` + mol_suffixes: [_r_u, _l_u, _x_u] + input_list: docking/input_list.txt + work_dir: ../bm-goes-here + +scenarios: + - name: true-interface + workflow: + topoaa: + autohis: true + rigidbody: + sampling: 5 + ambig_fname: _ti.tbl + unambig_fname: _unambig.tbl + ligand_top_fname: _ligand.top + ligand_param_fname: _ligand.param + seletop: + select: 2 + flexref: + ambig_fname: _ti.tbl + unambig_fname: _unambig.tbl + ligand_top_fname: _ligand.top + ligand_param_fname: _ligand.param + emref: + caprieval: + reference_fname: _ref.pdb