diff --git a/docs/config.md b/docs/config.md index 5adb99520e..70b70cc85d 100644 --- 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 index f8a10e9663..b7a33cafed 100644 --- 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.getBootDiskImage() ) + instancePolicy.setBootDisk( AllocationPolicy.Disk.newBuilder().setImage( executor.config.getBootDiskImage() ) ) + 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 index 9657754504..de9b6e2d38 100644 --- 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 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 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 + 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 diff --git a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy index b54a5ff7a8..ba21a84588 100644 --- a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy +++ b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy @@ -105,6 +105,7 @@ class GoogleBatchTaskHandlerTest extends Specification { !instancePolicy.getMachineType() !instancePolicy.getMinCpuPlatform() instancePolicy.getProvisioningModel().toString() == 'PROVISIONING_MODEL_UNSPECIFIED' + !instancePolicy.getBootDisk().getImage() and: allocationPolicy.getLocation().getAllowedLocationsCount() == 0 allocationPolicy.getNetwork().getNetworkInterfacesCount() == 0 @@ -121,6 +122,7 @@ class GoogleBatchTaskHandlerTest extends Specification { and: def ACCELERATOR = new AcceleratorResource(request: 1, type: 'nvidia-tesla-v100') def BOOT_DISK = MemoryUnit.of('10 GB') + def BOOT_IMAGE = 'batch-debian' def CONTAINER_IMAGE = 'ubuntu:22.1' def CONTAINER_OPTS = '--this --that' def CPU_PLATFORM = 'Intel Skylake' @@ -134,6 +136,7 @@ class GoogleBatchTaskHandlerTest extends Specification { getConfig() >> Mock(BatchConfig) { getAllowedLocations() >> ['zones/us-central1-a', 'zones/us-central1-c'] getBootDiskSize() >> BOOT_DISK + getBootDiskImage() >> BOOT_IMAGE getCpuPlatform() >> CPU_PLATFORM getMaxSpotAttempts() >> 5 getSpot() >> true @@ -211,6 +214,7 @@ class GoogleBatchTaskHandlerTest extends Specification { and: instancePolicy.getAccelerators(0).getCount() == 1 instancePolicy.getAccelerators(0).getType() == ACCELERATOR.type + instancePolicy.getBootDisk().getImage() == BOOT_IMAGE instancePolicy.getDisks(0).getNewDisk().getSizeGb() == DISK.request.toGiga() instancePolicy.getDisks(0).getNewDisk().getType() == DISK.type instancePolicy.getMachineType() == MACHINE_TYPE diff --git a/plugins/nf-google/src/test/nextflow/cloud/google/batch/client/BatchConfigTest.groovy b/plugins/nf-google/src/test/nextflow/cloud/google/batch/client/BatchConfigTest.groovy index c5dd59c76c..c6fc08f104 100644 --- a/plugins/nf-google/src/test/nextflow/cloud/google/batch/client/BatchConfigTest.groovy +++ b/plugins/nf-google/src/test/nextflow/cloud/google/batch/client/BatchConfigTest.groovy @@ -18,6 +18,7 @@ package nextflow.cloud.google.batch.client import nextflow.Session +import nextflow.util.MemoryUnit import spock.lang.Requires import spock.lang.Specification /** @@ -29,21 +30,20 @@ class BatchConfigTest extends Specification { @Requires({System.getenv('GOOGLE_APPLICATION_CREDENTIALS')}) def 'should create batch config' () { given: - def CONFIG = [google: [ - batch: [ - spot: true - ] - ] ] + def CONFIG = [:] def session = Mock(Session) { getConfig()>>CONFIG } when: def config = BatchConfig.create(session) then: - config.getSpot() + !config.getSpot() and: config.retryConfig.maxAttempts == 5 config.maxSpotAttempts == 0 config.autoRetryExitCodes == [50001] + and: + !config.bootDiskImage + !config.bootDiskSize } @Requires({System.getenv('GOOGLE_APPLICATION_CREDENTIALS')}) @@ -54,7 +54,9 @@ class BatchConfigTest extends Specification { spot: true, maxSpotAttempts: 8, autoRetryExitCodes: [50001, 50003, 50005], - retryPolicy: [maxAttempts: 10] + retryPolicy: [maxAttempts: 10], + bootDiskImage: 'batch-foo', + bootDiskSize: '100GB' ] ] ] def session = Mock(Session) { getConfig()>>CONFIG } @@ -67,6 +69,9 @@ class BatchConfigTest extends Specification { config.retryConfig.maxAttempts == 10 config.maxSpotAttempts == 8 config.autoRetryExitCodes == [50001, 50003, 50005] + and: + config.bootDiskImage == 'batch-foo' + config.bootDiskSize == MemoryUnit.of('100GB') } }