diff --git a/docs/executor/google-batch.mdx b/docs/executor/google-batch.mdx index b5e0b18658..6dc5c93c35 100644 --- a/docs/executor/google-batch.mdx +++ b/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 diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy index 792ec949d5..28fd90c647 100644 --- a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy +++ b/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 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 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.of() final priceModel = spot ? PriceModel.spot : PriceModel.standard diff --git a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy index fbe88ab060..21e0e5456c 100644 --- a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchTaskHandlerTest.groovy +++ b/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) + } + }