Add resource labels support for Azure Batch (nextflow-io#4178)

Signed-off-by: Ben Sherman <[email protected]>
Signed-off-by: Paolo Di Tommaso <[email protected]>
Co-authored-by: Paolo Di Tommaso <[email protected]>

docs/executor.md

@@ -28,6 +28,7 @@ Resource requests and other job characteristics can be controlled via the follow
 - {ref}`process-cpus`
 - {ref}`process-memory`
 - {ref}`process-queue`
+- {ref}`process-resourcelabels`
 - {ref}`process-time`
 
 See the {ref}`AWS Batch<aws-batch>` page for further configuration details.
@@ -55,6 +56,7 @@ Resource requests and other job characteristics can be controlled via the follow
 - {ref}`process-machineType`
 - {ref}`process-memory`
 - {ref}`process-queue`
+- {ref}`process-resourcelabels`
 - {ref}`process-time`
 
 See the {ref}`Azure Batch <azure-batch>` page for further configuration details.
@@ -191,8 +193,8 @@ Resource requests and other job characteristics can be controlled via the follow
 - {ref}`process-disk`
 - {ref}`process-machineType`
 - {ref}`process-memory`
-- {ref}`process-time`
 - {ref}`process-resourcelabels`
+- {ref}`process-time`
 
 See the {ref}`Google Cloud Batch <google-batch>` page for further configuration details.
 
@@ -218,6 +220,7 @@ Resource requests and other job characteristics can be controlled via the follow
 - {ref}`process-disk`
 - {ref}`process-machineType`
 - {ref}`process-memory`
+- {ref}`process-resourcelabels`
 - {ref}`process-time`
 
 See the {ref}`Google Life Sciences <google-lifesciences>` page for further configuration details.
@@ -312,6 +315,7 @@ Resource requests and other job characteristics can be controlled via the follow
 - {ref}`process-disk`
 - {ref}`process-memory`
 - {ref}`process-pod`
+- {ref}`process-resourcelabels`
 - {ref}`process-time`
 
 See the {ref}`Kubernetes <k8s-page>` page to learn how to set up a Kubernetes cluster to run Nextflow pipelines.

docs/process.md

@@ -2121,8 +2121,18 @@ process my_task {
 
 The limits and the syntax of the corresponding cloud provider should be taken into consideration when using resource labels.
 
-:::{note}
-Resource labels are currently only supported by the {ref}`awsbatch-executor`, {ref}`google-lifesciences-executor`, Google Cloud Batch and {ref}`k8s-executor` executors.
+Resource labels are currently supported by the following executors:
+
+- {ref}`awsbatch-executor`
+- {ref}`azurebatch-executor`
+- {ref}`google-batch-executor`
+- {ref}`google-lifesciences-executor`
+- {ref}`k8s-executor`
+
+:::{versionadded} 23.09.0-edge
+Resource labels are supported for Azure Batch when using automatic pool creation.
+
+Resource labels in Azure are added to pools, rather than jobs, in order to facilitate cost analysis. A new pool will be created for each new set of resource labels, therefore it is recommended to also set `azure.batch.deletePoolsOnCompletion = true` when using process-specific resource labels.
 :::
 
 See also: [label](#label)

plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy

@@ -39,6 +39,7 @@ import com.microsoft.azure.batch.protocol.models.ContainerRegistry
 import com.microsoft.azure.batch.protocol.models.ElevationLevel
 import com.microsoft.azure.batch.protocol.models.ImageInformation
 import com.microsoft.azure.batch.protocol.models.JobUpdateParameter
+import com.microsoft.azure.batch.protocol.models.MetadataItem
 import com.microsoft.azure.batch.protocol.models.MountConfiguration
 import com.microsoft.azure.batch.protocol.models.NetworkConfiguration
 import com.microsoft.azure.batch.protocol.models.OnAllTasksComplete
@@ -563,9 +564,10 @@ class AzBatchService implements Closeable {
             throw new IllegalArgumentException(msg)
         }
 
-        final key = CacheHelper.hasher([vmType.name, opts]).hash().toString()
+        final metadata = task.config.getResourceLabels()
+        final key = CacheHelper.hasher([vmType.name, opts, metadata]).hash().toString()
         final poolId = "nf-pool-$key-$vmType.name"
-        return new AzVmPoolSpec(poolId: poolId, vmType: vmType, opts: opts)
+        return new AzVmPoolSpec(poolId: poolId, vmType: vmType, opts: opts, metadata: metadata)
     }
 
     protected void checkPool(CloudPool pool, AzVmPoolSpec spec) {
@@ -698,6 +700,16 @@ class AzBatchService implements Closeable {
                 .withTaskSlotsPerNode(spec.vmType.numberOfCores)
                 .withStartTask(poolStartTask)
 
+        // resource labels
+        if( spec.metadata ) {
+            final metadata = spec.metadata.collect { name, value ->
+                new MetadataItem()
+                    .withName(name)
+                    .withValue(value)
+            }
+            poolParams.withMetadata(metadata)
+        }
+
         // virtual network
         if( spec.opts.virtualNetwork )
             poolParams.withNetworkConfiguration( new NetworkConfiguration().withSubnetId(spec.opts.virtualNetwork) )

plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzVmPoolSpec.groovy

@@ -36,4 +36,5 @@ class AzVmPoolSpec {
     String poolId
     AzVmType vmType
     AzPoolOpts opts
+    Map<String,String> metadata
 }

plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy

@@ -338,6 +338,7 @@ class AzBatchServiceTest extends Specification {
                 getMemory() >> MEM
                 getCpus() >> CPUS
                 getMachineType() >> TYPE
+                getResourceLabels() >> [foo: 'bar']
             }
         }
 
@@ -346,7 +347,8 @@ class AzBatchServiceTest extends Specification {
         then:
         1 * svc.guessBestVm(LOC, CPUS, MEM, TYPE) >> VM
         and:
-        spec.poolId == 'nf-pool-ddb1223ab79edfe07c0af2be7fceeb13-Standard_X1'
+        spec.poolId == 'nf-pool-9022a3fbfb5f93028d78fefaea5e21ab-Standard_X1'
+        spec.metadata == [foo: 'bar']
 
     }