Add azure batch pool virtualNetwork option (#3723)

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

Author: bentsherman

docs/azure.rst

@@ -418,6 +418,7 @@ azure.batch.pools.<name>.scaleInterval          Specify the interval at which to
 azure.batch.pools.<name>.schedulePolicy         Specify the scheduling policy for the pool identified with ``<name>``. It can be either ``spread`` or ``pack`` (default: ``spread``).
 azure.batch.pools.<name>.privileged             Enable the task to run with elevated access. Ignored if `runAs` is set (default: ``false``).
 azure.batch.pools.<name>.runAs                  Specify the username under which the task is run. The user must already exist on each node of the pool.
+azure.batch.pools.<name>.virtualNetwork         Specify the subnet ID of a virtual network in which to create the pool (requires version ``23.03.0-edge`` or later).
 azure.registry.server                           Specify the container registry from which to pull the Docker images (default: ``docker.io``, requires ``[email protected]``).
 azure.registry.userName                         Specify the username to connect to a private container registry (requires ``[email protected]``).
 azure.registry.password                         Specify the password to connect to a private container registry (requires ``[email protected]``).

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

@@ -40,6 +40,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.MountConfiguration
+import com.microsoft.azure.batch.protocol.models.NetworkConfiguration
 import com.microsoft.azure.batch.protocol.models.OutputFile
 import com.microsoft.azure.batch.protocol.models.OutputFileBlobContainerDestination
 import com.microsoft.azure.batch.protocol.models.OutputFileDestination
@@ -663,7 +664,6 @@ class AzBatchService implements Closeable {
                 .withCommandLine('bash -c "chmod +x azcopy && mkdir \$AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy \$AZ_BATCH_NODE_SHARED_DIR/bin/" ')
                 .withResourceFiles(resourceFiles)
 
-
         final poolParams = new PoolAddParameter()
                 .withId(spec.poolId)
                 .withVirtualMachineConfiguration(poolVmConfig(spec.opts))
@@ -674,6 +674,10 @@ class AzBatchService implements Closeable {
                 .withTaskSlotsPerNode(spec.vmType.numberOfCores)
                 .withStartTask(poolStartTask)
 
+        // virtual network
+        if( spec.opts.virtualNetwork )
+            poolParams.withNetworkConfiguration( new NetworkConfiguration().withSubnetId(spec.opts.virtualNetwork) )
+
         // scheduling policy
         if( spec.opts.schedulePolicy ) {
             final pol = ComputeNodeFillType.fromString(spec.opts.schedulePolicy)

plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy

@@ -65,6 +65,8 @@ class AzPoolOpts implements CacheFunnel {
     String userName
     String password
 
+    String virtualNetwork
+
     AzPoolOpts() {
         this(Collections.emptyMap())
     }
@@ -86,6 +88,7 @@ class AzPoolOpts implements CacheFunnel {
         this.registry = opts.registry
         this.userName = opts.userName
         this.password = opts.password
+        this.virtualNetwork = opts.virtualNetwork
     }
 
     @Override
@@ -104,6 +107,7 @@ class AzPoolOpts implements CacheFunnel {
         hasher.putBoolean(autoScale)
         hasher.putUnencodedChars(scaleFormula ?: '')
         hasher.putUnencodedChars(schedulePolicy ?: '')
+        hasher.putUnencodedChars(virtualNetwork ?: '')
         return hasher
     }