Add Google Batch NOT_FOUND error management

modules/nextflow/src/main/groovy/nextflow/exception/ProcessStageException.groovy

@@ -24,5 +24,5 @@ import groovy.transform.InheritConstructors
  * @author Paolo Di Tommaso <paolo.ditommaso@gmail.com>
  */
 @InheritConstructors
-class ProcessStageException extends ProcessException implements ShowOnlyExceptionMessage {
+class ProcessStageException extends ProcessUnrecoverableException implements ShowOnlyExceptionMessage {
 }

modules/nextflow/src/main/groovy/nextflow/file/FilePorter.groovy

@@ -331,7 +331,7 @@ class FilePorter {
                     // remove the target file that could be have partially downloaded
                     cleanup(stagePath)
                     // check if a stage/download retry is allowed
-                    if( count++ < maxRetries && e !instanceof NoSuchFileException && e !instanceof InterruptedIOException && !Thread.currentThread().isInterrupted() ) {
+                    if( count++ < maxRetries && recoverableError(e) && !Thread.currentThread().isInterrupted() ) {
                         def message = "Unable to stage foreign file: ${filePath.toUriString()} (try ${count} of ${maxRetries}) -- Cause: $e.message"
                         log.isDebugEnabled() ? log.warn(message, e) : log.warn(message)
 
@@ -344,6 +344,12 @@ class FilePorter {
             }
         }
 
+        private boolean recoverableError(Throwable e){
+            return e !instanceof NoSuchFileException
+                && (e instanceof SocketTimeoutException || e !instanceof InterruptedIOException)
+                && e !instanceof SocketException
+        }
+
         private String fmtError(Path filePath, Exception e) {
             def message = "Can't stage file ${FilesEx.toUriString(filePath)}"
             if( e instanceof NoSuchFileException )

modules/nextflow/src/main/groovy/nextflow/processor/TaskArrayCollector.groovy

@@ -92,7 +92,7 @@ class TaskArrayCollector {
         try {
             // submit task directly if the collector is closed
             // or if the task is retried (since it might have dynamic resources)
-            if( closed || task.config.getAttempt() > 1 ) {
+            if( closed ) {
                 executor.submit(task)
                 return
             }
@@ -137,6 +137,9 @@ class TaskArrayCollector {
      * @param tasks
      */
     protected TaskArrayRun createTaskArray(List<TaskRun> tasks) {
+        // mark the task as a child
+        for( TaskRun t : tasks )
+            t.isChild = true
         // prepare child job launcher scripts
         final handlers = tasks.collect( t -> executor.createTaskHandler(t) )
         for( TaskHandler handler : handlers ) {

modules/nextflow/src/main/groovy/nextflow/processor/TaskProcessor.groovy

@@ -2185,7 +2185,7 @@ class TaskProcessor {
         session.filePorter.transfer(batch)
     }
 
-    final protected void makeTaskContextStage3( TaskRun task, HashCode hash, Path folder ) {
+    protected void makeTaskContextStage3( TaskRun task, HashCode hash, Path folder ) {
 
         // set hash-code & working directory
         task.hash = hash
@@ -2340,12 +2340,12 @@ class TaskProcessor {
 
         makeTaskContextStage3(task, hash, folder)
 
+        // when no collector is define OR it's a task retry, then submit directly for execution
+        if( !arrayCollector || task.config.getAttempt() > 1 )
+            executor.submit(task)
         // add the task to the collection of running tasks
-        if( arrayCollector )
-            arrayCollector.collect(task)
         else
-            executor.submit(task)
-
+            arrayCollector.collect(task)
     }
 
     protected boolean checkWhenGuard(TaskRun task) {

modules/nextflow/src/main/groovy/nextflow/processor/TaskRun.groovy

@@ -324,6 +324,11 @@ class TaskRun implements Cloneable {
      */
     volatile boolean aborted
 
+    /**
+     * Mark the task if belongs to a TaskArrayRun
+     */
+    volatile boolean isChild
+
     /**
      * The action {@link ErrorStrategy} action applied if task has failed
      */
@@ -350,6 +355,7 @@ class TaskRun implements Cloneable {
         taskClone.config = config.clone()
         taskClone.config.setContext(taskClone.context)
         taskClone.cache0.clear()
+        taskClone.isChild = false
         return taskClone
     }
 

modules/nextflow/src/test/groovy/nextflow/processor/TaskArrayCollectorTest.groovy

@@ -87,23 +87,6 @@ class TaskArrayCollectorTest extends Specification {
         1 * executor.submit(task)
     }
 
-    def 'should submit retried tasks directly' () {
-        given:
-        def executor = Mock(DummyExecutor)
-        def collector = Spy(new TaskArrayCollector(null, executor, 5))
-        and:
-        def task = Mock(TaskRun) {
-            getConfig() >> Mock(TaskConfig) {
-                getAttempt() >> 2
-            }
-        }
-
-        when:
-        collector.collect(task)
-        then:
-        1 * executor.submit(task)
-    }
-
     def 'should create task array' () {
         given:
         def exec = Mock(DummyExecutor) {

modules/nextflow/src/test/groovy/nextflow/processor/TaskProcessorTest.groovy

@@ -22,6 +22,7 @@ import java.nio.file.Path
 import java.nio.file.Paths
 import java.util.concurrent.ExecutorService
 
+import com.google.common.hash.HashCode
 import groovyx.gpars.agent.Agent
 import nextflow.Global
 import nextflow.ISession
@@ -401,9 +402,7 @@ class TaskProcessorTest extends Specification {
 
     }
 
-
     def 'should update agent state'() {
-
         when:
         def state = new Agent<StateObj>(new StateObj())
         int i = 0
@@ -1172,4 +1171,51 @@ class TaskProcessorTest extends Specification {
 
     }
 
+    def 'should submit a task' () {
+        given:
+        def exec = Mock(Executor)
+        def proc = Spy(new TaskProcessor(executor: exec))
+        and:
+        def task = Mock(TaskRun)
+        def hash = Mock(HashCode)
+        def path = Mock(Path)
+
+        when:
+        proc.submitTask(task, hash, path)
+        then:
+        1 * proc.makeTaskContextStage3(task, hash, path) >> null
+        and:
+        1 * exec.submit(task)
+    }
+
+    def 'should collect a task' () {
+        given:
+        def exec = Mock(Executor)
+        def collector = Mock(TaskArrayCollector)
+        def proc = Spy(new TaskProcessor(executor: exec, arrayCollector: collector))
+        and:
+        def task = Mock(TaskRun)
+        def hash = Mock(HashCode)
+        def path = Mock(Path)
+
+        when:
+        proc.submitTask(task, hash, path)
+        then:
+        task.getConfig()>>Mock(TaskConfig) { getAttempt()>>1 }
+        and:
+        1 * proc.makeTaskContextStage3(task, hash, path) >> null
+        and:
+        1 * collector.collect(task)
+        0 * exec.submit(task)
+
+        when:
+        proc.submitTask(task, hash, path)
+        then:
+        task.getConfig()>>Mock(TaskConfig) { getAttempt()>>2 }
+        and:
+        1 * proc.makeTaskContextStage3(task, hash, path) >> null
+        and:
+        0 * collector.collect(task)
+        1 * exec.submit(task)
+    }
 }

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

@@ -17,6 +17,10 @@
 
 package nextflow.cloud.google.batch
 
+import com.google.api.gax.rpc.NotFoundException
+import com.google.cloud.batch.v1.JobStatus
+import com.google.cloud.batch.v1.Task
+
 import java.nio.file.Path
 import java.util.regex.Pattern
 
@@ -452,49 +456,73 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
      * @return Retrieve the submitted task state
      */
     protected String getTaskState() {
-        final tasks = client.listTasks(jobId)
-        if( !tasks.iterator().hasNext() ) {
-            // if there are no tasks checks the job status
-            return checkJobStatus()
-        }
+        return task.isChild
+            ? getStateFromTaskStatus()
+            : getStateFromJobStatus()
+    }
+
+    protected String getStateFromTaskStatus() {
         final now = System.currentTimeMillis()
         final delta =  now - timestamp;
         if( !taskState || delta >= 1_000) {
-            final status = client.getTaskStatus(jobId, taskId)
-            final newState = status?.state as String
-            if( newState ) {
-                log.trace "[GOOGLE BATCH] Get job=$jobId task=$taskId state=$newState"
-                taskState = newState
-                timestamp = now
-            }
-            if( newState == 'PENDING' ) {
-                final eventsCount = status.getStatusEventsCount()
-                final lastEvent = eventsCount > 0 ? status.getStatusEvents(eventsCount - 1) : null
-                if( lastEvent?.getDescription()?.contains('CODE_GCE_QUOTA_EXCEEDED') )
-                    log.warn1 "Batch job cannot be run: ${lastEvent.getDescription()}"
+            final status = client.getTaskInArrayStatus(jobId, taskId)
+            if( status ) {
+                inspectTaskStatus(status)
+            } else {
+                // If no task status retrieved check job status
+                final jobStatus = client.getJobStatus(jobId)
+                inspectJobStatus(jobStatus)
             }
         }
         return taskState
     }
 
-    protected String checkJobStatus() {
-        final jobStatus = client.getJobStatus(jobId)
-        final newState = jobStatus?.state as String
+    protected String getStateFromJobStatus() {
+        final now = System.currentTimeMillis()
+        final delta =  now - timestamp;
+        if( !taskState || delta >= 1_000) {
+            final status = client.getJobStatus(jobId)
+            inspectJobStatus(status)
+        }
+        return taskState
+    }
+
+    private void inspectTaskStatus(com.google.cloud.batch.v1.TaskStatus status) {
+        final newState = status?.state as String
+        if (newState) {
+            log.trace "[GOOGLE BATCH] Get job=$jobId task=$taskId state=$newState"
+            taskState = newState
+            timestamp = System.currentTimeMillis()
+        }
+        if (newState == 'PENDING') {
+            final eventsCount = status.getStatusEventsCount()
+            final lastEvent = eventsCount > 0 ? status.getStatusEvents(eventsCount - 1) : null
+            if (lastEvent?.getDescription()?.contains('CODE_GCE_QUOTA_EXCEEDED'))
+                log.warn1 "Batch job cannot be run: ${lastEvent.getDescription()}"
+        }
+    }
+
+    protected String inspectJobStatus(JobStatus status) {
+        final newState = status?.state as String
         if (newState) {
+            log.trace "[GOOGLE BATCH] Get job=$jobId state=$newState"
             taskState = newState
             timestamp = System.currentTimeMillis()
             if (newState == "FAILED") {
                 noTaskJobfailure = true
             }
-            return taskState
-        } else {
-            return "PENDING"
+        }
+        if (newState == 'SCHEDULED') {
+            final eventsCount = status.getStatusEventsCount()
+            final lastEvent = eventsCount > 0 ? status.getStatusEvents(eventsCount - 1) : null
+            if (lastEvent?.getDescription()?.contains('CODE_GCE_QUOTA_EXCEEDED'))
+                log.warn1 "Batch job cannot be run: ${lastEvent.getDescription()}"
         }
     }
 
-    static private final List<String> RUNNING_OR_COMPLETED = ['RUNNING', 'SUCCEEDED', 'FAILED']
+    static private final List<String> RUNNING_OR_COMPLETED = ['RUNNING', 'SUCCEEDED', 'FAILED', 'DELETION_IN_PROGRESS']
 
-    static private final List<String> COMPLETED = ['SUCCEEDED', 'FAILED']
+    static private final List<String> COMPLETED = ['SUCCEEDED', 'FAILED', 'DELETION_IN_PROGRESS']
 
     @Override
     boolean checkIfRunning() {
@@ -526,6 +554,8 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
                 task.stderr = errorFile
             }
             status = TaskStatus.COMPLETED
+            if( task.isChild )
+                client.removeFromArrayTasks(jobId, taskId)
             return true
         }
 

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

@@ -50,11 +50,12 @@ import groovy.util.logging.Slf4j
 @Slf4j
 @CompileStatic
 class BatchClient {
-
+    private final static long TASK_STATE_INVALID_TIME = 1_000
     protected String projectId
     protected String location
     protected BatchServiceClient batchServiceClient
     protected BatchConfig config
+    private Map<String, TaskStatusRecord> arrayTaskStatus = new HashMap<String, TaskStatusRecord>()
 
     BatchClient(BatchConfig config) {
         this.config = config
@@ -111,7 +112,7 @@ class BatchClient {
     }
 
     Task describeTask(String jobId, String taskId) {
-        final name = TaskName.of(projectId, location, jobId, 'group0', taskId)
+        final name = generateTaskName(jobId, taskId)
         return apply(()-> batchServiceClient.getTask(name))
     }
 
@@ -141,6 +142,10 @@ class BatchClient {
         return location
     }
 
+    String generateTaskName(String jobId, String taskId) {
+        TaskName.of(projectId, location, jobId, 'group0', taskId)
+    }
+
     /**
      * Creates a retry policy using the configuration specified by {@link BatchRetryConfig}
      *
@@ -194,4 +199,39 @@ class BatchClient {
         // apply the action with
         return Failsafe.with(policy).get(action)
     }
+
+
+    TaskStatus getTaskInArrayStatus(String jobId, String taskId) {
+        final taskName = generateTaskName(jobId,taskId)
+        final now = System.currentTimeMillis()
+        TaskStatusRecord record = arrayTaskStatus.get(taskName)
+        if( !record || now - record.timestamp > TASK_STATE_INVALID_TIME ){
+            log.debug("[GOOGLE BATCH] Updating tasks status for job $jobId")
+            updateArrayTasks(jobId, now)
+            record = arrayTaskStatus.get(taskName)
+        }
+        return record?.status
+    }
+
+    private void updateArrayTasks(String jobId, long now){
+        for( Task t: listTasks(jobId) ){
+            arrayTaskStatus.put(t.name, new TaskStatusRecord(t.status, now))
+        }
+    }
+
+    void removeFromArrayTasks(String jobId, String taskId){
+        final taskName = generateTaskName(jobId,taskId)
+        TaskStatusRecord record = arrayTaskStatus.remove(taskName)
+    }
 }
+
+@CompileStatic
+class TaskStatusRecord {
+    protected TaskStatus status
+    protected long timestamp
+
+    TaskStatusRecord(TaskStatus status, long timestamp) {
+        this.status = status
+        this.timestamp = timestamp
+    }
+}