Allow SLURM executor option --mem-per-cpu (#4023) [ci fast]

Signed-off-by: Fritjof Lammers <[email protected]> Signed-off-by: Fritjof Lammers <[email protected]> Signed-off-by: Ben Sherman <[email protected]> Co-authored-by: Ben Sherman <[email protected]>
nextflow-io · Jul 6, 2023 · 96c04e3 · 96c04e3
1 parent 0f4d886
commit 96c04e3
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 3 deletions.
diff --git a/docs/config.md b/docs/config.md
@@ -608,6 +608,12 @@ The following settings are available:
 `executor.name`
 : The name of the executor to be used (default: `local`).
 
+`executor.perCpuMemAllocation`
+: :::{versionadded} 23.07.0-edge
+  :::
+: *Used only by the {ref}`slurm-executor` executor.*
+: When `true`, specifies memory allocations for SLURM jobs as `--mem-per-cpu <task.memory / task.cpus>` instead of `--mem <task.memory>`.
+
 `executor.perJobMemLimit`
 : Specifies Platform LSF *per-job* memory limit mode. See {ref}`lsf-executor`.
 

diff --git a/docs/executor.md b/docs/executor.md
@@ -513,6 +513,10 @@ Resource requests and other job characteristics can be controlled via the follow
 SLURM partitions can be specified with the `queue` directive.
 :::
 
-:::{tip}
+:::{note}
 Nextflow does not provide direct support for SLURM multi-clusters. If you need to submit workflow executions to a cluster other than the current one, specify it with the `SLURM_CLUSTERS` variable in the launch environment.
 :::
+
+:::{versionadded} 23.07.0-edge
+Some SLURM clusters require memory allocations to be specified with `--mem-per-cpu` instead of `--mem`. You can specify `executor.perCpuMemAllocation = true` in the Nextflow configuration to enable this behavior. Nextflow will automatically compute the memory per CPU for each task (by default 1 CPU is used).
+:::
diff --git a/modules/nextflow/src/main/groovy/nextflow/executor/SlurmExecutor.groovy b/modules/nextflow/src/main/groovy/nextflow/executor/SlurmExecutor.groovy
@@ -37,6 +37,8 @@ class SlurmExecutor extends AbstractGridExecutor {
 
     static private Pattern SUBMIT_REGEX = ~/Submitted batch job (\d+)/
 
+    private boolean perCpuMemAllocation
+
     private boolean hasSignalOpt(Map config) {
         def opts = config.clusterOptions?.toString()
         return opts ? opts.contains('--signal ') || opts.contains('--signal=') : false
@@ -76,7 +78,11 @@ class SlurmExecutor extends AbstractGridExecutor {
             // be stored, just collected). In both cases memory use is based upon the job's
             // Resident Set Size (RSS). A task may exceed the memory limit until the next periodic
             // accounting sample. -- https://slurm.schedmd.com/sbatch.html
-            result << '--mem' << task.config.getMemory().toMega().toString() + 'M'
+            final mem = task.config.getMemory().toMega()
+            if( perCpuMemAllocation )
+                result << '--mem-per-cpu' << mem.intdiv(task.config.getCpus()).toString() + 'M'
+            else
+                result << '--mem' << mem.toString() + 'M'
         }
 
         // the requested partition (a.k.a queue) name
@@ -190,6 +196,12 @@ class SlurmExecutor extends AbstractGridExecutor {
         return result
     }
 
+    @Override
+    void register() {
+        super.register()
+        perCpuMemAllocation = session.getExecConfigProp(name, 'perCpuMemAllocation', false)
+    }
+
     @Override
     protected boolean pipeLauncherScript() {
         return isFusionEnabled()

diff --git a/modules/nextflow/src/test/groovy/nextflow/executor/SlurmExecutorTest.groovy b/modules/nextflow/src/test/groovy/nextflow/executor/SlurmExecutorTest.groovy
@@ -80,7 +80,7 @@ class SlurmExecutorTest extends Specification {
     def testGetHeaders() {
 
         setup:
-        // LSF executor
+        // SLURM executor
         def executor = [:] as SlurmExecutor
 
         // mock process
@@ -186,6 +186,23 @@ class SlurmExecutorTest extends Specification {
                 #SBATCH -x 3
                 '''
                 .stripIndent().leftTrim()
+
+        // test perCpuMemAllocation
+        when:
+        executor.@perCpuMemAllocation = true
+        task.config = new TaskConfig()
+        task.config.cpus = 8
+        task.config.memory = '24 GB'
+        then:
+        executor.getHeaders(task) == '''
+                #SBATCH -J nf-the_task_name
+                #SBATCH -o /work/path/.command.log
+                #SBATCH --no-requeue
+                #SBATCH --signal B:USR2@30
+                #SBATCH -c 8
+                #SBATCH --mem-per-cpu 3072M
+                '''
+                .stripIndent().leftTrim()
     }
 
     def testWorkDirWithBlanks() {