Add scheduling.provisioningModel per-process hint for Google Batch

docs/executor.mdx

@@ -131,6 +131,16 @@ Resource requests and other job characteristics can be controlled via the follow
 - [resourcelabels][process-resourcelabels]
 - [time][process-time]
 
+The following [hints][process-hints] are supported:
+
+- `scheduling.provisioningModel`: Select the [provisioning model](https://cloud.google.com/batch/docs/create-run-job-spot) for a specific process, overriding the global `google.batch.spot` / `google.batch.preemptible` config options. The value is one of `spot`, `standard` (or `ondemand`) for on-demand, or `preemptible` (case-insensitive). For example:
+
+  ```nextflow
+  hints 'scheduling.provisioningModel': 'standard'
+  ```
+
+  The key may be used as-is or with the `google-batch/` prefix to restrict it to this executor (e.g. `hints 'google-batch/scheduling.provisioningModel': 'spot'`).
+
 See [Cloud Batch][google-batch]for further configuration details.
 
 ## HTCondor

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

@@ -358,6 +358,89 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
         boolean requiresScratchVolume
     }
 
+    private static final String HINT_PREFIX = 'google-batch/'
+    private static final String PROVISIONING_MODEL_HINT = 'scheduling.provisioningModel'
+    private static final Set<String> KNOWN_HINTS = Set.of(PROVISIONING_MODEL_HINT)
+    private static final String SUPPORTED_HINTS_MSG =
+        KNOWN_HINTS.collect { HINT_PREFIX + it }.sort().join(', ')
+
+    /**
+     * Resolve the effective Google Batch {@link AllocationPolicy.ProvisioningModel} for a task.
+     *
+     * The per-process `scheduling.provisioningModel` hint overrides the global `google.batch.spot`
+     * / `google.batch.preemptible` config when set. The executor-prefixed form 
+     * `google-batch/scheduling.provisioningModel` takes precedence over the bare form.
+     *
+     * Hint value is a case-insensitive string naming a Google Batch provisioning model:
+     * `spot`, `standard` (aka `ondemand`), or `preemptible`. Any other value throws an
+     * {@link IllegalArgumentException} (strict validation, no fail-open).
+     *
+     * @param config The task config
+     * @return The effective provisioning model (never {@code null})
+     */
+    protected AllocationPolicy.ProvisioningModel resolveProvisioningModel(TaskConfig config) {
+        final hints = config.getHints()
+        if( !hints )
+            return globalProvisioningModel()
+        // a prefixed hint takes precedence over the unprefixed form
+        final value = hints.get(HINT_PREFIX + PROVISIONING_MODEL_HINT) ?: hints.get(PROVISIONING_MODEL_HINT)
+        if( value == null )
+            return globalProvisioningModel()
+        if( !(value instanceof CharSequence) )
+            throw new IllegalArgumentException("Invalid '${PROVISIONING_MODEL_HINT}' hint value: $value -- it must be one of: spot, standard (aka ondemand), preemptible")
+        final str = value.toString().trim().toLowerCase()
+        switch( str ) {
+            case 'spot':
+                return AllocationPolicy.ProvisioningModel.SPOT
+            case 'standard':
+            case 'ondemand':
+                // `standard` is Google's canonical enum term; `ondemand` is accepted as an
+                // alias and resolves to the same STANDARD model.
+                return AllocationPolicy.ProvisioningModel.STANDARD
+            case 'preemptible':
+                return AllocationPolicy.ProvisioningModel.PREEMPTIBLE
+            default:
+                throw new IllegalArgumentException("Invalid '${PROVISIONING_MODEL_HINT}' hint value: '$value' -- it must be one of: spot, standard (aka ondemand), preemptible")
+        }
+    }
+
+    /**
+     * The provisioning model implied by the global `google.batch` config. Spot takes precedence
+     * over preemptible (matching the previous ordering), and STANDARD (on-demand) is the default
+     * when neither is enabled.
+     *
+     * @return The effective provisioning model
+     */
+    private AllocationPolicy.ProvisioningModel globalProvisioningModel() {
+        if( batchConfig.spot )
+            return AllocationPolicy.ProvisioningModel.SPOT
+        if( batchConfig.preemptible )
+            return AllocationPolicy.ProvisioningModel.PREEMPTIBLE
+        return AllocationPolicy.ProvisioningModel.STANDARD
+    }
+
+    /**
+     * Validate `google-batch/`-prefixed hints against the set of hints supported by the Google
+     * Batch executor. Per the hints ADR, an unrecognized *prefixed* hint is a hard error, so this method
+     * throws {@link IllegalArgumentException}. Only hints explicitly routed to this executor are validated
+     * here.
+     *
+     * @param hints the full 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}")
+    }
+
     /**
      * Build the instance policy or template for job allocation.
      * Note: This method sets machineInfo field as a side effect.
@@ -393,6 +476,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(PROVISIONING_MODEL_HINT) || hints?.containsKey(HINT_PREFIX + PROVISIONING_MODEL_HINT) )
+                log.warn1 "Google Batch hint '${PROVISIONING_MODEL_HINT}' ignored because an instance template was specified"
+
             instancePolicyOrTemplate
                 .setInstallGpuDrivers(batchConfig.getInstallGpuDrivers())
                 .setInstanceTemplate(task.config.getMachineType().minus('template://'))
@@ -479,11 +566,10 @@ 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)
+            // resolve the effective provisioning model
+            final provisioningModel = resolveProvisioningModel(task.config)
+            if( provisioningModel != AllocationPolicy.ProvisioningModel.STANDARD )
+                instancePolicy.setProvisioningModel(provisioningModel)
 
             instancePolicyOrTemplate.setPolicy(instancePolicy)
         }
@@ -553,6 +639,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 +1011,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
@@ -362,6 +363,53 @@ class GoogleBatchTaskHandlerTest extends Specification {
         taskGroup.getTaskSpec().getVolumes(0) == GCS_VOL
     }
 
+    def 'should use instance template and ignore the provisioningModel hint' () {
+        given:
+        def GCS_VOL = Volume.newBuilder().setGcs(GCS.newBuilder().setRemotePath('foo').build() ).build()
+        def WORK_DIR = CloudStorageFileSystem.forBucket('foo').getPath('/scratch')
+        def CONTAINER_IMAGE = 'debian:latest'
+        def INSTANCE_TEMPLATE = 'instance-template'
+        def exec = Mock(GoogleBatchExecutor) {
+            getBatchConfig() >> Mock(BatchConfig) {
+                getInstallGpuDrivers() >> true
+            }
+        }
+        and:
+        def bean = new TaskBean(workDir: WORK_DIR, inputFiles: [:])
+        def task = Mock(TaskRun) {
+            toTaskBean() >> bean
+            getHashLog() >> 'abcd1234'
+            getWorkDir() >> WORK_DIR
+            getContainer() >> CONTAINER_IMAGE
+            getConfig() >> Mock(TaskConfig) {
+                getCpus() >> 2
+                getMachineType() >> "template://${INSTANCE_TEMPLATE}"
+                getResourceLabels() >> [:]
+                getHints() >> ['scheduling.provisioningModel': 'spot']
+            }
+        }
+        and:
+        def mounts = ['/mnt/disks/foo/scratch:/mnt/disks/foo/scratch:rw']
+        def volumes = [GCS_VOL]
+        def launcher = new GoogleBatchLauncherSpecMock('bash .command.run', mounts, volumes)
+
+        and:
+        def handler = Spy(new GoogleBatchTaskHandler(task, exec))
+
+        when:
+        def req = handler.newSubmitRequest(task, launcher)
+        then:
+        handler.fusionEnabled() >> false
+        handler.findBestMachineType(_, false) >> null
+
+        and:
+        // the template is honored wholesale; the provisioningModel hint is silently ignored (only warned)
+        def allocationPolicy = req.getAllocationPolicy()
+        def instancePolicyOrTemplate = allocationPolicy.getInstances(0)
+        instancePolicyOrTemplate.getInstanceTemplate() == INSTANCE_TEMPLATE
+        instancePolicyOrTemplate.getInstallGpuDrivers() == true
+    }
+
     def 'should create the trace record' () {
         given:
         def exec = Mock(Executor) { getName() >> 'google-batch' }
@@ -910,6 +958,118 @@ class GoogleBatchTaskHandlerTest extends Specification {
         2        | true   | true      | 2
     }
 
+    def 'should resolve provisioning model from scheduling.provisioningModel hint' () {
+        given:
+        def task = Mock(TaskRun) { hashLog >> '1234567890'; getWorkDir() >> Path.of('/work/dir') }
+        def exec = Mock(GoogleBatchExecutor) {
+            getClient() >> Mock(BatchClient)
+            getBatchConfig() >> Mock(BatchConfig) {
+                getSpot() >> CONFIG_SPOT; isSpot() >> CONFIG_SPOT
+                getPreemptible() >> CONFIG_PREEMPTIBLE; isPreemptible() >> CONFIG_PREEMPTIBLE
+            }
+        }
+        def handler = Spy(new GoogleBatchTaskHandler(task, exec))
+        def config = Mock(TaskConfig) { getHints() >> HINTS }
+
+        expect:
+        handler.resolveProvisioningModel(config) == AllocationPolicy.ProvisioningModel.valueOf(EXPECTED)
+
+        where:
+        HINTS                                                              | CONFIG_SPOT | CONFIG_PREEMPTIBLE | EXPECTED
+        [:]                                                               | false       | false              | 'STANDARD'     // unset -> global default (on-demand)
+        [:]                                                               | true        | false              | 'SPOT'         // unset -> global spot
+        [:]                                                               | false       | true               | 'PREEMPTIBLE'  // unset -> global preemptible
+        [:]                                                               | true        | true               | 'SPOT'         // spot wins over preemptible
+        ['scheduling.provisioningModel': 'spot']                          | false       | false              | 'SPOT'         // bare hint enables spot
+        ['scheduling.provisioningModel': 'standard']                      | true        | false              | 'STANDARD'     // bare hint overrides global spot -> on-demand
+        ['scheduling.provisioningModel': 'preemptible']                   | false       | false              | 'PREEMPTIBLE'  // bare hint preemptible
+        ['google-batch/scheduling.provisioningModel': 'spot']             | false       | false              | 'SPOT'         // prefixed hint enables spot
+        ['google-batch/scheduling.provisioningModel': 'standard']         | true        | false              | 'STANDARD'     // prefixed hint overrides global spot
+        ['google-batch/scheduling.provisioningModel': 'standard',
+         'scheduling.provisioningModel': 'spot']                          | true        | false              | 'STANDARD'     // prefixed wins over bare
+        ['scheduling.provisioningModel': 'SPOT']                          | false       | false              | 'SPOT'         // value is case-insensitive
+        ['scheduling.provisioningModel': ' Preemptible ']                 | false       | false              | 'PREEMPTIBLE'  // value is trimmed
+        ['scheduling.provisioningModel': 'ondemand']                      | false       | false              | 'STANDARD'     // `ondemand` is an alias for `standard` (on-demand)
+        ['scheduling.provisioningModel': 'OnDemand']                      | true        | false              | 'STANDARD'     // `ondemand` alias is case-insensitive, overrides global spot
+    }
+
+    def 'should fail on invalid scheduling.provisioningModel hint value' () {
+        given:
+        def task = Mock(TaskRun) { hashLog >> '1234567890'; getWorkDir() >> Path.of('/work/dir') }
+        def exec = Mock(GoogleBatchExecutor) {
+            getClient() >> Mock(BatchClient)
+            getBatchConfig() >> Mock(BatchConfig) { getSpot() >> false; isSpot() >> false; getPreemptible() >> false; isPreemptible() >> false }
+        }
+        def handler = Spy(new GoogleBatchTaskHandler(task, exec))
+        def config = Mock(TaskConfig) { getHints() >> ['scheduling.provisioningModel': BAD_VALUE] }
+
+        when:
+        handler.resolveProvisioningModel(config)
+        then:
+        thrown(IllegalArgumentException)
+
+        where:
+        BAD_VALUE << [123, true, '', 'yes', '1', 'spotty', 'STANDARD_ISH']   // wrong type and bogus strings must throw, not fail open ('ondemand' is now a valid alias)
+    }
+
+    def 'should throw on unknown google-batch prefixed hint' () {
+        given:
+        def task = Mock(TaskRun) { hashLog >> '1234567890'; getWorkDir() >> Path.of('/work/dir') }
+        def exec = Mock(GoogleBatchExecutor) {
+            getClient() >> Mock(BatchClient)
+            getBatchConfig() >> Mock(BatchConfig) { getSpot() >> false; isSpot() >> false; getPreemptible() >> false; isPreemptible() >> false }
+        }
+        def handler = Spy(new GoogleBatchTaskHandler(task, exec))
+
+        when:
+        handler.validateHints(HINTS)
+        then:
+        thrown(IllegalArgumentException)
+
+        where:
+        HINTS                                                        | _
+        ['google-batch/scheduling.provisioningmodel': 'spot']       | _   // realistic case-typo of a known key
+        ['google-batch/bogus': 'x']                                 | _   // wholly unknown prefixed key
+    }
+
+    def 'should report all unknown google-batch hints together' () {
+        given:
+        def task = Mock(TaskRun) { hashLog >> '1234567890'; getWorkDir() >> Path.of('/work/dir') }
+        def exec = Mock(GoogleBatchExecutor) {
+            getClient() >> Mock(BatchClient)
+            getBatchConfig() >> Mock(BatchConfig) { getSpot() >> false; isSpot() >> false; getPreemptible() >> false; isPreemptible() >> false }
+        }
+        def handler = Spy(new GoogleBatchTaskHandler(task, exec))
+
+        when:
+        handler.validateHints(['google-batch/bogus': 'x', 'google-batch/scheduling.provisioningmodel': 'spot'])
+        then:
+        def e = thrown(IllegalArgumentException)
+        e.message.contains('google-batch/bogus')
+        e.message.contains('google-batch/scheduling.provisioningmodel')
+    }
+
+    def 'should ignore unknown bare and foreign-executor hints' () {
+        given:
+        def task = Mock(TaskRun) { hashLog >> '1234567890'; getWorkDir() >> Path.of('/work/dir') }
+        def exec = Mock(GoogleBatchExecutor) {
+            getClient() >> Mock(BatchClient)
+            getBatchConfig() >> Mock(BatchConfig) { getSpot() >> false; isSpot() >> false; getPreemptible() >> false; isPreemptible() >> false }
+        }
+        def handler = Spy(new GoogleBatchTaskHandler(task, exec))
+
+        when:
+        // only google-batch/-prefixed keys are validated; bare and foreign-executor keys are left untouched
+        handler.validateHints(HINTS)
+        then:
+        notThrown(IllegalArgumentException)
+
+        where:
+        HINTS                                          | _
+        ['scheduling.bogus': 'x']                      | _   // bare unknown key -> ignored here (owned by no specific executor)
+        ['awsbatch/consumableResources': [:]]          | _   // foreign-executor prefixed key -> not ours to validate
+    }
+
     def 'should return zero when no status events exist'() {
         given:
         def handler = Spy(GoogleBatchTaskHandler)