Azure jobs correctly deleted after completion

docs/config.md

@@ -313,6 +313,9 @@ The following settings are available:
 `azure.batch.copyToolInstallMode`
 : Specify where the `azcopy` tool used by Nextflow. When `node` is specified it's copied once during the pool creation. When `task` is provider, it's installed for each task execution (default: `node`).
 
+`azure.batch.terminateJobsOnCompletion`
+: Enables the Batch Job to automatically terminate a job once all tasks have completed (default: `true`).
+
 `azure.batch.deleteJobsOnCompletion`
 : Enable the automatic deletion of jobs created by the pipeline execution (default: `true`).
 

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

@@ -38,8 +38,10 @@ import com.microsoft.azure.batch.protocol.models.ContainerConfiguration
 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.MountConfiguration
 import com.microsoft.azure.batch.protocol.models.NetworkConfiguration
+import com.microsoft.azure.batch.protocol.models.OnAllTasksComplete
 import com.microsoft.azure.batch.protocol.models.OutputFile
 import com.microsoft.azure.batch.protocol.models.OutputFileBlobContainerDestination
 import com.microsoft.azure.batch.protocol.models.OutputFileDestination
@@ -783,14 +785,37 @@ class AzBatchService implements Closeable {
         apply(() -> client.taskOperations().deleteTask(key.jobId, key.taskId))
     }
 
-    protected void cleanupJobs() {
+    protected void terminateJobs() {
+        /* 
+        We set the job to terminate when all tasks are complete rather than directly terminating, this allows Azure Batch to handle the termination for us.
+        */
+
         for( Map.Entry<TaskProcessor,String> entry : allJobIds ) {
             final proc = entry.key
             final jobId = entry.value
-            if( proc.hasErrors() ) {
-                log.debug "Preserving Azure job with error: ${jobId}"
-                continue
+
+            try {
+                log.trace "Terminating Azure job ${jobId}"
+
+                CloudJob job = apply(() -> client.jobOperations().getJob(jobId))
+                final poolInfo = job.poolInfo()
+
+                JobUpdateParameter jobParameter = new JobUpdateParameter()
+                        .withOnAllTasksComplete(OnAllTasksComplete.TERMINATE_JOB)
+                        .withPoolInfo(poolInfo)
+
+                apply(() -> client.jobOperations().updateJob(jobId, jobParameter))
+            }
+            catch (Exception e) {
+                log.warn "Unable to terminate Azure Batch job ${jobId} - Reason: ${e.message ?: e}"
             }
+        }
+    }
+
+    protected void cleanupJobs() {
+        for( Map.Entry<TaskProcessor,String> entry : allJobIds ) {
+            final proc = entry.key
+            final jobId = entry.value
 
             try {
                 log.trace "Deleting Azure job ${jobId}"
@@ -826,6 +851,11 @@ class AzBatchService implements Closeable {
     }
     @Override
     void close() {
+        // Terminate existing jobs to prevent them occupying quota
+        if( config.batch().terminateJobsOnCompletion!=Boolean.FALSE ) {
+            terminateJobs()
+        }
+
         // cleanup app successful jobs
         if( config.batch().deleteJobsOnCompletion!=Boolean.FALSE ) {
             cleanupJobs()

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

@@ -47,6 +47,7 @@ class AzBatchOpts implements CloudTransferOptions {
     String location
     Boolean autoPoolMode
     Boolean allowPoolCreation
+    Boolean terminateJobsOnCompletion
     Boolean deleteJobsOnCompletion
     Boolean deletePoolsOnCompletion
     CopyToolInstallMode copyToolInstallMode
@@ -62,6 +63,7 @@ class AzBatchOpts implements CloudTransferOptions {
         location = config.location
         autoPoolMode = config.autoPoolMode
         allowPoolCreation = config.allowPoolCreation
+        terminateJobsOnCompletion = config.terminateJobsOnCompletion
         deleteJobsOnCompletion = config.deleteJobsOnCompletion
         deletePoolsOnCompletion = config.deletePoolsOnCompletion
         pools = parsePools(config.pools instanceof Map ? config.pools as Map<String,Map> : Collections.<String,Map>emptyMap())