From 35b067a199bb5437e2a910f2256012f7d30d4485 Mon Sep 17 00:00:00 2001 From: Vu Date: Sat, 2 May 2026 13:56:43 +0700 Subject: [PATCH] Add shm-size and ipc-mode options to docker run script generation Extends TaskCommandFileGenerator and the dockerRunScript SSP template with two new optional parameters: shmSize and ipcMode. When set, the generated docker run command includes --shm-size and --ipc flags respectively. Both default to None, leaving existing behavior unchanged. Flags are emitted as quoted direct argv tokens in the docker run invocation rather than through shell variable expansion, preventing shell-meta injection even if upstream validation were bypassed. --- RELEASE_NOTES.md | 6 +- .../templates/exec/dockerRunScript.ssp | 8 ++ .../scala/wdlTools/exec/TaskExecutor.scala | 18 ++++- .../exec/TaskCommandFileGeneratorTest.scala | 76 +++++++++++++++++++ 4 files changed, 103 insertions(+), 5 deletions(-) create mode 100644 src/test/scala/wdlTools/exec/TaskCommandFileGeneratorTest.scala diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 005cd8fe..2ca82d28 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -2,7 +2,11 @@ ## in develop -... +* `TaskCommandFileGenerator.writeDockerRunScript` and `apply` now accept optional + `shmSize` and `ipcMode` parameters that emit `--shm-size=` and `--ipc=` + 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) diff --git a/src/main/resources/templates/exec/dockerRunScript.ssp b/src/main/resources/templates/exec/dockerRunScript.ssp index 2b3ee087..6d2f22aa 100644 --- a/src/main/resources/templates/exec/dockerRunScript.ssp +++ b/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} diff --git a/src/main/scala/wdlTools/exec/TaskExecutor.scala b/src/main/scala/wdlTools/exec/TaskExecutor.scala index 4d278357..cadfff1c 100644 --- a/src/main/scala/wdlTools/exec/TaskExecutor.scala +++ b/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) } diff --git a/src/test/scala/wdlTools/exec/TaskCommandFileGeneratorTest.scala b/src/test/scala/wdlTools/exec/TaskCommandFileGeneratorTest.scala new file mode 100644 index 00000000..6f84c4eb --- /dev/null +++ b/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 \\") + } +}