Add Google Batch native retry on spot termination (#4500)

This commit adds the automatic retry made by Google Batch when 
a spot reclaim terminates a job. 

This feature can be controlled by the setting:

```
google.batch.maxSpotAttempts = <number>
```  

By default, the execution is attempted 5 times. When 0 is specified, 
no attempt is made on a spot reclaim.

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

Author: bentsherman

docs/config.md

@@ -775,6 +775,11 @@ The following settings are available for Google Cloud Batch:
 `google.location`
 : The Google Cloud location where jobs are executed (default: `us-central1`).
 
+`google.batch.maxSpotAttempts`
+: :::{versionadded} 23.11.0-edge
+  :::
+: Max number of execution attempts of a job interrupted by a Compute Engine spot reclaim event (default: `5`).
+
 `google.project`
 : The Google Cloud project ID to use for pipeline execution
 

plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy

@@ -24,6 +24,7 @@ import com.google.cloud.batch.v1.AllocationPolicy
 import com.google.cloud.batch.v1.ComputeResource
 import com.google.cloud.batch.v1.Environment
 import com.google.cloud.batch.v1.Job
+import com.google.cloud.batch.v1.LifecyclePolicy
 import com.google.cloud.batch.v1.LogsPolicy
 import com.google.cloud.batch.v1.Runnable
 import com.google.cloud.batch.v1.ServiceAccount
@@ -207,6 +208,23 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
             )
             .addAllVolumes( launcher.getVolumes() )
 
+        // retry on spot reclaim
+        if( executor.config.maxSpotAttempts ) {
+            // Note: Google Batch uses the special exit status 50001 to signal
+            // the execution was terminated due a spot reclaim. When this happens
+            // The policy re-execute the jobs automatically up to `maxSpotAttempts` times
+            taskSpec
+                .setMaxRetryCount( executor.config.maxSpotAttempts )
+                .addLifecyclePolicies(
+                    LifecyclePolicy.newBuilder()
+                        .setActionCondition(
+                            LifecyclePolicy.ActionCondition.newBuilder()
+                                .addExitCodes(50001)
+                        )
+                        .setAction(LifecyclePolicy.Action.RETRY_TASK)
+                )
+        }
+
         // instance policy
         final allocationPolicy = AllocationPolicy.newBuilder()
         final instancePolicyOrTemplate = AllocationPolicy.InstancePolicyOrTemplate.newBuilder()

plugins/nf-google/src/main/nextflow/cloud/google/batch/client/BatchConfig.groovy

@@ -37,8 +37,9 @@ class BatchConfig {
     private List<String> allowedLocations
     private MemoryUnit bootDiskSize
     private String cpuPlatform
-    private boolean spot
+    private int maxSpotAttempts
     private boolean preemptible
+    private boolean spot
     private boolean usePrivateAddress
     private String network
     private String subnetwork
@@ -49,6 +50,7 @@ class BatchConfig {
     List<String> getAllowedLocations() { allowedLocations }
     MemoryUnit getBootDiskSize() { bootDiskSize }
     String getCpuPlatform() { cpuPlatform }
+    int getMaxSpotAttempts() { maxSpotAttempts }
     boolean getPreemptible() { preemptible }
     boolean getSpot() { spot }
     boolean getUsePrivateAddress() { usePrivateAddress }
@@ -63,8 +65,9 @@ class BatchConfig {
         result.allowedLocations = session.config.navigate('google.batch.allowedLocations', List.of()) as List<String>
         result.bootDiskSize = session.config.navigate('google.batch.bootDiskSize') as MemoryUnit
         result.cpuPlatform = session.config.navigate('google.batch.cpuPlatform')
-        result.spot = session.config.navigate('google.batch.spot',false)
+        result.maxSpotAttempts = session.config.navigate('google.batch.maxSpotAttempts',5) as int
         result.preemptible = session.config.navigate('google.batch.preemptible',false)
+        result.spot = session.config.navigate('google.batch.spot',false)
         result.usePrivateAddress = session.config.navigate('google.batch.usePrivateAddress',false)
         result.network = session.config.navigate('google.batch.network')
         result.subnetwork = session.config.navigate('google.batch.subnetwork')

plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy

@@ -130,6 +130,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
                 getAllowedLocations() >> ['zones/us-central1-a', 'zones/us-central1-c']
                 getBootDiskSize() >> BOOT_DISK
                 getCpuPlatform() >> CPU_PLATFORM
+                getMaxSpotAttempts() >> 5
                 getSpot() >> true
                 getNetwork() >> 'net-1'
                 getServiceAccountEmail() >> '[email protected]'
@@ -170,16 +171,20 @@ class GoogleBatchTaskHandlerTest extends Specification {
 
         and:
         def taskGroup = req.getTaskGroups(0)
-        def runnable = taskGroup.getTaskSpec().getRunnables(0)
+        def taskSpec = taskGroup.getTaskSpec()
+        def runnable = taskSpec.getRunnables(0)
         def allocationPolicy = req.getAllocationPolicy()
         def instancePolicy = allocationPolicy.getInstances(0).getPolicy()
         def networkInterface = allocationPolicy.getNetwork().getNetworkInterfaces(0)
         and:
-        taskGroup.getTaskSpec().getComputeResource().getBootDiskMib() == BOOT_DISK.toMega()
-        taskGroup.getTaskSpec().getComputeResource().getCpuMilli() == CPUS * 1_000
-        taskGroup.getTaskSpec().getComputeResource().getMemoryMib() == MEM.toMega()
-        taskGroup.getTaskSpec().getMaxRunDuration().getSeconds() == TIMEOUT.seconds
-        taskGroup.getTaskSpec().getVolumes(0).getMountPath() == '/tmp'
+        taskSpec.getComputeResource().getBootDiskMib() == BOOT_DISK.toMega()
+        taskSpec.getComputeResource().getCpuMilli() == CPUS * 1_000
+        taskSpec.getComputeResource().getMemoryMib() == MEM.toMega()
+        taskSpec.getMaxRunDuration().getSeconds() == TIMEOUT.seconds
+        taskSpec.getVolumes(0).getMountPath() == '/tmp'
+        taskSpec.getMaxRetryCount() == 5
+        taskSpec.getLifecyclePolicies(0).getActionCondition().getExitCodes(0) == 50001
+        taskSpec.getLifecyclePolicies(0).getAction().toString() == 'RETRY_TASK'
         and:
         runnable.getContainer().getCommandsList().join(' ') == '/bin/bash -o pipefail -c bash .command.run'
         runnable.getContainer().getImageUri() == CONTAINER_IMAGE