Skip to content

Commit 10c4324

Browse files
Copilotrvhonorato
andauthored
Reintroduce configurable SLURM partition in general config and job script generation (#189)
* Initial plan * Add optional SLURM partition config and tests Agent-Logs-Url: https://github.com/haddocking/haddock-runner/sessions/2f8cea17-19ad-493c-bbee-5a926e30a1b8 Co-authored-by: rvhonorato <9445814+rvhonorato@users.noreply.github.com> * Refine SLURM header tests and directive ordering Agent-Logs-Url: https://github.com/haddocking/haddock-runner/sessions/2f8cea17-19ad-493c-bbee-5a926e30a1b8 Co-authored-by: rvhonorato <9445814+rvhonorato@users.noreply.github.com> * docs: use long partition in bm5 SLURM example Agent-Logs-Url: https://github.com/haddocking/haddock-runner/sessions/bdd522d9-220e-4270-a6d4-849c6215121d Co-authored-by: rvhonorato <9445814+rvhonorato@users.noreply.github.com> * test: add no-partition deserialize coverage and tighten slurm header checks Agent-Logs-Url: https://github.com/haddocking/haddock-runner/sessions/074643b0-b68f-48de-9079-96800ebf3c67 Co-authored-by: rvhonorato <9445814+rvhonorato@users.noreply.github.com> * Validate non-empty general.partition in input validation * Docs: align partition example value to long --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: rvhonorato <9445814+rvhonorato@users.noreply.github.com>
1 parent bc30bb5 commit 10c4324

5 files changed

Lines changed: 200 additions & 7 deletions

File tree

docs/src/reference.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ The `general` section contains global settings that apply to all scenarios and t
2929
| `max_concurrent` | integer | Yes | Maximum number of jobs to run simultaneously. Controls how many target-scenario combinations execute in parallel. |
3030
| `ncores` | integer | Yes | Number of CPU cores to allocate per job. |
3131
| `execution` | string | Yes | Execution backend. Valid values: `local`, `slurm`. |
32+
| `partition` | string | No | SLURM partition to submit jobs to when `execution: slurm`. If omitted, the cluster default partition is used. |
3233
| `mol_suffixes` | array of strings | Yes | File suffixes used to identify molecule files. Must contain at least 2 suffixes (typically receptor and ligand). |
3334
| `input_list` | string | Yes | Path to the input list file containing file paths for all targets. |
3435
| `work_dir` | string | Yes | Directory where benchmark results will be stored. Created automatically if it doesn't exist. |
@@ -198,6 +199,7 @@ general:
198199
max_concurrent: 4
199200
ncores: 4
200201
execution: slurm
202+
partition: long
201203
mol_suffixes: [_r_u, _l_u, _shape]
202204
input_list: shape/input.txt
203205
work_dir: shape-results

docs/src/setting-up-bm5.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,8 @@ For SLURM clusters, modify your config:
164164
```yaml
165165
general:
166166
execution: slurm
167-
cpus_per_task: 4
167+
ncores: 4
168+
partition: long
168169
```
169170
170171
## Step 6: Monitor and Manage the Benchmark

src/input.rs

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,12 @@ impl Input {
116116
anyhow::bail!("max_concurrent must be greater than 0");
117117
}
118118

119+
if let Some(partition) = &self.general.partition
120+
&& partition.trim().is_empty()
121+
{
122+
anyhow::bail!("partition must be non-empty when defined");
123+
}
124+
119125
match &self.general.execution {
120126
Execution::Local => {
121127
validate_haddock3()?;
@@ -147,6 +153,7 @@ pub struct General {
147153
pub max_concurrent: u16,
148154
pub ncores: u16,
149155
pub execution: Execution,
156+
pub partition: Option<String>,
150157
}
151158

152159
#[derive(Debug, Deserialize, Serialize, Clone)]
@@ -259,6 +266,7 @@ mod tests {
259266
max_concurrent: 1,
260267
ncores: 1,
261268
execution: Execution::Local,
269+
partition: None,
262270
},
263271
scenarios: vec![],
264272
};
@@ -279,6 +287,7 @@ mod tests {
279287
max_concurrent: 1,
280288
ncores: 1,
281289
execution: Execution::Local,
290+
partition: None,
282291
},
283292
scenarios: vec![],
284293
};
@@ -299,6 +308,7 @@ mod tests {
299308
max_concurrent: 1,
300309
ncores: 1,
301310
execution: Execution::Local,
311+
partition: None,
302312
},
303313
scenarios: vec![],
304314
};
@@ -319,6 +329,7 @@ mod tests {
319329
max_concurrent: 1,
320330
ncores: 1,
321331
execution: Execution::Local,
332+
partition: None,
322333
},
323334
scenarios: vec![],
324335
};
@@ -340,6 +351,7 @@ mod tests {
340351
max_concurrent: 1,
341352
ncores: 1,
342353
execution: Execution::Local,
354+
partition: None,
343355
},
344356
scenarios: vec![],
345357
};
@@ -360,6 +372,7 @@ mod tests {
360372
max_concurrent: 0,
361373
ncores: 1,
362374
execution: Execution::Local,
375+
partition: None,
363376
},
364377
scenarios: vec![],
365378
};
@@ -369,6 +382,29 @@ mod tests {
369382
assert!(result.is_err());
370383
}
371384

385+
#[test]
386+
fn test_input_validate_general_empty_partition() {
387+
let input = Input {
388+
general: General {
389+
mol_suffixes: vec!["_r".to_string(), "_l".to_string()],
390+
input_list: "test.txt".to_string(),
391+
work_dir: PathBuf::from("/tmp"),
392+
max_concurrent: 1,
393+
ncores: 1,
394+
execution: Execution::Slurm,
395+
partition: Some(" ".to_string()),
396+
},
397+
scenarios: vec![],
398+
};
399+
400+
let result = input.validate_general();
401+
assert!(result.is_err());
402+
assert_eq!(
403+
result.unwrap_err().to_string(),
404+
"partition must be non-empty when defined"
405+
);
406+
}
407+
372408
#[test]
373409
fn test_input_deserialize_unknown_top_level_field() {
374410
let yaml = r#"
@@ -413,6 +449,51 @@ scenarios:
413449
assert!(result.is_err());
414450
}
415451

452+
#[test]
453+
fn test_input_deserialize_with_partition() {
454+
let yaml = r#"
455+
general:
456+
mol_suffixes: ["_r", "_l"]
457+
input_list: input_list.txt
458+
work_dir: ./work
459+
max_concurrent: 1
460+
ncores: 1
461+
execution: slurm
462+
partition: gpu
463+
scenarios:
464+
- name: test
465+
workflow:
466+
topoaa:
467+
autohis: true
468+
"#;
469+
470+
let result: Result<Input, _> = serde_yaml::from_str(yaml);
471+
assert!(result.is_ok());
472+
assert_eq!(result.unwrap().general.partition, Some("gpu".to_string()));
473+
}
474+
475+
#[test]
476+
fn test_input_deserialize_without_partition() {
477+
let yaml = r#"
478+
general:
479+
mol_suffixes: ["_r", "_l"]
480+
input_list: input_list.txt
481+
work_dir: ./work
482+
max_concurrent: 1
483+
ncores: 1
484+
execution: slurm
485+
scenarios:
486+
- name: test
487+
workflow:
488+
topoaa:
489+
autohis: true
490+
"#;
491+
492+
let result: Result<Input, _> = serde_yaml::from_str(yaml);
493+
assert!(result.is_ok());
494+
assert_eq!(result.unwrap().general.partition, None);
495+
}
496+
416497
#[test]
417498
fn test_input_deserialize_unknown_scenario_field() {
418499
let yaml = r#"

src/job.rs

Lines changed: 110 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -372,12 +372,7 @@ impl Job {
372372
let absolute_wd = canonicalize(&self.wd).unwrap_or_else(|_| self.wd.to_path_buf());
373373

374374
// Write SLURM header
375-
let header = "#!/bin/bash\n".to_string()
376-
+ "#SBATCH --job-name=haddock\n"
377-
+ "#SBATCH --output=haddock-%j.out\n"
378-
+ "#SBATCH --error=haddock-%j.err\n"
379-
+ &format!("#SBATCH --cpus-per-task={}\n", self.general.ncores)
380-
+ "#SBATCH --ntasks=1\n";
375+
let header = self.generate_slurm_header();
381376

382377
// Write job body
383378
let body = format!(
@@ -393,6 +388,25 @@ impl Job {
393388
Ok(())
394389
}
395390

391+
fn generate_slurm_header(&self) -> String {
392+
let mut header = "#!/bin/bash\n".to_string()
393+
+ "#SBATCH --job-name=haddock\n"
394+
+ "#SBATCH --output=haddock-%j.out\n"
395+
+ "#SBATCH --error=haddock-%j.err\n";
396+
397+
header.push_str("#SBATCH --ntasks=1\n");
398+
header.push_str(&format!(
399+
"#SBATCH --cpus-per-task={}\n",
400+
self.general.ncores
401+
));
402+
403+
if let Some(partition) = &self.general.partition {
404+
header.push_str(&format!("#SBATCH --partition={partition}\n"));
405+
}
406+
407+
header
408+
}
409+
396410
/// Resolves _fname patterns to actual file paths
397411
fn resolve_fname_pattern(&self, pattern_str: &str, files: &[PathBuf]) -> Option<String> {
398412
let pattern = Regex::new(pattern_str).ok()?;
@@ -448,6 +462,7 @@ mod tests {
448462
max_concurrent: 1,
449463
ncores: 1,
450464
execution: Execution::Local,
465+
partition: None,
451466
};
452467

453468
let scenario = Scenario {
@@ -485,6 +500,7 @@ mod tests {
485500
max_concurrent: 1,
486501
ncores: 1,
487502
execution: Execution::Local,
503+
partition: None,
488504
},
489505
scenarios: vec![
490506
Scenario {
@@ -546,6 +562,7 @@ mod tests {
546562
max_concurrent: 1,
547563
ncores: 1,
548564
execution: Execution::Local,
565+
partition: None,
549566
};
550567

551568
let scenario = Scenario {
@@ -623,6 +640,7 @@ mod tests {
623640
max_concurrent: 1,
624641
ncores: 1,
625642
execution: Execution::Local,
643+
partition: None,
626644
},
627645
};
628646

@@ -682,6 +700,7 @@ mod tests {
682700
max_concurrent: 1,
683701
ncores: 1,
684702
execution: Execution::Local,
703+
partition: None,
685704
},
686705
};
687706

@@ -695,4 +714,89 @@ mod tests {
695714
assert!(all_files.contains(&misc_file));
696715
assert!(all_files.contains(&shape_file));
697716
}
717+
718+
#[test]
719+
fn test_generate_slurm_header_without_partition() {
720+
let job = Job {
721+
name: "test".to_string(),
722+
status: Status::Unknown,
723+
wd: PathBuf::from("/tmp"),
724+
target: Target {
725+
id: "target".to_string(),
726+
molecules: vec![],
727+
restraints: vec![],
728+
toppar: vec![],
729+
misc: vec![],
730+
shape: None,
731+
size: 0,
732+
},
733+
scenario: Scenario {
734+
name: "scenario".to_string(),
735+
workflow: Workflow {
736+
modules: IndexMap::new(),
737+
},
738+
},
739+
general: General {
740+
mol_suffixes: vec!["_r".to_string(), "_l".to_string()],
741+
input_list: "test.txt".to_string(),
742+
work_dir: PathBuf::from("/tmp"),
743+
max_concurrent: 1,
744+
ncores: 4,
745+
execution: Execution::Slurm,
746+
partition: None,
747+
},
748+
};
749+
750+
let header = job.generate_slurm_header();
751+
let expected = "#!/bin/bash\n\
752+
#SBATCH --job-name=haddock\n\
753+
#SBATCH --output=haddock-%j.out\n\
754+
#SBATCH --error=haddock-%j.err\n\
755+
#SBATCH --ntasks=1\n\
756+
#SBATCH --cpus-per-task=4\n";
757+
assert_eq!(header, expected);
758+
}
759+
760+
#[test]
761+
fn test_generate_slurm_header_with_partition() {
762+
let job = Job {
763+
name: "test".to_string(),
764+
status: Status::Unknown,
765+
wd: PathBuf::from("/tmp"),
766+
target: Target {
767+
id: "target".to_string(),
768+
molecules: vec![],
769+
restraints: vec![],
770+
toppar: vec![],
771+
misc: vec![],
772+
shape: None,
773+
size: 0,
774+
},
775+
scenario: Scenario {
776+
name: "scenario".to_string(),
777+
workflow: Workflow {
778+
modules: IndexMap::new(),
779+
},
780+
},
781+
general: General {
782+
mol_suffixes: vec!["_r".to_string(), "_l".to_string()],
783+
input_list: "test.txt".to_string(),
784+
work_dir: PathBuf::from("/tmp"),
785+
max_concurrent: 1,
786+
ncores: 4,
787+
execution: Execution::Slurm,
788+
partition: Some("gpu".to_string()),
789+
},
790+
};
791+
792+
let header = job.generate_slurm_header();
793+
let expected = "#!/bin/bash\n\
794+
#SBATCH --job-name=haddock\n\
795+
#SBATCH --output=haddock-%j.out\n\
796+
#SBATCH --error=haddock-%j.err\n\
797+
#SBATCH --ntasks=1\n\
798+
#SBATCH --cpus-per-task=4\n\
799+
#SBATCH --partition=gpu\n";
800+
assert_eq!(header, expected);
801+
}
698802
}

src/queue.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ mod tests {
205205
max_concurrent: 1,
206206
ncores: 1,
207207
execution: crate::input::Execution::Local,
208+
partition: None,
208209
},
209210
},
210211
Job {
@@ -233,6 +234,7 @@ mod tests {
233234
max_concurrent: 1,
234235
ncores: 1,
235236
execution: crate::input::Execution::Local,
237+
partition: None,
236238
},
237239
},
238240
];
@@ -286,6 +288,7 @@ mod tests {
286288
max_concurrent: 1,
287289
ncores: 1,
288290
execution: crate::input::Execution::Local,
291+
partition: None,
289292
},
290293
},
291294
Job {
@@ -314,6 +317,7 @@ mod tests {
314317
max_concurrent: 1,
315318
ncores: 1,
316319
execution: crate::input::Execution::Local,
320+
partition: None,
317321
},
318322
},
319323
Job {
@@ -342,6 +346,7 @@ mod tests {
342346
max_concurrent: 1,
343347
ncores: 1,
344348
execution: crate::input::Execution::Local,
349+
partition: None,
345350
},
346351
},
347352
];

0 commit comments

Comments
 (0)