support boot disk size config

weng-lab · Mar 14, 2023 · 74d4dec · 74d4dec
1 parent f1e1d64
commit 74d4dec
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -8,7 +8,7 @@ plugins {
 }
 
 group = "io.krews"
-version = "0.13.0"
+version = "0.13.1"
 
 repositories {
     maven { setUrl("http://dl.bintray.com/kotlin/kotlin-eap") }

diff --git a/src/main/kotlin/krews/config/GoogleConfig.kt b/src/main/kotlin/krews/config/GoogleConfig.kt
@@ -36,7 +36,9 @@ data class GoogleTaskConfig(
     // GPUs to attach to the VM
     val gpus: GoogleGPUConfig? = null,
     // Image for the machine boot disk
-    val bootImage: String? = null
+    val bootImage: String? = null,
+    // Size of the boot disk
+    val bootDiskSizeGb: Int? = null
 )
 
 data class GoogleGPUConfig(

diff --git a/src/main/kotlin/krews/executor/google/GoogleLocalExecutor.kt b/src/main/kotlin/krews/executor/google/GoogleLocalExecutor.kt
@@ -74,6 +74,8 @@ class GoogleLocalExecutor(workflowConfig: WorkflowConfig) : LocallyDirectedExecu
         val maxCpus = taskRunContexts.map { it.cpus }.maxBy { it ?: -1 } ?: taskConfig.google?.cpus
         val maxMemory = taskRunContexts.map { it.memory }.maxBy { it?.bytes ?: -1 } ?: taskConfig.google?.mem
         virtualMachine.machineType = googleMachineType(taskConfig.google, maxCpus, maxMemory)
+        if (taskConfig.google?.bootDiskSizeGb != null)
+            virtualMachine.setBootDiskSizeGb(taskConfig.google?.bootDiskSizeGb)
 
         if (taskConfig.google?.gpus != null) {
             val acceleratorConfig = Accelerator()