Azure Batch configurable startTask

docs/azure.md

@@ -380,6 +380,19 @@ The value of the setting must be the identifier of a subnet available in the vir
 Batch Authentication with Shared Keys does not allow to link external resources (like Virtual Networks) to the pool. Therefore, Active Directory Authentication must be used in conjunction with the `virtualNetwork` setting.
 :::
 
+### Start Task
+
+Nextflow uses azcopy to stage files in and out of the worker nodes. To do this, it installs azcopy to a shared directory by running a start task. If you have additional requirements for the worker nodes, you can modify this start task by setting the property `startTask`. This is the default shell script:
+
+```shell
+bash -c "chmod +x azcopy && mkdir $AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy $AZ_BATCH_NODE_SHARED_DIR/bin/"
+```
+
+:::{warning}
+If you modify the `startTask` causing azcopy to be unavailable on the worker machine your tasks be unable to correctly stage files in and out and they may fail.
+:::
+
+
 ## Microsoft Entra (formerly Active Directory Authentication)
 
 :::{versionadded} 22.11.0-edge

docs/config.md

@@ -414,6 +414,13 @@ The following settings are available:
 : *New in `nf-azure` version `0.11.0`*
 : Specify the ID of the Compute Node agent SKU which the pool identified with `<name>` supports (default: `batch.node.centos 8`).
 
+`azure.batch.pools.<name>.startTask`
+: *New in `nf-azure` version `1.5.2`*
+: Specify the startTask that is executed as the node joins the Azure Batch node pool (default: `bash -c "chmod +x azcopy && mkdir \$AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy \$AZ_BATCH_NODE_SHARED_DIR/bin/"`).
+
+`azure.batch.pools.<name>.startTaskPrivileged`
+: Enable the task to run with elevated access. Ignored if `runAs` is set (default: `false`).
+
 `azure.batch.pools.<name>.virtualNetwork`
 : :::{versionadded} 23.03.0-edge
   :::

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

@@ -430,7 +430,7 @@ class AzBatchService implements Closeable {
 
         return new TaskAddParameter()
                 .withId(taskId)
-                .withUserIdentity(userIdentity(pool.opts.privileged, pool.opts.runAs))
+                .withUserIdentity(userIdentity(pool.opts.privileged, pool.opts.runAs, AutoUserScope.TASK))
                 .withContainerSettings(containerOpts)
                 .withCommandLine(cmd)
                 .withResourceFiles(resourceFileUrls(task,sas))
@@ -660,7 +660,7 @@ class AzBatchService implements Closeable {
          *
          * https://github.com/MicrosoftDocs/azure-docs/blob/master/articles/batch/batch-docker-container-workloads.md#:~:text=Run%20container%20applications%20on%20Azure,compatible%20containers%20on%20the%20nodes.
          */
-        final containerConfig = new ContainerConfiguration().withType(DOCKER_COMPATIBLE);
+        final containerConfig = new ContainerConfiguration().withType(DOCKER_COMPATIBLE)
         final registryOpts = config.registry()
 
         if( registryOpts && registryOpts.isConfigured() ) {
@@ -681,18 +681,40 @@ class AzBatchService implements Closeable {
                 .withContainerConfiguration(containerConfig)
     }
 
-    protected StartTask createStartTask() {
-        if( config.batch().getCopyToolInstallMode() != CopyToolInstallMode.node )
-            return null
 
-        final resourceFiles = new ArrayList(10)
-        resourceFiles << new ResourceFile()
-            .withHttpUrl(AZCOPY_URL)
-            .withFilePath('azcopy')
+    protected StartTask createStartTask(AzPoolOpts opts) {
+
+        def startCmd = [] as ArrayList
+        final resourceFiles = [] as ArrayList
 
-        return new StartTask()
-            .withCommandLine('bash -c "chmod +x azcopy && mkdir \$AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy \$AZ_BATCH_NODE_SHARED_DIR/bin/" ')
-            .withResourceFiles(resourceFiles)
+        // Get any custom start task command
+        if ( opts.startTask ) {
+            startCmd << opts.startTask
+            log.debug "Adding custom start task to command: ${opts.startTask}"
+        }
+
+        // If enabled, append azcopy installer to start task command
+        if( config.batch().getCopyToolInstallMode() == CopyToolInstallMode.node ) {
+            startCmd << 'chmod +x azcopy && mkdir \$AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy \$AZ_BATCH_NODE_SHARED_DIR/bin/'
+            resourceFiles << new ResourceFile()
+                .withHttpUrl(AZCOPY_URL)
+                .withFilePath('azcopy')
+            log.debug "Adding azcopy installer to start task"
+        }
+
+        final startTaskCmd = "bash -c \"${startCmd.join('; ')}\""
+        log.debug "Start task final command: $startTaskCmd"
+
+        // If there is no start task contents we return a null to indicate no start task
+        if ( startCmd.isEmpty() && resourceFiles.isEmpty() ) {
+            return null
+        } else {
+        // else return a StartTask object with the start task command and resource files
+            return new StartTask()
+                .withCommandLine(startTaskCmd)
+                .withResourceFiles(resourceFiles)
+                .withUserIdentity(userIdentity(opts.startTaskPrivileged, null, AutoUserScope.POOL))
+        }
     }
 
     protected void createPool(AzVmPoolSpec spec) {
@@ -706,7 +728,7 @@ class AzBatchService implements Closeable {
                 // https://docs.microsoft.com/en-us/azure/batch/batch-parallel-node-tasks
                 .withTaskSlotsPerNode(spec.vmType.numberOfCores)
 
-        final startTask = createStartTask()
+        final startTask = createStartTask(spec.opts)
         if( startTask ) {
             poolParams .withStartTask(startTask)
         }
@@ -860,14 +882,14 @@ class AzBatchService implements Closeable {
         }
     }
 
-    protected UserIdentity userIdentity(boolean  privileged, String runAs) {
+    protected UserIdentity userIdentity(boolean  privileged, String runAs, AutoUserScope scope) {
         UserIdentity identity = new UserIdentity()
         if (runAs) {
             identity.withUserName(runAs)
         } else  {
             identity.withAutoUser(new AutoUserSpecification()
                     .withElevationLevel(privileged ? ElevationLevel.ADMIN : ElevationLevel.NON_ADMIN)
-                    .withScope(AutoUserScope.TASK))
+                    .withScope(scope))
         }
         return identity
     }

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

@@ -60,6 +60,9 @@ class AzPoolOpts implements CacheFunnel {
     Duration scaleInterval
     Integer maxVmCount
 
+    String startTask
+    boolean startTaskPrivileged
+
     String schedulePolicy // spread | pack
     String registry
     String userName
@@ -86,6 +89,8 @@ class AzPoolOpts implements CacheFunnel {
         this.schedulePolicy = opts.schedulePolicy
         this.scaleInterval = opts.scaleInterval as Duration ?: DEFAULT_SCALE_INTERVAL
         this.maxVmCount = opts.maxVmCount as Integer ?: vmCount *3
+        this.startTask = opts.startTask
+        this.startTaskPrivileged = opts.startTaskPrivileged ?: false
         this.registry = opts.registry
         this.userName = opts.userName
         this.password = opts.password
@@ -111,6 +116,8 @@ class AzPoolOpts implements CacheFunnel {
         hasher.putUnencodedChars(schedulePolicy ?: '')
         hasher.putUnencodedChars(virtualNetwork ?: '')
         hasher.putBoolean(lowPriority)
+        hasher.putUnencodedChars(startTask ?: '')
+        hasher.putBoolean(startTaskPrivileged)
         return hasher
     }
 

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

@@ -225,6 +225,81 @@ class AzBatchServiceTest extends Specification {
     }
 
 
+    def 'should configure default startTask' () {
+        given:
+        def CONFIG = [batch:[copyToolInstallMode: 'node']]
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        def svc = new AzBatchService(exec)
+
+        when:
+        def configuredStartTask = svc.createStartTask( new AzPoolOpts() )
+        then:
+        configuredStartTask.commandLine == 'bash -c "chmod +x azcopy && mkdir $AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy $AZ_BATCH_NODE_SHARED_DIR/bin/"'
+    }
+
+    def 'should configure custom startTask' () {
+        given:
+        def CONFIG = [batch:[copyToolInstallMode: 'node']]
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        def svc = new AzBatchService(exec)
+
+        when:
+        def configuredStartTask = svc.createStartTask( new AzPoolOpts(startTask: 'echo hello-world') )
+        then:
+        configuredStartTask.commandLine == 'bash -c "echo hello-world; chmod +x azcopy && mkdir $AZ_BATCH_NODE_SHARED_DIR/bin/ && cp azcopy $AZ_BATCH_NODE_SHARED_DIR/bin/"'
+    }
+
+    def 'should configure not install AzCopy because copyToolInstallMode is off' () {
+        given:
+        def CONFIG = [batch:[copyToolInstallMode: 'off']]
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        def svc = new AzBatchService(exec)
+
+        when:
+        def configuredStartTask = svc.createStartTask( new AzPoolOpts(startTask: 'echo hello-world') )
+        then:
+        configuredStartTask.commandLine == 'bash -c "echo hello-world"'
+        configuredStartTask.resourceFiles == []
+    }
+
+    def 'should configure not install AzCopy because copyToolInstallMode is task' () {
+        given:
+        def CONFIG = [batch:[copyToolInstallMode: 'task']]
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        def svc = new AzBatchService(exec)
+
+        when:
+        def configuredStartTask = svc.createStartTask( new AzPoolOpts(startTask: 'echo hello-world') )
+        then:
+        configuredStartTask.commandLine == 'bash -c "echo hello-world"'
+        configuredStartTask.resourceFiles == []
+    }
+
+    def 'should create null startTask because no options are enabled' () {
+        given:
+        def CONFIG = [batch:[copyToolInstallMode: 'off']]
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        def svc = new AzBatchService(exec)
+
+        when:
+        def configuredStartTask = svc.createStartTask( new AzPoolOpts() )
+        then:
+        configuredStartTask == null
+    }
+
+    def 'should configure privileged startTask' () {
+        given:
+        def CONFIG = [batch:[copyToolInstallMode: 'node']]
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        def svc = new AzBatchService(exec)
+        and:
+
+        when:
+        def configuredStartTask = svc.createStartTask( new AzPoolOpts(startTaskPrivileged: true) )
+        then:
+        configuredStartTask.userIdentity().autoUser().elevationLevel().value == 'admin'
+    }
+
     def 'should check scaling formula' () {
         given:
         def exec = Mock(AzBatchExecutor) { getConfig() >> new AzConfig([:]) }
@@ -359,7 +434,7 @@ class AzBatchServiceTest extends Specification {
         then:
         1 * svc.guessBestVm(LOC, CPUS, MEM, TYPE) >> VM
         and:
-        spec.poolId == 'nf-pool-6e9cf97d3d846621464131d3842265ce-Standard_X1'
+        spec.poolId == 'nf-pool-289d374ac1622e709cf863bce2570cab-Standard_X1'
         spec.metadata == [foo: 'bar']
 
     }

plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy

@@ -47,6 +47,8 @@ class AzPoolOptsTest extends Specification {
         !opts.password
         !opts.virtualNetwork
         !opts.lowPriority
+        !opts.startTask
+        !opts.startTaskPrivileged
     }
 
     def 'should create pool with custom options' () {
@@ -68,7 +70,9 @@ class AzPoolOptsTest extends Specification {
             userName: 'some-user',
             password: 'some-pwd',
             virtualNetwork: 'some-vnet',
-            lowPriority: true
+            lowPriority: true,
+            startTask: 'echo hello-world',
+            startTaskPrivileged: true
         ])
         then:
         opts.runAs == 'foo'
@@ -89,6 +93,8 @@ class AzPoolOptsTest extends Specification {
         opts.password == 'some-pwd'
         opts.virtualNetwork == 'some-vnet'
         opts.lowPriority
+        opts.startTask == 'echo hello-world'
+        opts.startTaskPrivileged
     }
 
 }

plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzureConfigTest.groovy

@@ -116,7 +116,8 @@ class AzureConfigTest extends Specification {
                                                      runAs: 'root',
                                                      scaleFormula: 'x + y + z',
                                                      scaleInterval:  '15 min',
-                                                     schedulePolicy: 'pack' ]]
+                                                     schedulePolicy: 'pack',
+                                                     startTask: 'echo hello-world']]
                                      ]] ]
         }
 
@@ -145,6 +146,7 @@ class AzureConfigTest extends Specification {
         cfg.batch().pool('myPool').scaleInterval == Duration.of('15 min')
         cfg.batch().pool('myPool').privileged == true
         cfg.batch().pool('myPool').runAs == 'root'
+        cfg.batch().pool('myPool').startTask == 'echo hello-world'
     }
 
     def 'should get azure batch endpoint from account and location' () {
@@ -171,7 +173,7 @@ class AzureConfigTest extends Specification {
         cfg.batch().location == LOCATION
     }
 
-	def 'should get azure batch file share root point' () {
+    def 'should get azure batch file share root point' () {
 
         given:
         def KEY = 'xyz1343'
@@ -188,12 +190,12 @@ class AzureConfigTest extends Specification {
                                      endpoint: ENDPOINT,
                                      location: LOCATION,
                                      pools: [
-            		                     myPool1: [ fileShareRootPath: '/somewhere/over/the/rainbow' ],
-            		                     myPool2: [ sku:'batch.node.centos 8' ],
-            		                     myPool3: [ sku:'batch.node.ubuntu 20.04' ],
-            		                     myPool4: [ sku:'batch.node.debian 10', fileShareRootPath: '/mounting/here' ],
-            		                     myPool5: [ : ]]
-            	                 ]] ]
+                                         myPool1: [ fileShareRootPath: '/somewhere/over/the/rainbow' ],
+                                         myPool2: [ sku:'batch.node.centos 8' ],
+                                         myPool3: [ sku:'batch.node.ubuntu 20.04' ],
+                                         myPool4: [ sku:'batch.node.debian 10', fileShareRootPath: '/mounting/here' ],
+                                         myPool5: [ : ]]
+                                 ]] ]
         }
         def cfg = AzConfig.getConfig(session)
         then: