diff --git a/docs/reference/config.mdx b/docs/reference/config.mdx index 15960aefc9..b31d77cc90 100644 --- a/docs/reference/config.mdx +++ b/docs/reference/config.mdx @@ -513,6 +513,12 @@ The client ID for an Azure [managed identity](https://learn.microsoft.com/en-us/ Enable autoscaling feature for the pool identified with ``. +##### `azure.batch.pools..ephemeralOsDisk` + + + +Use an ephemeral OS disk for the compute nodes in the pool identified with ``. The OS disk is placed on the node's local/cache storage instead of remote Azure storage, which can improve performance and reduce cost, but the disk and its data are lost when the node is deallocated. The VM size must support an ephemeral OS disk with a cache large enough for the OS image (default`false`). + ##### `azure.batch.pools..fileShareRootPath` The internal root mount point when mounting File Shares. Must be `/mnt/resource/batch/tasks/fsmounts` for CentOS nodes or `/mnt/batch/tasks/fsmounts` for Ubuntu nodes (defaultCentOS). diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy index b35fc0c41f..b02ef9681c 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy @@ -47,6 +47,9 @@ import com.azure.compute.batch.models.BatchTaskContainerSettings import com.azure.compute.batch.models.BatchTaskCreateContent import com.azure.compute.batch.models.BatchTaskSchedulingPolicy import com.azure.compute.batch.models.ContainerConfiguration +import com.azure.compute.batch.models.DiffDiskPlacement +import com.azure.compute.batch.models.DiffDiskSettings +import com.azure.compute.batch.models.OSDisk import com.azure.compute.batch.models.ContainerRegistryReference import com.azure.compute.batch.models.ContainerType import com.azure.compute.batch.models.ElevationLevel @@ -858,8 +861,18 @@ class AzBatchService implements Closeable { final image = getImage(opts) - new VirtualMachineConfiguration(image.imageReference, image.nodeAgentSkuId) + final vmConfig = new VirtualMachineConfiguration(image.imageReference, image.nodeAgentSkuId) .setContainerConfiguration(containerConfig) + + // use an ephemeral OS disk placed on the node's local cache disk instead of remote storage + if( opts.ephemeralOsDisk ) { + log.debug "[AZURE BATCH] Enabling ephemeral OS disk for Azure Batch pool" + vmConfig.setOsDisk( + new OSDisk().setEphemeralOSDiskSettings( + new DiffDiskSettings().setPlacement(DiffDiskPlacement.CACHE_DISK) ) ) + } + + return vmConfig } diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy index 9714374481..4a0bc6c6b8 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy @@ -64,6 +64,12 @@ class AzPoolOpts implements CacheFunnel, ConfigScope { """) final boolean lowPriority + @ConfigOption + @Description(""" + Use an ephemeral OS disk for the compute nodes in the pool identified with ``. The OS disk is placed on the node's local/cache storage instead of remote Azure storage, which can improve performance and reduce cost, but the disk and its data are lost when the node is deallocated. The VM size must support an ephemeral OS disk with a cache large enough for the OS image (default: `false`). + """) + final boolean ephemeralOsDisk + @ConfigOption @Description(""" The max number of virtual machines when using auto scaling. @@ -175,6 +181,7 @@ class AzPoolOpts implements CacheFunnel, ConfigScope { this.password = opts.password this.virtualNetwork = opts.virtualNetwork this.lowPriority = opts.lowPriority as boolean + this.ephemeralOsDisk = opts.ephemeralOsDisk as boolean } @Override @@ -195,6 +202,7 @@ class AzPoolOpts implements CacheFunnel, ConfigScope { hasher.putUnencodedChars(schedulePolicy ?: '') hasher.putUnencodedChars(virtualNetwork ?: '') hasher.putBoolean(lowPriority) + hasher.putBoolean(ephemeralOsDisk) hasher.putUnencodedChars(startTask.script ?: '') hasher.putBoolean(startTask.privileged) return hasher diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy index 0cafb05aa0..8182acb75d 100644 --- a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy @@ -24,11 +24,14 @@ import dev.failsafe.function.CheckedPredicate import com.azure.compute.batch.models.BatchPool import com.azure.compute.batch.models.BatchJobCreateContent +import com.azure.compute.batch.models.BatchSupportedImage +import com.azure.compute.batch.models.DiffDiskPlacement import com.azure.compute.batch.models.ElevationLevel import com.azure.compute.batch.models.EnvironmentSetting import com.azure.core.exception.HttpResponseException import com.azure.core.http.HttpResponse import com.azure.identity.ManagedIdentityCredential +import com.azure.json.JsonProviders import com.google.common.hash.HashCode import nextflow.Global import nextflow.Session @@ -516,7 +519,7 @@ class AzBatchServiceTest extends Specification { then: 1 * svc.guessBestVm(LOC, CPUS, MEM, null, TYPE) >> VM and: - spec.poolId == 'nf-pool-42f3635f3fb8b71160900efa959f7809-Standard_X1' + spec.poolId == 'nf-pool-e3a346bceb6506d94e3da8f523d68cb6-Standard_X1' spec.metadata == [foo: 'bar'] } @@ -1185,4 +1188,38 @@ class AzBatchServiceTest extends Specification { thrown(IllegalArgumentException) createCalls == 1 } + + private static BatchSupportedImage supportedImage() { + final json = '{"nodeAgentSKUId":"batch.node.ubuntu 24.04","imageReference":{"publisher":"microsoft-dsvm","offer":"ubuntu-hpc","sku":"2204"}}' + JsonProviders.createReader(json).withCloseable { BatchSupportedImage.fromJson(it) } + } + + def 'should not set OS disk when ephemeralOsDisk is disabled' () { + given: + def exec = createExecutor() + AzBatchService svc = Spy(new AzBatchService(exec)) + def opts = new AzPoolOpts([:]) + + when: + def result = svc.poolVmConfig(opts) + then: + 1 * svc.getImage(opts) >> supportedImage() + and: + result.osDisk == null + } + + def 'should set ephemeral OS disk when ephemeralOsDisk is enabled' () { + given: + def exec = createExecutor() + AzBatchService svc = Spy(new AzBatchService(exec)) + def opts = new AzPoolOpts([ephemeralOsDisk: true]) + + when: + def result = svc.poolVmConfig(opts) + then: + 1 * svc.getImage(opts) >> supportedImage() + and: + result.osDisk != null + result.osDisk.ephemeralOSDiskSettings.placement == DiffDiskPlacement.CACHE_DISK + } } diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy index f47d8d673a..e5b27592eb 100644 --- a/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy @@ -46,6 +46,7 @@ class AzPoolOptsTest extends Specification { !opts.password !opts.virtualNetwork !opts.lowPriority + !opts.ephemeralOsDisk !opts.startTask.script !opts.startTask.privileged } @@ -70,6 +71,7 @@ class AzPoolOptsTest extends Specification { password: 'some-pwd', virtualNetwork: 'some-vnet', lowPriority: true, + ephemeralOsDisk: true, startTask: [ script: 'echo hello-world', privileged: true @@ -94,6 +96,7 @@ class AzPoolOptsTest extends Specification { opts.password == 'some-pwd' opts.virtualNetwork == 'some-vnet' opts.lowPriority + opts.ephemeralOsDisk opts.startTask.script == 'echo hello-world' opts.startTask.privileged }