From 61063cb2b27e7c32aef322a7f8f59e708ebcd286 Mon Sep 17 00:00:00 2001
From: Khalid Shakir <kshakir@broadinstitute.org>
Date: Fri, 10 May 2024 21:04:24 -0400
Subject: [PATCH] Optionally export MEM_SIZE and MEM_UNIT for all standard
 backends.

---
 .../StandardAsyncExecutionActor.scala         | 33 +++++++++++++++++--
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/backend/src/main/scala/cromwell/backend/standard/StandardAsyncExecutionActor.scala b/backend/src/main/scala/cromwell/backend/standard/StandardAsyncExecutionActor.scala
index bbb796220c4..57cc86a82df 100644
--- a/backend/src/main/scala/cromwell/backend/standard/StandardAsyncExecutionActor.scala
+++ b/backend/src/main/scala/cromwell/backend/standard/StandardAsyncExecutionActor.scala
@@ -39,6 +39,7 @@ import org.apache.commons.lang3.exception.ExceptionUtils
 import shapeless.Coproduct
 import wom.callable.{AdHocValue, CommandTaskDefinition, ContainerizedInputExpression}
 import wom.expression.WomExpression
+import wom.format.MemorySize
 import wom.graph.LocalName
 import wom.values._
 import wom.{CommandSetupSideEffectFile, InstantiatedCommand, WomFileMapper}
@@ -427,9 +428,25 @@ trait StandardAsyncExecutionActor
     val errorOrGlobFiles: ErrorOr[List[WomGlobFile]] =
       backendEngineFunctions.findGlobOutputs(call, jobDescriptor)
 
-    lazy val environmentVariables = instantiatedCommand.environmentVariables map { case (k, v) =>
-      s"""export $k="$v""""
-    } mkString ("", "\n", "\n")
+    lazy val environmentVariables = {
+      /*
+      Add `MEM_SIZE` and `MEM_UNIT` before the other environment variables on all backends that define a `memory`
+      runtime attribute. As of May 2022 some backends may expose these same environment variables via other means where
+      they are accessible elsewhere, for example within sidecar containers used for resource monitoring.
+       */
+      val memoryEnvironmentVariables: List[(String, String)] =
+        runtimeMemoryOption.toList.flatMap(runtimeMemory =>
+          List(
+            "MEM_SIZE" -> runtimeMemory.amount.toString,
+            "MEM_UNIT" -> runtimeMemory.unit.toString
+          )
+        )
+      val environmentVariables: List[(String, String)] =
+        memoryEnvironmentVariables ++ instantiatedCommand.environmentVariables
+      environmentVariables map { case (key, value) =>
+        s"""export $key="$value""""
+      } mkString ("", "\n", "\n")
+    }
 
     val shortId = jobDescriptor.workflowDescriptor.id.shortString
     // Give the out and error FIFO variables names that are unlikely to conflict with anything the user is doing.
@@ -755,6 +772,16 @@ trait StandardAsyncExecutionActor
   lazy val continueOnReturnCode: ContinueOnReturnCode =
     RuntimeAttributesValidation.extract(ContinueOnReturnCodeValidation.instance, validatedRuntimeAttributes)
 
+  /**
+    * Returns the memory size for the job.
+    *
+    * @return the memory size for the job.
+    */
+  lazy val runtimeMemoryOption: Option[MemorySize] = RuntimeAttributesValidation.extractOption(
+    runtimeAttributesValidation = MemoryValidation.instance(),
+    validatedRuntimeAttributes = validatedRuntimeAttributes
+  )
+
   /**
     * Returns the max number of times that a failed job should be retried, obtained by converting `maxRetries` to an Int.
     */