add GPUs to BSUB test

weng-lab · Oct 21, 2023 · 4d1fe79 · 4d1fe79
1 parent b1f3f9e
commit 4d1fe79
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -13,12 +13,13 @@ plugins {
 }
 
 group = "io.krews"
-version = "0.14.4"
+version = "0.14.5"
 
 repositories {
     maven { setUrl("https://dl.bintray.com/kotlin/kotlin-eap") }
     mavenCentral()
     jcenter()
+    mavenLocal()
 }
 
 dependencies {

diff --git a/src/main/kotlin/krews/config/BsubConfig.kt b/src/main/kotlin/krews/config/BsubConfig.kt
@@ -13,12 +13,14 @@ data class BsubWorkflowConfig(
 data class BsubTaskConfig(
     // Number of cpus. Can be used to override the runtime value.
     val cpus: Int? = null,
+    // Whether or not to request a GPU.
+    val gpu: Boolean? = null,
     // Memory capacity. Can be used to override the runtime value.
     val mem: Capacity? = null,
     // Time limit on the run time for the job in minutes.
     val time: Int? = null,
     // SBatch partition to use.
     val partition: String? = null,
-    // Additional sbatch arguments
+    // Additional bsub arguments
     val sbatchArgs: Map<String, String>? = null
 )
diff --git a/src/main/kotlin/krews/executor/bsub/BsubExecutor.kt b/src/main/kotlin/krews/executor/bsub/BsubExecutor.kt
@@ -76,6 +76,7 @@ class BsubExecutor(private val workflowConfig: WorkflowConfig) : LocallyDirected
         val cpus = taskConfig.bsub?.cpus ?: taskRunContexts.map { it.cpus }.maxBy { it ?: -1 }
         val timePerTask = taskConfig.bsub?.time ?: taskRunContexts.map { it.cpus }.maxBy { it ?: -1 }
         val time = if (timePerTask != null) timePerTask * taskRunContexts.size else null
+        val gpus = taskConfig.bsub?.gpu ?: "\"\""
 
         appendBsubParam(bsubScript, "J", bsubWorkflowJobName)
         appendBsubParam(bsubScript, "o", logsPath.resolve("out.txt"))
@@ -84,6 +85,7 @@ class BsubExecutor(private val workflowConfig: WorkflowConfig) : LocallyDirected
         appendBsubParam(bsubScript, "n", cpus)
         appendBsubParam(bsubScript, "W", time)
         appendBsubParam(bsubScript, "q", taskConfig.bsub?.partition)
+        appendBsubParam(bsubScript, "gpu", gpus)
         for ((argName, argVal) in taskConfig.bsub?.sbatchArgs ?: mapOf()) {
             appendBsubParam(bsubScript, argName, argVal)
         }