Merge branch 'develop' into wx_1819_reference_disks_in_batch

broadinstitute · Sep 11, 2024 · a2b6352 · a2b6352
2 parents bfa1f67 + abc410f
commit a2b6352
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -34,6 +34,7 @@ be found [here](https://cromwell.readthedocs.io/en/stable/backends/HPC/#optional
 - Fixes pulling Docker image metadata from private GCR repositories.
 - Fixed `google_project` and `google_compute_service_account` workflow options not taking effect when using GCP Batch backend
 - Added a way to use a custom LogsPolicy for the job execution, setting `backend.providers.batch.config.batch.logs-policy` to "CLOUD_LOGGING" (default) keeps the current behavior, or, set it to "PATH" to save the logs into the the mounted disk, at the end, this log file gets copied to the google cloud storage bucket with "task.log" as the name.
+- When "CLOUD_LOGGING" is used, many more Cromwell / WDL labels for workflow, root workflow, call, shard etc. are now assigned to GCP Batch log entries.
 
 ### Improved handling of Life Sciences API quota errors
 

diff --git a/docs/backends/GCPBatch.md b/docs/backends/GCPBatch.md
@@ -171,12 +171,15 @@ backend.providers.GCPBATCH.config {
 Every call run on the GCP Batch backend is given certain labels by default, so that Google resources can be queried by these labels later. 
 The current default label set automatically applied is:
 
-| Key | Value | Example | Notes |
-|-----|-------|---------|-------|
-| cromwell-workflow-id | The Cromwell ID given to the root workflow (i.e. the ID returned by Cromwell on submission) | cromwell-d4b412c5-bf3d-4169-91b0-1b635ce47a26 | To fit the required [format](#label-format), we prefix with 'cromwell-' |
-| cromwell-sub-workflow-name | The name of this job's sub-workflow | my-sub-workflow | Only present if the task is called in a subworkflow. |
-| wdl-task-name | The name of the WDL task | my-task | |
-| wdl-call-alias | The alias of the WDL call that created this job | my-task-1 | Only present if the task was called with an alias. |
+| Key                        | Value                                                                                       | Example                                           | Notes                                                                                                                             |
+|----------------------------|---------------------------------------------------------------------------------------------|---------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------|
+| cromwell-workflow-id       | The Cromwell ID given to the root workflow (i.e. the ID returned by Cromwell on submission) | cromwell-d4b412c5-bf3d-4169-91b0-1b635ce47a26     | To fit the required [format](#label-format), we prefix with 'cromwell-'                                                           |
+| cromwell-sub-workflow-id   | The Cromwell ID given to this job's sub-workflow (immediate parent workflow)                | cromwell-sub-d4b412c5-bf3d-4169-91b0-1b635ce47a26 | To fit the required [format](#label-format), we prefix with 'cromwell-sub-'. Only present if the task is called in a subworkflow. |
+| cromwell-sub-workflow-name | The name of this job's sub-workflow                                                         | my-sub-workflow                                   | Only present if the task is called in a subworkflow.                                                                              |
+| wdl-task-name              | The name of the WDL task                                                                    | my-task                                           |                                                                                                                                   |
+| wdl-call-alias             | The alias of the WDL call that created this job                                             | my-task-1                                         | Only present if the task was called with an alias.                                                                                |
+| wdl-attempt                | Attempt number for this call                                                                | 1                                                 |                                                                                                                                   |
+| wdl-shard-index            | Index of this job within a scatter,                                                         |                                                   | Only present if the task was called within a scatter.                                                                             |
 
 Any custom labels provided as '`google_labels`' in the [workflow options](../wf_options/Google) are also applied to Google resources by GCP Batch.
 

diff --git a/...e/batch/src/main/scala/cromwell/backend/google/batch/api/GcpBatchRequestFactoryImpl.scala b/...e/batch/src/main/scala/cromwell/backend/google/batch/api/GcpBatchRequestFactoryImpl.scala
@@ -23,6 +23,7 @@ import cromwell.backend.google.batch.models.GcpBatchConfigurationAttributes.GcsT
 import cromwell.backend.google.batch.models.{GcpBatchLogsPolicy, GcpBatchRequest, VpcAndSubnetworkProjectLabelValues}
 import cromwell.backend.google.batch.runnable._
 import cromwell.backend.google.batch.util.{BatchUtilityConversions, GcpBatchMachineConstraints}
+import cromwell.core.labels.{Label, Labels}
 import cromwell.core.logging.JobLogger
 
 import scala.jdk.CollectionConverters._
@@ -248,14 +249,43 @@ class GcpBatchRequestFactoryImpl()(implicit gcsTransferConfiguration: GcsTransfe
           .build
     }
 
+    val googleLabels = data.createParameters.googleLabels.map(l => Label(l.key, l.value))
+
+    val jobDescriptor = data.createParameters.jobDescriptor
+    val backendJobDescriptorKey = jobDescriptor.key
+
+    val workflow = jobDescriptor.workflowDescriptor
+    val call = jobDescriptor.taskCall
+    val subWorkflow = workflow.callable
+    val subWorkflowLabels =
+      if (!subWorkflow.equals(workflow.rootWorkflow))
+        Labels("cromwell-sub-workflow-name" -> subWorkflow.name,
+               "cromwell-sub-workflow-id" -> s"cromwell-sub-${jobDescriptor.workflowDescriptor.id.toString}"
+        )
+      else
+        Labels.empty
+
+    val alias = call.localName
+    val aliasLabels =
+      if (!alias.equals(call.callable.name))
+        Labels("wdl-call-alias" -> alias)
+      else
+        Labels.empty
+
+    val shardLabels = Labels(backendJobDescriptorKey.index.map(l => Label("wdl-shard-index", l.toString)).toVector)
+
+    val allLabels = Labels(
+      "cromwell-workflow-id" -> s"cromwell-${workflow.rootWorkflowId}",
+      "wdl-task-name" -> call.callable.name,
+      "wdl-attempt" -> backendJobDescriptorKey.attempt.toString,
+      "goog-batch-worker" -> "true",
+      "submitter" -> "cromwell"
+    ) ++ shardLabels ++ subWorkflowLabels ++ aliasLabels ++ Labels(googleLabels.toVector)
+
     val job = Job.newBuilder
       .addTaskGroups(taskGroup)
       .setAllocationPolicy(allocationPolicy.build())
-      .putLabels("submitter",
-                 "cromwell"
-      ) // label to signify job submitted by cromwell for larger tracking purposes within GCP batch
-      .putLabels("goog-batch-worker", "true")
-      .putAllLabels(data.createParameters.googleLabels.map(label => label.key -> label.value).toMap.asJava)
+      .putAllLabels(allLabels.asJavaMap)
       .setLogsPolicy(logsPolicy)
 
     CreateJobRequest.newBuilder