@@ -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}
0 commit comments