In Google Batch, support user specified boot images

nextflow-io · Aug 30, 2024 · f99ec18 · f99ec18
1 parent 5a37e61
commit f99ec18
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/docs/config.md b/docs/config.md
@@ -886,6 +886,11 @@ The following settings are available for Google Cloud Batch:
   :::
 : Define the set of allowed locations for VMs to be provisioned. See [Google documentation](https://cloud.google.com/batch/docs/reference/rest/v1/projects.locations.jobs#locationpolicy) for details (default: no restriction).
 
+`google.batch.bootDiskImage`
+: :::{versionadded} 24.08.0-edge
+  :::
+: Set the image URI of the virtual machine boot disk, e.g `batch-debian`. See [Google documentation](https://cloud.google.com/batch/docs/vm-os-environment-overview#vm-os-image-options) for details (default: none).
+
 `google.batch.bootDiskSize`
 : Set the size of the virtual machine boot disk, e.g `50.GB` (default: none).
 

diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
@@ -332,6 +332,9 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
                 instancePolicyOrTemplate.setInstallGpuDrivers(true)
             }
 
+            if( executor.config.bootDiskImage )
+                instancePolicy.setBootDisk(AllocationPolicy.Disk.newBuilder().setImage( executor.config.bootDiskImage ))
+
             if( fusionEnabled() && !disk ) {
                 disk = new DiskResource(request: '375 GB', type: 'local-ssd')
                 log.debug "[GOOGLE BATCH] Process `${task.lazyName()}` - adding local volume as fusion scratch: $disk"

diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/client/BatchConfig.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/client/BatchConfig.groovy
@@ -38,6 +38,7 @@ class BatchConfig {
     private GoogleOpts googleOpts
     private GoogleCredentials credentials
     private List<String> allowedLocations
+    private String bootDiskImage
     private MemoryUnit bootDiskSize
     private String cpuPlatform
     private int maxSpotAttempts
@@ -54,6 +55,7 @@ class BatchConfig {
     GoogleOpts getGoogleOpts() { return googleOpts }
     GoogleCredentials getCredentials() { return credentials }
     List<String> getAllowedLocations() { allowedLocations }
+    String getBootDiskImage() { bootDiskImage }
     MemoryUnit getBootDiskSize() { bootDiskSize }
     String getCpuPlatform() { cpuPlatform }
     int getMaxSpotAttempts() { maxSpotAttempts }
@@ -72,6 +74,7 @@ class BatchConfig {
         result.googleOpts = GoogleOpts.create(session)
         result.credentials = result.googleOpts.credentials
         result.allowedLocations = session.config.navigate('google.batch.allowedLocations', List.of()) as List<String>
+        result.bootDiskImage = session.config.navigate('google.batch.bootDiskImage')
         result.bootDiskSize = session.config.navigate('google.batch.bootDiskSize') as MemoryUnit
         result.cpuPlatform = session.config.navigate('google.batch.cpuPlatform')
         result.maxSpotAttempts = session.config.navigate('google.batch.maxSpotAttempts', DEFAULT_MAX_SPOT_ATTEMPTS) as int