Add per-process Spot to on-demand fallback hint for Google Batch

docs/executor/google-batch.mdx

@@ -21,15 +21,36 @@ Use the following process directives to control resource requests and other job
 - [resourcelabels][process-resourcelabels]
 - [time][process-time]
 
+The following [hints][process-hints] are supported:
+
+- `scheduling.spotAttempts`: Run a process on Spot for its first N attempts, then automatically fall back to on-demand (`STANDARD`) for later attempts. The value is a positive integer N: attempts `1` to `N` request a Spot VM, attempts after `N` request on-demand. The attempt count is the greater of the task's execution attempt and its submit attempt, so the fallback is triggered both by a mid-run Spot reclaim and by a failure to obtain a Spot VM (the latter in combination with the [`maxSubmitAwait`][process-maxsubmitawait] directive). Requires `errorStrategy 'retry'` with `maxRetries` set high enough to reach the on-demand attempts. For example:
+
+  ```nextflow
+  process EXAMPLE {
+      errorStrategy 'retry'
+      maxRetries 3
+      hints 'scheduling.spotAttempts': '2'   // Spot for attempts 1 and 2, on-demand from attempt 3
+  }
+  ```
+
+  This key may be used as-is or with the `google-batch/` prefix to restrict it to this executor.
+
+  :::note
+  `scheduling.spotAttempts` is distinct from the [`google.batch.maxSpotAttempts`][config-maxspotattempts] config option, despite the similar name. `google.batch.maxSpotAttempts` sets the number of times Cloud Batch retries a job *internally* after a Spot reclaim, and each retry requests a Spot VM again; it applies to every process in the run. `scheduling.spotAttempts` instead changes the provisioning model across Nextflow-level attempts, and is set per process. The two can be combined.
+  :::
+
 See [Cloud Batch][google-batch] for further configuration details.
 
+[config-maxspotattempts]: ../reference/config/google#googlebatchmaxspotattempts
 [google-batch]: ../google#cloud-batch
 [process-accelerator]: ../reference/process#accelerator
 [process-container]: ../reference/process#container
 [process-containeroptions]: ../reference/process#containeroptions
 [process-cpus]: ../reference/process#cpus
 [process-disk]: ../reference/process#disk
+[process-hints]: ../reference/process#hints
 [process-machinetype]: ../reference/process#machinetype
+[process-maxsubmitawait]: ../reference/process#maxsubmitawait
 [process-memory]: ../reference/process#memory
 [process-resourcelabels]: ../reference/process#resourcelabels
 [process-time]: ../reference/process#time

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

@@ -358,6 +358,77 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
         boolean requiresScratchVolume
     }
 
+    private static final String HINT_PREFIX = 'google-batch/'
+    private static final String SPOT_ATTEMPTS_HINT = 'scheduling.spotAttempts'
+    private static final Set<String> KNOWN_HINTS = Set.of(SPOT_ATTEMPTS_HINT)
+    private static final String SUPPORTED_HINTS_MSG =
+        KNOWN_HINTS.collect { HINT_PREFIX + it }.sort().join(', ')
+
+    /**
+     * Validate the google-batch prefixed hints, throwing if a prefixed key is not supported
+     *
+     * @param hints The hints map from the task config
+     */
+    protected void validateHints(Map<String,Object> hints) {
+        if( !hints )
+            return
+        final unknown = []
+        for( final key : hints.keySet() ) {
+            if( !key?.startsWith(HINT_PREFIX) )
+                continue
+            if( !KNOWN_HINTS.contains(key.substring(HINT_PREFIX.length())) )
+                unknown.add(key)
+        }
+        if( unknown )
+            throw new IllegalArgumentException("Unknown Google Batch hint(s): ${unknown.collect { "'$it'" }.join(', ')} -- supported keys are: ${SUPPORTED_HINTS_MSG}")
+    }
+
+    /**
+     * Resolve the effective provisioning model for a task, honouring the scheduling.spotAttempts hint
+     *
+     * @param config The task config
+     * @return The effective provisioning model
+     */
+    protected AllocationPolicy.ProvisioningModel resolveProvisioningModel(TaskConfig config) {
+        final spotAttempts = resolveSpotAttempts(config)
+        if( spotAttempts != null ) {
+            // run on spot for the first N attempts, then fall back to on-demand (STANDARD).
+            // `attempt` covers a mid-run reclaim; `submitAttempt` covers a failure to obtain a VM
+            // (via `maxSubmitAwait`), so the larger of the two drives the escalation
+            final attempt = Math.max(config.getAttempt(), config.getSubmitAttempt())
+            return attempt <= spotAttempts
+                ? AllocationPolicy.ProvisioningModel.SPOT
+                : AllocationPolicy.ProvisioningModel.STANDARD
+        }
+        // otherwise use the global config: spot wins over preemptible, else on-demand
+        if( batchConfig.spot )
+            return AllocationPolicy.ProvisioningModel.SPOT
+        if( batchConfig.preemptible )
+            return AllocationPolicy.ProvisioningModel.PREEMPTIBLE
+        return AllocationPolicy.ProvisioningModel.STANDARD
+    }
+
+    /**
+     * The scheduling.spotAttempts hint value, or null when unset; throws on a non-positive integer
+     *
+     * @param config The task config
+     * @return The number of leading attempts to run on spot, or null
+     */
+    protected Integer resolveSpotAttempts(TaskConfig config) {
+        final hints = config?.getHints()
+        if( !hints )
+            return null
+        final value = hints.get(HINT_PREFIX + SPOT_ATTEMPTS_HINT) ?: hints.get(SPOT_ATTEMPTS_HINT)
+        if( value == null )
+            return null
+        final n = value instanceof Number
+            ? (value as Integer)
+            : (value.toString().trim().isInteger() ? value.toString().trim() as Integer : null)
+        if( n == null || n < 1 )
+            throw new IllegalArgumentException("Invalid '${SPOT_ATTEMPTS_HINT}' hint value: '${value}' -- it must be a positive integer")
+        return n
+    }
+
     /**
      * Build the instance policy or template for job allocation.
      * Note: This method sets machineInfo field as a side effect.
@@ -393,6 +464,10 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
             if( batchConfig.spot )
                 log.warn1 'Config option `google.batch.spot` ignored because an instance template was specified'
 
+            final hints = task.config.getHints()
+            if( hints?.containsKey(SPOT_ATTEMPTS_HINT) || hints?.containsKey(HINT_PREFIX + SPOT_ATTEMPTS_HINT) )
+                log.warn1 "Google Batch hint '${SPOT_ATTEMPTS_HINT}' ignored because an instance template was specified"
+
             instancePolicyOrTemplate
                 .setInstallGpuDrivers(batchConfig.getInstallGpuDrivers())
                 .setInstanceTemplate(task.config.getMachineType().minus('template://'))
@@ -479,11 +554,9 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
             if( batchConfig.cpuPlatform )
                 instancePolicy.setMinCpuPlatform(batchConfig.cpuPlatform)
 
-            if( batchConfig.preemptible )
-                instancePolicy.setProvisioningModel(AllocationPolicy.ProvisioningModel.PREEMPTIBLE)
-
-            if( batchConfig.spot )
-                instancePolicy.setProvisioningModel(AllocationPolicy.ProvisioningModel.SPOT)
+            final provisioningModel = resolveProvisioningModel(task.config)
+            if( provisioningModel != AllocationPolicy.ProvisioningModel.STANDARD )
+                instancePolicy.setProvisioningModel(provisioningModel)
 
             instancePolicyOrTemplate.setPolicy(instancePolicy)
         }
@@ -553,6 +626,9 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
     }
 
     protected Job newSubmitRequest(TaskRun task, GoogleBatchLauncherSpec launcher) {
+        // validate the task's hints once per submission
+        validateHints(task.config.getHints())
+
         // container validation
         if( !task.container )
             throw new ProcessUnrecoverableException("Process `${task.lazyName()}` failed because the container image was not specified")
@@ -922,7 +998,10 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
         final location = client.location
         final cpus = config.getCpus()
         final memory = config.getMemory() ? config.getMemory().toMega().toInteger() : 1024
-        final spot = batchConfig.spot ?: batchConfig.preemptible
+        // price as spot when the effective provisioning model is SPOT or PREEMPTIBLE
+        final provisioningModel = resolveProvisioningModel(config)
+        final spot = provisioningModel == AllocationPolicy.ProvisioningModel.SPOT \
+                  || provisioningModel == AllocationPolicy.ProvisioningModel.PREEMPTIBLE
         final machineType = config.getMachineType()
         final families = machineType ? machineType.tokenize(',') : List.<String>of()
         final priceModel = spot ? PriceModel.spot : PriceModel.standard

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

@@ -27,6 +27,7 @@ import nextflow.exception.ProcessException
 
 import java.nio.file.Path
 
+import com.google.cloud.batch.v1.AllocationPolicy
 import com.google.cloud.batch.v1.GCS
 import com.google.cloud.batch.v1.StatusEvent
 import com.google.cloud.batch.v1.TaskStatus
@@ -1377,4 +1378,89 @@ class GoogleBatchTaskHandlerTest extends Specification {
         'h?-standard-88'      | 'local-ssd'          // 'h?' does not match regex ^h3-.*$, so not classified as h3
     }
 
+    // provisioning model resolution: scheduling.spotAttempts spot -> on-demand fallback
+
+    private GoogleBatchTaskHandler provisioningHandler(boolean useSpot = false, boolean usePreemptible = false) {
+        def task = Mock(TaskRun) { hashLog >> '1234567890'; getWorkDir() >> Path.of('/work/dir') }
+        def cfg = Mock(BatchConfig) {
+            getSpot() >> useSpot; isSpot() >> useSpot
+            getPreemptible() >> usePreemptible; isPreemptible() >> usePreemptible
+        }
+        def exec = Mock(GoogleBatchExecutor) { getClient() >> Mock(BatchClient); getBatchConfig() >> cfg }
+        return Spy(new GoogleBatchTaskHandler(task, exec))
+    }
+
+    def 'should resolve provisioning model from spotAttempts hint' () {
+        given:
+        def handler = provisioningHandler(GLOBAL_SPOT, GLOBAL_PREEMPT)
+        def config = Mock(TaskConfig) { getHints() >> HINTS; getAttempt() >> ATTEMPT; getSubmitAttempt() >> SUBMIT }
+
+        expect:
+        handler.resolveProvisioningModel(config) == AllocationPolicy.ProvisioningModel.valueOf(EXPECTED)
+
+        where:
+        HINTS                                         | ATTEMPT | SUBMIT | GLOBAL_SPOT | GLOBAL_PREEMPT || EXPECTED
+        ['scheduling.spotAttempts': '2']              | 1       | 1      | false       | false          || 'SPOT'
+        ['scheduling.spotAttempts': '2']              | 2       | 1      | false       | false          || 'SPOT'
+        ['scheduling.spotAttempts': '2']              | 3       | 1      | false       | false          || 'STANDARD'
+        ['scheduling.spotAttempts': '2']              | 1       | 3      | false       | false          || 'STANDARD'   // submitAttempt (stockout) drives the fallback
+        ['scheduling.spotAttempts': '2']              | 1       | 2      | false       | false          || 'SPOT'
+        ['scheduling.spotAttempts': 2]                | 1       | 1      | false       | false          || 'SPOT'
+        ['google-batch/scheduling.spotAttempts': '1'] | 2       | 1      | false       | false          || 'STANDARD'
+        [:]                                           | 1       | 1      | true        | false          || 'SPOT'        // global fallback, unchanged
+        [:]                                           | 1       | 1      | false       | true           || 'PREEMPTIBLE'
+        [:]                                           | 1       | 1      | false       | false          || 'STANDARD'
+        [:]                                           | 1       | 1      | true        | true           || 'SPOT'        // spot wins over preemptible
+    }
+
+    def 'should throw on invalid spotAttempts hint' () {
+        given:
+        def handler = provisioningHandler()
+        def config = Mock(TaskConfig) { getHints() >> HINTS }
+
+        when:
+        handler.resolveSpotAttempts(config)
+
+        then:
+        thrown(IllegalArgumentException)
+
+        where:
+        HINTS << [
+            ['scheduling.spotAttempts': 'lots'],
+            ['scheduling.spotAttempts': '0'],
+            ['scheduling.spotAttempts': -1],
+        ]
+    }
+
+    def 'should accept known and unprefixed google-batch hints' () {
+        given:
+        def handler = provisioningHandler()
+
+        when:
+        handler.validateHints(HINTS)
+
+        then:
+        noExceptionThrown()
+
+        where:
+        HINTS << [
+            null,
+            [:],
+            ['google-batch/scheduling.spotAttempts': '2'],
+            ['scheduling.spotAttempts': '2'],    // bare (unprefixed) keys are not validated
+            ['other-executor/foo': 'bar'],       // foreign-executor keys are left untouched
+        ]
+    }
+
+    def 'should reject unknown prefixed google-batch hints' () {
+        given:
+        def handler = provisioningHandler()
+
+        when:
+        handler.validateHints(['google-batch/scheduling.unknown': 'x'])
+
+        then:
+        thrown(IllegalArgumentException)
+    }
+
 }