From 33cf281aa559eab50f915e92a3bdeb58cc59e3a3 Mon Sep 17 00:00:00 2001
From: hpratt <hpratt1@proton.me>
Date: Mon, 23 Oct 2023 12:04:46 -0400
Subject: [PATCH] add local singularity image parameter

---
 build.gradle.kts                                      | 2 +-
 src/main/kotlin/krews/config/BaseConfig.kt            | 8 +++++++-
 src/main/kotlin/krews/executor/bsub/BsubExecutor.kt   | 5 +++--
 src/main/kotlin/krews/executor/slurm/SlurmExecutor.kt | 4 +++-
 4 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/build.gradle.kts b/build.gradle.kts
index a66df23..6afc8a7 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -13,7 +13,7 @@ plugins {
 }
 
 group = "io.krews"
-version = "0.14.15"
+version = "0.15.0"
 
 repositories {
     maven { setUrl("https://dl.bintray.com/kotlin/kotlin-eap") }
diff --git a/src/main/kotlin/krews/config/BaseConfig.kt b/src/main/kotlin/krews/config/BaseConfig.kt
index 8af547b..1dfe3df 100644
--- a/src/main/kotlin/krews/config/BaseConfig.kt
+++ b/src/main/kotlin/krews/config/BaseConfig.kt
@@ -1,5 +1,9 @@
 package krews.config
 
+data class SingularityConfig (
+    val localImage: String? = null
+)
+
 data class TaskConfig (
     // Task level input parameters
     val params: Map<String, Any> = mapOf(),
@@ -14,7 +18,9 @@ data class TaskConfig (
     // The number of tasks "executions" that will be run with the same job / vm.
     val grouping: Int = 1,
     // Docker image; defaults to one provided by task
-    val dockerImage: String? = null
+    val dockerImage: String? = null,
+    // optional Singularity config
+    val singularity: SingularityConfig? = null
 )
 
 data class WorkflowConfig (
diff --git a/src/main/kotlin/krews/executor/bsub/BsubExecutor.kt b/src/main/kotlin/krews/executor/bsub/BsubExecutor.kt
index 6e2f51f..8abc475 100644
--- a/src/main/kotlin/krews/executor/bsub/BsubExecutor.kt
+++ b/src/main/kotlin/krews/executor/bsub/BsubExecutor.kt
@@ -191,10 +191,11 @@ class BsubExecutor(private val workflowConfig: WorkflowConfig) : LocallyDirected
             bsubScript.append("export SINGULARITY_BIND=\"$binds\"\n")
 
             // Add running the task to script
-            val remoteImage = !taskRunContext.dockerImage.endsWith(".sif")
+            val trueImage = taskConfig.singularity?.localImage ?: taskRunContext.dockerImage
+            val remoteImage = !trueImage.endsWith(".sif")
             bsubScript.append("\n")
             bsubScript.append("# Run task command.\n")
-            bsubScript.append("singularity exec --containall ${if (remoteImage) "docker://" else ""}${taskRunContext.dockerImage} $containerCommand")
+            bsubScript.append("singularity exec --containall ${if (remoteImage) "docker://" else ""}${trueImage} $containerCommand")
             bsubScript.append("\n")
 
             // Add copying output files into output dir to script
diff --git a/src/main/kotlin/krews/executor/slurm/SlurmExecutor.kt b/src/main/kotlin/krews/executor/slurm/SlurmExecutor.kt
index 282bc76..5bcb7c4 100644
--- a/src/main/kotlin/krews/executor/slurm/SlurmExecutor.kt
+++ b/src/main/kotlin/krews/executor/slurm/SlurmExecutor.kt
@@ -195,9 +195,11 @@ class SlurmExecutor(private val workflowConfig: WorkflowConfig) : LocallyDirecte
             sbatchScript.append("export SINGULARITY_BIND=\"$binds\"\n")
 
             // Add running the task to script
+            val trueImage = taskConfig.singularity?.localImage ?: taskRunContext.dockerImage
+            val remoteImage = !trueImage.endsWith(".sif")
             sbatchScript.append("\n")
             sbatchScript.append("# Run task command.\n")
-            sbatchScript.append("singularity exec --containall docker://${taskRunContext.dockerImage} $containerCommand")
+            sbatchScript.append("singularity exec --containall ${if (remoteImage) "docker://" else ""}${trueImage} $containerCommand")
             sbatchScript.append("\n")
 
             // Add copying output files into output dir to script