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 \\") + } +}