Feature Request: add option to change default scheduling policy with nf-google

New feature

A configurable flag within the nf-google plugin to change the scheduling policy between AS_SOON_AS_POSSIBLE and IN_ORDER - see docs here

Use case

The default scheduling policy is AS_SOON_AS_POSSIBLE, which attempts to run all tasks in an array job simultaneously. Workflows with high resource requirements can result in excessive resource requests that exceed GCP quotas and cause job failures.

Suggested implementation

I guess there are a few ways of setting this. I made a small update so the scheduling policy can now be set per process via resourceLabels, with options of either AS_SOON_AS_POSSIBLE or IN_ORDER. Using IN_ORDER will force tasks within each array job to execute sequentially rather than in parallel.

For example:

process {
    withName: EXAMPLE_PROCESS {
        resourceLabels = ['scheduling-policy': 'in-order']
    }
}

// or

process EXAMPLE_PROCESS {
    cpus 1
    memory 1.GB
    maxForks 10
    array 2
    resourceLabels 'scheduling-policy': 'in-order'
}

Code change

file GoogleBatchTaskHandler.groovy

// task group
final taskGroup = TaskGroup.newBuilder()
   .setTaskSpec(taskSpec)

/*
* Added modification to allow changes to scheduling policy
*/

// Read scheduling policy from resourceLabels
def labels = task.config.getResourceLabels()
def schedulingPolicyLabel = labels?.get('scheduling-policy')?.toString()?.toLowerCase()?.replace('-', '_')

if( task instanceof TaskArrayRun ) {
   final arraySize = task.getArraySize()
   taskGroup.setTaskCount(arraySize)

   if( schedulingPolicyLabel == 'in_order' ) {
      log.info "[GOOGLE BATCH] Process `${task.lazyName()}` using IN_ORDER scheduling (parallelism forced to 1)"
      taskGroup.setSchedulingPolicy(TaskGroup.SchedulingPolicy.IN_ORDER)
      taskGroup.setParallelism(1)
   }
}

// create the job
return Job.newBuilder()
   .addTaskGroups(taskGroup)
   .setAllocationPolicy(allocationPolicy)
   .setLogsPolicy(createLogsPolicy())
   .putAllLabels(task.config.getResourceLabels())
   .build()

Start in GoogleBatchTaskHandler.groovy and read the task-group construction, especially how TaskArrayRun jobs are configured and how resourceLabels are passed through. Compare the supported Google Batch scheduling policies in the linked documentation. Done means the configured policy is accepted for array jobs without changing the existing default behavior.