Add ephemeral OS disk support to Azure Batch pools

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 `<name>`.
 
+##### `azure.batch.pools.<name>.ephemeralOsDisk`
+
+<AddedInVersion version="26.10.0" />
+
+Use an ephemeral OS disk for the compute nodes in the pool identified with `<name>`. 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.<name>.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).

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
     }
 
 

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 `<name>`. 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

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
+    }
 }

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
     }