Add max cpu and memory factor support.

Add a factor of max cpu and memory to restrict the max instance type. The default value is set to 2.
MemVerge · Oct 7, 2023 · b7a2c72 · b7a2c72
1 parent 7f881a6
commit b7a2c72
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 9 deletions.
diff --git a/plugins/nf-float/src/main/com/memverge/nextflow/FloatConf.groovy b/plugins/nf-float/src/main/com/memverge/nextflow/FloatConf.groovy
@@ -58,6 +58,9 @@ class FloatConf {
     float cpuFactor = 1
     float memoryFactory = 1
 
+    float maxCpuFactor = 2
+    float maxMemoryFactor = 2
+
     /**
      * Create a FloatConf instance and initialize the content from the
      * configuration.  The configuration should contain a "float" node.
@@ -152,6 +155,12 @@ class FloatConf {
         if (floatNode.cpuFactor) {
             this.cpuFactor = floatNode.cpuFactor as Float
         }
+        if (floatNode.maxCpuFactor) {
+            this.maxCpuFactor = floatNode.maxCpuFactor as Float
+        }
+        if (floatNode.maxMemoryFactor) {
+            this.maxMemoryFactor = floatNode.maxMemoryFactor as Float
+        }
         if (floatNode.memoryFactor) {
             this.memoryFactory = floatNode.memoryFactor as Float
         }

diff --git a/plugins/nf-float/src/main/com/memverge/nextflow/FloatGridExecutor.groovy b/plugins/nf-float/src/main/com/memverge/nextflow/FloatGridExecutor.groovy
@@ -28,8 +28,6 @@ import nextflow.util.ServiceName
 
 import java.nio.file.Path
 import java.nio.file.StandardCopyOption
-import java.util.regex.Matcher
-import java.util.regex.Pattern
 import java.util.stream.Collectors
 
 /**
@@ -124,7 +122,7 @@ class FloatGridExecutor extends AbstractGridExecutor {
         log.info "[float] sync the float binary, $res"
     }
 
-    private String getMemory(TaskRun task) {
+    private int getMemory(TaskRun task) {
         final mem = task.config.getMemory()
         Long giga = mem?.toGiga()
         if (!giga) {
@@ -134,7 +132,7 @@ class FloatGridExecutor extends AbstractGridExecutor {
         }
         giga = (long) ((float) (giga) * floatConf.memoryFactory)
         giga = Math.max(giga, DFT_MEM_GB)
-        return giga.toString()
+        return giga
     }
 
     private Collection<String> getExtra(TaskRun task) {
@@ -151,7 +149,7 @@ class FloatGridExecutor extends AbstractGridExecutor {
     private static List<String> splitWithQuotes(String input) {
         List<String> ret = new ArrayList<String>()
         int start = 0
-        boolean  inQuotes = false
+        boolean inQuotes = false
         for (int i = 0; i < input.size(); i++) {
             if (input[i] == '"') {
                 inQuotes = !inQuotes
@@ -335,8 +333,13 @@ class FloatGridExecutor extends AbstractGridExecutor {
         cmd << 'submit'
         getMountVols(task).forEach { cmd << '--dataVolume' << it }
         cmd << '--image' << container
-        cmd << '--cpu' << getCpu(task).toString()
-        cmd << '--mem' << getMemory(task)
+
+        int cpu = getCpu(task)
+        int maxCpu = (floatConf.maxCpuFactor * cpu.doubleValue()).intValue()
+        cmd << '--cpu' << "${cpu}:${maxCpu}".toString()
+        int memGiga = getMemory(task)
+        int maxMemGiga = (floatConf.maxMemoryFactor * memGiga.doubleValue()).intValue()
+        cmd << '--mem' << "${memGiga}:${maxMemGiga}".toString()
         cmd << '--job' << getScriptFilePath(handler, scriptFile)
         getEnv(handler).each { key, val ->
             cmd << '--env' << "${key}=${val}".toString()

diff --git a/plugins/nf-float/src/test/com/memverge/nextflow/FloatBaseTest.groovy b/plugins/nf-float/src/test/com/memverge/nextflow/FloatBaseTest.groovy
@@ -100,14 +100,16 @@ class FloatBaseTest extends BaseTest {
 
     def submitCmd(Map param = [:]) {
         def taskIndex = param.taskIndex?:'1'
+        def realCpu = param.cpu ?: cpu
+        def realMem = param.memory ?: mem
         return [bin, '-a', param.addr ?: addr,
                 '-u', user,
                 '-p', pass,
                 'submit',
                 '--dataVolume', param.nfs ?: nfs + ':' + workDir,
                 '--image', param.image ?: image,
-                '--cpu', param.cpu ?: cpu.toString(),
-                '--mem', param.memory ?: mem.toString(),
+                '--cpu', realCpu + ':' + realCpu * 2,
+                '--mem', realMem + ':' + realMem * 2,
                 '--job', script,
                 '--customTag', jobID(taskID),
                 '--customTag', "${FloatConf.NF_SESSION_ID}:uuid-$uuid",

diff --git a/plugins/nf-float/src/test/com/memverge/nextflow/FloatGridExecutorTest.groovy b/plugins/nf-float/src/test/com/memverge/nextflow/FloatGridExecutorTest.groovy
@@ -99,6 +99,24 @@ class FloatGridExecutorTest extends FloatBaseTest {
         cmd.join(' ') == expected.join(' ')
     }
 
+    def "max cpu and memory factor"() {
+        given:
+        final dataVolume = nfs + ':/data'
+        final exec = newTestExecutor([float: [
+                address : addr,
+                maxCpuFactor: 3,
+                maxMemoryFactor: 2.51,
+                nfs     : dataVolume]])
+        final task = newTask(exec)
+
+        when:
+        final cmd = exec.getSubmitCommandLine(task, Paths.get(script))
+        final expected = submitCmd(nfs: dataVolume)
+
+        then:
+        cmd.join(' ').contains("--cpu 5:15 --mem 10:25")
+    }
+
     def "add default local mount point"() {
         given:
         final exec = newTestExecutor()