Add shm-size and ipc-mode options to docker run script generation

RELEASE_NOTES.md

@@ -2,7 +2,11 @@
 
 ## in develop
 
-...
+* `TaskCommandFileGenerator.writeDockerRunScript` and `apply` now accept optional
+  `shmSize` and `ipcMode` parameters that emit `--shm-size=<size>` and `--ipc=<mode>`
+  flags in the generated `docker run` script. Both default to `None` (no flag emitted),
+  so existing callers are unaffected. Enables multi-GPU NCCL workloads that need more
+  than Docker's 64 MB default `/dev/shm`. (APPS-3954)
 
 ## 0.17.17 (2024-02-29)
 

src/main/resources/templates/exec/dockerRunScript.ssp

@@ -5,6 +5,8 @@
 <%@ val commandFile: String %>
 <%@ val imageName: String %>
 <%@ val maxMemory: Long %>
+<%@ val shmSize: Option[String] %>
+<%@ val ipcMode: Option[String] %>
 <% val bashDollar: String = "$" %>
 #!/bin/bash -x
 
@@ -22,6 +24,12 @@ docker run \\
 --memory=${maxMemory.toString} \\
 --cidfile ${containerIdFile} \\
 ${bashDollar}{extraFlags} \\
+#if (shmSize.isDefined)
+"--shm-size=${shmSize.get}" \\
+#end
+#if (ipcMode.isDefined)
+"--ipc=${ipcMode.get}" \\
+#end
 --entrypoint /bin/bash \\
 -v ${hostRootDir}:${containerRootDir} \\
 ${imageName} ${commandFile}

src/main/scala/wdlTools/exec/TaskExecutor.scala

@@ -71,7 +71,9 @@ case class TaskCommandFileGenerator(logger: Logger = Logger.get) {
   def writeDockerRunScript(imageName: String,
                            hostPaths: ExecPaths,
                            guestPaths: ExecPaths,
-                           maxMemory: Long = SysUtils.totalMemorySize): Path = {
+                           maxMemory: Long = SysUtils.totalMemorySize,
+                           shmSize: Option[String] = None,
+                           ipcMode: Option[String] = None): Path = {
     val dockerRunScript = renderer.render(
         TaskCommandFileGenerator.DefaultDockerRunScript,
         Map(
@@ -82,7 +84,9 @@ case class TaskCommandFileGenerator(logger: Logger = Logger.get) {
             "stdoutFile" -> guestPaths.getStdoutFile().toString,
             "stderrFile" -> guestPaths.getStderrFile().toString,
             "imageName" -> imageName,
-            "maxMemory" -> maxMemory
+            "maxMemory" -> maxMemory,
+            "shmSize" -> shmSize,
+            "ipcMode" -> ipcMode
         )
     )
     val commandFile = hostPaths.getContainerCommandFile(true).asJavaPath
@@ -93,11 +97,17 @@ case class TaskCommandFileGenerator(logger: Logger = Logger.get) {
 
   def apply(command: Option[String],
             hostPaths: ExecPaths,
-            container: Option[(String, ExecPaths)] = None): Path = {
+            container: Option[(String, ExecPaths)] = None,
+            shmSize: Option[String] = None,
+            ipcMode: Option[String] = None): Path = {
     if (container.isDefined) {
       val (containerImage, guestPaths) = container.get
       writeCommandScript(command, hostPaths, Some(guestPaths))
-      writeDockerRunScript(containerImage, hostPaths, guestPaths)
+      writeDockerRunScript(containerImage,
+                           hostPaths,
+                           guestPaths,
+                           shmSize = shmSize,
+                           ipcMode = ipcMode)
     } else {
       writeCommandScript(command, hostPaths)
     }

src/test/scala/wdlTools/exec/TaskCommandFileGeneratorTest.scala

@@ -0,0 +1,76 @@
+package wdlTools.exec
+
+import dx.util.{FileUtils, Logger, PosixPath}
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class TaskCommandFileGeneratorTest extends AnyFlatSpec with Matchers {
+  private val logger = Logger.Quiet
+  private val containerMountDir = PosixPath("/home/wdlTools")
+
+  private def renderDockerScript(shmSize: Option[String] = None,
+                                 ipcMode: Option[String] = None): String = {
+    val (hostPaths, guestPaths) =
+      DefaultExecPaths.createLocalContainerPair(containerMountDir = containerMountDir)
+    val generator = TaskCommandFileGenerator(logger)
+    val scriptPath = generator.writeDockerRunScript(
+        imageName = "test-image:latest",
+        hostPaths = hostPaths,
+        guestPaths = guestPaths,
+        maxMemory = 1024L * 1024L * 1024L,
+        shmSize = shmSize,
+        ipcMode = ipcMode
+    )
+    FileUtils.readFileContent(scriptPath)
+  }
+
+  it should "omit --shm-size and --ipc when neither is set" in {
+    val script = renderDockerScript()
+    script should not include "--shm-size"
+    script should not include "--ipc="
+    script should include("--memory=1073741824")
+    script should include("test-image:latest")
+  }
+
+  it should "include quoted --shm-size when shmSize is set" in {
+    val script = renderDockerScript(shmSize = Some("8g"))
+    script should include("\"--shm-size=8g\"")
+    script should not include "--ipc="
+  }
+
+  it should "include quoted --ipc when ipcMode is set" in {
+    val script = renderDockerScript(ipcMode = Some("host"))
+    script should include("\"--ipc=host\"")
+    script should not include "--shm-size"
+  }
+
+  it should "include both flags when both are set" in {
+    val script = renderDockerScript(shmSize = Some("4g"), ipcMode = Some("host"))
+    script should include("\"--shm-size=4g\"")
+    script should include("\"--ipc=host\"")
+  }
+
+  it should "preserve --user/--hostname extraFlags when shm/ipc are set" in {
+    val script = renderDockerScript(shmSize = Some("2g"), ipcMode = Some("host"))
+    script should include("--user $(id -u):$(id -g)")
+    script should include("--hostname $(hostname)")
+  }
+
+  it should "emit shm/ipc as direct docker run args, not via extraFlags shell var" in {
+    // Defense-in-depth: callers (e.g. dxCompiler) are responsible for validating values, but
+    // the rendered script also avoids the unquoted ${extraFlags} expansion path so that even
+    // if an unvalidated value contained shell metas, it would appear as a single argv token.
+    val script = renderDockerScript(shmSize = Some("2g"), ipcMode = Some("host"))
+    script should not include "extraFlags=\"${extraFlags} --shm-size"
+    script should not include "extraFlags=\"${extraFlags} --ipc"
+  }
+
+  it should "render shm-size on its own backslash-continued docker run line" in {
+    // Pin the conditional render so a stray newline or missing trailing backslash in the
+    // template would break this test rather than producing a broken bash command at runtime.
+    val script = renderDockerScript(shmSize = Some("8g"), ipcMode = Some("host"))
+    script should include("\"--shm-size=8g\" \\")
+    script should include("\"--ipc=host\" \\")
+    script should include("--entrypoint /bin/bash \\")
+  }
+}