Add Google Batch LogsPolicy PATH option for logging to GCS (#6431)

---------

Signed-off-by: David Glazer <[email protected]>
Signed-off-by: Emma Rogge <[email protected]>
Signed-off-by: Ben Sherman <[email protected]>

Author: dglazer

docs/reference/config.md

@@ -922,6 +922,12 @@ The following settings are available:
   :::
 : List of custom mount options for `gcsfuse` (default: `['-o rw', '-implicit-dirs']`).
 
+`google.batch.logsPath`
+: :::{versionadded} 25.11.0-edge
+  :::
+: 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.maxSpotAttempts`
 : :::{versionadded} 23.11.0-edge
   :::

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

@@ -186,9 +186,17 @@ class GoogleBatchScriptLauncher extends BashWrapperBuilder implements GoogleBatc
 
     GoogleBatchScriptLauncher withConfig(GoogleOpts config) {
         this.config = config
+        addLogsBucket(config.batch.logsPath())
         return this
     }
 
+    protected void addLogsBucket(Path path) {
+        if( path instanceof CloudStoragePath )
+            buckets.add(path.bucket())
+        else if( path != null )
+            throw new IllegalArgumentException("Unexpected value for Google Batch logs path: ${path.toUriString()}")
+    }
+
     GoogleBatchScriptLauncher withIsArray(boolean value) {
         this.isArray = value
         return this

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

@@ -32,6 +32,7 @@ import com.google.cloud.batch.v1.ServiceAccount
 import com.google.cloud.batch.v1.TaskGroup
 import com.google.cloud.batch.v1.TaskSpec
 import com.google.cloud.batch.v1.Volume
+import com.google.cloud.storage.contrib.nio.CloudStoragePath
 import com.google.protobuf.Duration
 import groovy.transform.CompileStatic
 import groovy.transform.PackageScope
@@ -447,14 +448,26 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
         return Job.newBuilder()
             .addTaskGroups(taskGroup)
             .setAllocationPolicy(allocationPolicy)
-            .setLogsPolicy(
-                LogsPolicy.newBuilder()
-                    .setDestination(LogsPolicy.Destination.CLOUD_LOGGING)
-            )
+            .setLogsPolicy(createLogsPolicy())
             .putAllLabels(task.config.getResourceLabels())
             .build()
     }
 
+    protected LogsPolicy createLogsPolicy() {
+        final logsPath = executor.batchConfig.logsPath()
+        if( logsPath instanceof CloudStoragePath ) {
+            return LogsPolicy.newBuilder()
+                .setDestination(LogsPolicy.Destination.PATH)
+                .setLogsPath(GoogleBatchScriptLauncher.containerMountPath(logsPath))
+                .build()
+        }
+        else {
+            return LogsPolicy.newBuilder()
+                .setDestination(LogsPolicy.Destination.CLOUD_LOGGING)
+                .build()
+        }
+    }
+
     /**
      * @return Retrieve the submitted task state
      */

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

@@ -16,11 +16,10 @@
 
 package nextflow.cloud.google.batch.client
 
-import com.google.auth.oauth2.GoogleCredentials
+import java.nio.file.Path
+
 import groovy.transform.CompileStatic
 import groovy.util.logging.Slf4j
-import nextflow.Session
-import nextflow.cloud.google.GoogleOpts
 import nextflow.config.spec.ConfigOption
 import nextflow.config.spec.ConfigScope
 import nextflow.script.dsl.Description
@@ -85,6 +84,12 @@ class BatchConfig implements ConfigScope {
     """)
     final boolean installGpuDrivers
 
+    @ConfigOption
+    @Description("""
+        The Google Cloud Storage path where job logs should be stored, e.g. `gs://my-logs-bucket/logs`.
+    """)
+    final String logsPath
+
     @ConfigOption
     @Description("""
         Max number of execution attempts of a job interrupted by a Compute Engine Spot reclaim event (default: `0`).
@@ -142,6 +147,7 @@ class BatchConfig implements ConfigScope {
         cpuPlatform = opts.cpuPlatform
         gcsfuseOptions = opts.gcsfuseOptions as List<String> ?: DEFAULT_GCSFUSE_OPTS
         installGpuDrivers = opts.installGpuDrivers as boolean
+        logsPath = opts.logsPath
         maxSpotAttempts = opts.maxSpotAttempts != null ? opts.maxSpotAttempts as int : DEFAULT_MAX_SPOT_ATTEMPTS
         network = opts.network
         networkTags = opts.networkTags as List<String> ?: Collections.emptyList()
@@ -155,4 +161,8 @@ class BatchConfig implements ConfigScope {
 
     BatchRetryConfig getRetryConfig() { retry }
 
+    Path logsPath() {
+        return logsPath as Path
+    }
+
 }

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

@@ -142,6 +142,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
         def MACHINE_TYPE = 'vm-type-2'
         def MEM = MemoryUnit.of('8 GB')
         def TIMEOUT = Duration.of('1 hour')
+        def LOGS_PATH = CloudStorageFileSystem.forBucket('my-logs-bucket').getPath('/logs')
         and:
         def exec = Mock(GoogleBatchExecutor) {
             getBatchConfig() >> Mock(BatchConfig) {
@@ -157,6 +158,7 @@ class GoogleBatchTaskHandlerTest extends Specification {
                 getServiceAccountEmail() >> '[email protected]'
                 getSubnetwork() >> 'subnet-1'
                 usePrivateAddress >> true
+                logsPath() >> LOGS_PATH
             }
         }
         and:
@@ -239,7 +241,8 @@ class GoogleBatchTaskHandlerTest extends Specification {
         networkInterface.getSubnetwork() == 'subnet-1'
         networkInterface.getNoExternalIpAddress() == true
         and:
-        req.getLogsPolicy().getDestination().toString() == 'CLOUD_LOGGING'
+        req.getLogsPolicy().getDestination().toString() == 'PATH'
+        req.getLogsPolicy().getLogsPath() == '/mnt/disks/my-logs-bucket/logs'
         and:
         req.getLabelsMap() == [foo: 'bar']
 

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

@@ -17,17 +17,14 @@
 
 package nextflow.cloud.google.batch.client
 
-import nextflow.Session
 import nextflow.util.MemoryUnit
-import spock.lang.Requires
 import spock.lang.Specification
 /**
  *
  * @author Paolo Di Tommaso <[email protected]>
  */
 class BatchConfigTest extends Specification {
 
-    @Requires({System.getenv('GOOGLE_APPLICATION_CREDENTIALS')})
     def 'should create batch config' () {
         when:
         def config = new BatchConfig([:])
@@ -40,9 +37,9 @@ class BatchConfigTest extends Specification {
         and:
         !config.bootDiskImage
         !config.bootDiskSize
+        !config.logsPath
     }
 
-    @Requires({System.getenv('GOOGLE_APPLICATION_CREDENTIALS')})
     def 'should create batch config with custom settings' () {
         given:
         def opts = [
@@ -51,7 +48,8 @@ class BatchConfigTest extends Specification {
             autoRetryExitCodes: [50001, 50003, 50005],
             retryPolicy: [maxAttempts: 10],
             bootDiskImage: 'batch-foo',
-            bootDiskSize: '100GB'
+            bootDiskSize: '100GB',
+            logsPath: 'gs://my-logs-bucket/logs'
         ]
 
         when:
@@ -65,6 +63,8 @@ class BatchConfigTest extends Specification {
         and:
         config.bootDiskImage == 'batch-foo'
         config.bootDiskSize == MemoryUnit.of('100GB')
+        and:
+        config.logsPath == 'gs://my-logs-bucket/logs'
     }
 
 }