Fix AWS Batch task polling forever when the job is missing from the DescribeJobs response

docs/reference/config.mdx

@@ -931,9 +931,9 @@ Determines how often to log the executor status (default`5 min`).
 
 ##### `executor.exitReadTimeout`
 
-*Used only by grid executors.*
+*Used only by grid executors and AWS Batch.*
 
-Determines how long to wait for the `.exitcode` file to be created after the task has completed, before returning an error status (default`270 sec`).
+Determines how long to wait for the `.exitcode` file to be created after the task has completed, before returning an error status (default`270 sec`). For AWS Batch, it determines how long to wait when the job status cannot be resolved via the Batch API (e.g. the job is missing from the `DescribeJobs` response), before returning an error status.
 
 ##### `executor.jobName`
 

plugins/nf-amazon/src/main/nextflow/cloud/aws/batch/AwsBatchTaskHandler.groovy

@@ -133,6 +133,22 @@ class AwsBatchTaskHandler extends TaskHandler implements BatchHandler<String,Job
      */
     private BatchContext<String,JobDetail> context
 
+    /**
+     * Timestamp of the first poll of the current streak of polls for which the
+     * job status could not be resolved i.e. {@link #describeJob(String)}
+     * returned {@code null}; zero while the status is resolvable. Used to bound
+     * an unresolvable job status with {@link #exitReadTimeoutMillis} instead of
+     * polling forever -- see {@link #checkIfCompleted()}
+     */
+    private long unresolvedTimestampMillis
+
+    /**
+     * Max time to wait for the job status to become resolvable again before
+     * failing the task, from the `executor.exitReadTimeout` config option --
+     * the same setting bounding the analogous condition in the grid executors
+     */
+    private long exitReadTimeoutMillis
+
     @TestOnly
     protected AwsBatchTaskHandler() {}
 
@@ -155,6 +171,7 @@ class AwsBatchTaskHandler extends TaskHandler implements BatchHandler<String,Job
         this.exitFile = task.workDir.resolve(TaskRun.CMD_EXIT)
         this.wrapperFile = task.workDir.resolve(TaskRun.CMD_RUN)
         this.traceFile = task.workDir.resolve(TaskRun.CMD_TRACE)
+        this.exitReadTimeoutMillis = executor.config.getExitReadTimeout(executor.name).toMillis()
     }
 
     protected String getJobId() { jobId }
@@ -304,7 +321,28 @@ class AwsBatchTaskHandler extends TaskHandler implements BatchHandler<String,Job
         if( !isRunning() )
             return false
         final job = describeJob(jobId)
-        final done = job?.status() in [JobStatus.SUCCEEDED, JobStatus.FAILED]
+        if( job == null ) {
+            // the job status cannot be resolved -- either the describe request failed or
+            // the job is missing from the response (e.g. it aged out of the Batch job
+            // history, which retains jobs for ~24 hours). Bound this condition with
+            // `exitReadTimeout`, symmetrically with the grid executors, instead of
+            // polling forever -- see https://github.com/nextflow-io/nextflow/issues/7301
+            final now = System.currentTimeMillis()
+            if( !unresolvedTimestampMillis )
+                unresolvedTimestampMillis = now
+            final delta = now - unresolvedTimestampMillis
+            if( delta < exitReadTimeoutMillis )
+                return false
+            log.debug "[AWS BATCH] Failed to get status for job=$jobId -- exitReadTimeoutMillis: $exitReadTimeoutMillis; delta: $delta"
+            task.exitStatus = Integer.MAX_VALUE
+            task.error = new ProcessException("Task '${task.lazyName()}' status could not be resolved -- the job was missing from the DescribeJobs response for longer than `exitReadTimeout` ($exitReadTimeoutMillis ms). Likely it has been terminated by the external system")
+            task.stdout = outputFile
+            task.stderr = errorFile
+            status = TaskStatus.COMPLETED
+            return true
+        }
+        unresolvedTimestampMillis = 0
+        final done = job.status() in [JobStatus.SUCCEEDED, JobStatus.FAILED]
         if( done ) {
             // take the exit code of the container, if missing (null)
             // take the exit code from the `.exitcode` file create by nextflow

plugins/nf-amazon/src/test/nextflow/cloud/aws/batch/AwsBatchTaskHandlerTest.groovy

@@ -28,8 +28,10 @@ import nextflow.cloud.aws.config.AwsConfig
 import nextflow.cloud.aws.util.S3PathFactory
 import nextflow.cloud.types.CloudMachineInfo
 import nextflow.cloud.types.PriceModel
+import nextflow.exception.ProcessException
 import nextflow.exception.ProcessUnrecoverableException
 import nextflow.executor.Executor
+import nextflow.executor.ExecutorConfig
 import nextflow.fusion.FusionScriptLauncher
 import nextflow.processor.Architecture
 import nextflow.processor.BatchContext
@@ -1237,7 +1239,7 @@ class AwsBatchTaskHandlerTest extends Specification {
         def session = Mock(Session) {getConfig() >> config }
         def opts = new AwsOptions(session)
         and:
-        def executor = Mock(AwsBatchExecutor) { getAwsOptions() >> opts; isFusionEnabled()>>FUSION }
+        def executor = Mock(AwsBatchExecutor) { getAwsOptions() >> opts; isFusionEnabled()>>FUSION; getConfig() >> new ExecutorConfig([:]) }
         def proc = Mock(TaskProcessor) { getExecutor()>>executor }
         def task = Mock(TaskRun) {workDir>> Path.of('/foo'); getProcessor()>>proc }
         def handler = new AwsBatchTaskHandler(task, executor)
@@ -1365,6 +1367,67 @@ class AwsBatchTaskHandlerTest extends Specification {
 
     }
 
+    def 'should not fail the task while the job status is unresolvable within the exit read timeout'() {
+        given:
+        def task = new TaskRun()
+        task.name = 'hello'
+        def jobId = 'job-123'
+        def handler = Spy(new AwsBatchTaskHandler(task: task, jobId: jobId, status: TaskStatus.RUNNING, exitReadTimeoutMillis: 3_600_000))
+
+        when:
+        def result = handler.checkIfCompleted()
+        then:
+        1 * handler.describeJob(jobId) >> null
+        and:
+        result == false
+        handler.status == TaskStatus.RUNNING
+        handler.@unresolvedTimestampMillis > 0
+    }
+
+    def 'should fail the task when the job status is unresolvable for longer than the exit read timeout'() {
+        given:
+        def task = new TaskRun()
+        task.name = 'hello'
+        def jobId = 'job-123'
+        def handler = Spy(new AwsBatchTaskHandler(task: task, jobId: jobId, status: TaskStatus.RUNNING, exitReadTimeoutMillis: 200))
+
+        when:
+        def result1 = handler.checkIfCompleted()
+        // wait longer than the timeout defined by `exitReadTimeoutMillis`
+        sleep 300
+        def result2 = handler.checkIfCompleted()
+        then:
+        2 * handler.describeJob(jobId) >> null
+        and:
+        result1 == false
+        result2 == true
+        handler.status == TaskStatus.COMPLETED
+        handler.task.exitStatus == Integer.MAX_VALUE
+        handler.task.error instanceof ProcessException
+        handler.task.error.message.contains('exitReadTimeout')
+    }
+
+    def 'should reset the unresolved status window when the job status becomes resolvable again'() {
+        given:
+        def task = new TaskRun()
+        task.name = 'hello'
+        def jobId = 'job-123'
+        def handler = Spy(new AwsBatchTaskHandler(task: task, jobId: jobId, status: TaskStatus.RUNNING, exitReadTimeoutMillis: 3_600_000))
+        and:
+        def job = JobDetail.builder().status(JobStatus.RUNNING).build()
+
+        when:
+        def result1 = handler.checkIfCompleted()
+        def result2 = handler.checkIfCompleted()
+        then:
+        2 * handler.describeJob(jobId) >>> [null, job]
+        and:
+        result1 == false
+        result2 == false
+        handler.status == TaskStatus.RUNNING
+        handler.@unresolvedTimestampMillis == 0L
+    }
+
     def 'should return zero spot interruptions when no attempts or non-spot terminations exist'() {
         given:
         def handler = Spy(AwsBatchTaskHandler)