diff --git a/ats/atsMachines/slurmProcessorScheduled.py b/ats/atsMachines/slurmProcessorScheduled.py index c73c62d..465f5d5 100644 --- a/ats/atsMachines/slurmProcessorScheduled.py +++ b/ats/atsMachines/slurmProcessorScheduled.py @@ -157,6 +157,7 @@ def examineOptions(self, options): print("%s options.sequential = %s " % (DEBUG_SLURM, options.sequential)) print("%s options.nosrun = %s " % (DEBUG_SLURM, options.nosrun)) print("%s options.salloc = %s " % (DEBUG_SLURM, options.salloc)) + print("%s options.sargs = %s " % (DEBUG_SLURM, options.sargs)) print("%s options.checkForAtsProc = %s " % (DEBUG_SLURM, options.checkForAtsProc)) print("%s options.showGroupStartOnly = %s " % (DEBUG_SLURM, options.showGroupStartOnly)) print("%s options.removeStartNote = %s " % (DEBUG_SLURM, options.removeStartNote)) @@ -203,6 +204,7 @@ def examineOptions(self, options): self.exclusive = options.exclusive self.mpibind = options.mpibind self.salloc = options.salloc + self.sargs = options.sargs self.toss_nn = options.toss_nn self.strict_nn = options.strict_nn self.useMinNodes = options.useMinNodes @@ -229,6 +231,10 @@ def addOptions(self, parser): default = add_partition, help = "Partition in which to run jobs with np > 0") + parser.add_option("--sargs", action="store", type="string", dest='sargs', + default = "", + help = "extra srun/salloc args") + parser.add_option("--numNodes", action="store", type="int", dest='numNodes', default = 2, help="Number of nodes to use") @@ -357,9 +363,15 @@ def calculateCommandList(self, test): # Set the --partition option here # If not within an salloc, then set the partition to be used for the srun line - if self.runningWithinSalloc == False: + if self.runningWithinSalloc == False and self.partition: srun_partition="--partition=%s" % self.partition + if self.sargs: + extra_sargs = self.sargs + else: + extra_sargs = "--comment=nosargs" + + # Set the --nodes srun option here # # print("DEBUG SAD JULY num_nodes=%i minNodes=%i self.numNodes=%i " % (num_nodes, minNodes, self.numNodes)) @@ -479,11 +491,11 @@ def calculateCommandList(self, test): # End of Coding suggested by Chris Scroeder if self.salloc : - return ["salloc", srun_partition, srun_ex_or_sh, srun_nodes] + commandList + return ["salloc", srun_partition, srun_ex_or_sh, srun_nodes, extra_sargs] + commandList else: return ["srun", srun_mpi_type, "--label", "-J", test.jobname, srun_partition, srun_ex_or_sh, srun_unbuffered, srun_mpibind, srun_distribution, srun_nodes, srun_cpus_per_task, - "--ntasks=%i" % np \ + "--ntasks=%i" % np, extra_sargs \ ] + commandList # ---------------------------------------------------------------------------------------------------------------------------- @@ -516,11 +528,11 @@ def calculateCommandList(self, test): print("SAD DEBUG SRUN800 ") if self.salloc : - return ["salloc", srun_partition, srun_ex_or_sh, srun_nodes] + commandList + return ["salloc", srun_partition, srun_ex_or_sh, srun_nodes, extra_sargs] + commandList else: return ["srun", srun_mpi_type, "--label", "-J", test.jobname, srun_partition, srun_ex_or_sh, srun_unbuffered, srun_mpibind, srun_distribution, srun_nodes, srun_cpus_per_task, - "--ntasks=%i" % np \ + "--ntasks=%i" % np, extra_sargs, \ ] + commandList def canRun(self, test):