Implement image streaming for Google Batch

docs/reference/config.md

@@ -971,6 +971,15 @@ The following settings are available:
 : The Google Cloud Storage path where job logs should be stored, e.g. `gs://my-logs-bucket/logs`.
 : When specified, Google Batch will write job logs to this location instead of [Cloud Logging](https://cloud.google.com/logging/docs). The bucket must be accessible and writable by the service account.
 
+`google.batch.enableImageStreaming`
+: :::{versionadded} 26.02.1-edge
+  :::
+: Enable container image streaming to speed up job start-up (default: `false`). This can reduce latency for large images but comes with some [limitations](https://cloud.google.com/batch/docs/use-image-streaming), including:
+
+  - `process.containerOptions` is ignored when image streaming is enabled.
+  - Image streaming only supports container images stored in Artifact Registry. If you currently use Container Registry, you should [transition to Artifact Registry](https://cloud.google.com/artifact-registry/docs/transition/transition-from-gcr).
+  - You must run your Batch job's VMs in the same [location](https://cloud.google.com/artifact-registry/docs/repositories/repo-locations) as the Artifact Registry storing the container image.
+
 `google.batch.maxSpotAttempts`
 : :::{versionadded} 23.11.0-edge
   :::

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

@@ -313,6 +313,7 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
             .setImageUri(task.container)
             .addAllCommands(cmd)
             .addAllVolumes(launcher.getContainerMounts())
+            .setEnableImageStreaming(batchConfig?.enableImageStreaming ?: false)
 
         def containerOptions = task.config.getContainerOptions() ?: ''
         if( fusionEnabled() ) {

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

@@ -90,6 +90,22 @@ class BatchConfig implements ConfigScope {
     """)
     final String logsPath
 
+    @ConfigOption
+    @Description("""
+        Enable container image streaming to speed up job start-up (default: `false`).
+        This can reduce latency for large images but comes with some
+        [limitations](https://cloud.google.com/batch/docs/use-image-streaming), including:
+
+        - `process.containerOptions` is ignored when image streaming is enabled.
+        - Image streaming only supports container images stored in Artifact Registry.
+          If you currently use Container Registry, you should 
+          [transition to Artifact Registry](https://cloud.google.com/artifact-registry/docs/transition/transition-from-gcr).
+        - You must run your Batch job's VMs in the same
+          [location](https://cloud.google.com/artifact-registry/docs/repositories/repo-locations)
+          as the Artifact Registry storing the container image.
+    """)
+    final boolean enableImageStreaming
+
     @ConfigOption
     @Description("""
         Max number of execution attempts of a job interrupted by a Compute Engine Spot reclaim event (default: `0`).
@@ -148,6 +164,7 @@ class BatchConfig implements ConfigScope {
         gcsfuseOptions = opts.gcsfuseOptions as List<String> ?: DEFAULT_GCSFUSE_OPTS
         installGpuDrivers = opts.installGpuDrivers as boolean
         logsPath = opts.logsPath
+        enableImageStreaming = opts.enableImageStreaming != null ? opts.enableImageStreaming as boolean : false
         maxSpotAttempts = opts.maxSpotAttempts != null ? opts.maxSpotAttempts as int : DEFAULT_MAX_SPOT_ATTEMPTS
         network = opts.network
         networkTags = opts.networkTags as List<String> ?: Collections.emptyList()

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

@@ -31,6 +31,8 @@ class BatchConfigTest extends Specification {
         then:
         !config.getSpot()
         and:
+        !config.enableImageStreaming
+        and:
         config.retryConfig.maxAttempts == 5
         config.maxSpotAttempts == 0
         config.autoRetryExitCodes == [50001]
@@ -44,6 +46,7 @@ class BatchConfigTest extends Specification {
         given:
         def opts = [
             spot: true,
+            enableImageStreaming: true,
             maxSpotAttempts: 8,
             autoRetryExitCodes: [50001, 50003, 50005],
             retryPolicy: [maxAttempts: 10],