diff --git a/src/main/scala/io/viash/Main.scala b/src/main/scala/io/viash/Main.scala index f98102dd3..4c7d7b2f7 100644 --- a/src/main/scala/io/viash/Main.scala +++ b/src/main/scala/io/viash/Main.scala @@ -68,7 +68,13 @@ object Main { cli.subcommands match { case List(cli.run) => val config = readConfig(cli.run) - ViashRun(config, args = runArgs.dropWhile(_ == "--"), keepFiles = cli.run.keep.toOption.map(_.toBoolean)) + ViashRun( + config, + args = runArgs.dropWhile(_ == "--"), + keepFiles = cli.run.keep.toOption.map(_.toBoolean), + cpus = cli.run.cpus.toOption, + memory = cli.run.memory.toOption + ) case List(cli.build) => val config = readConfig(cli.build) ViashBuild( @@ -82,7 +88,12 @@ object Main { 0 // Exceptions are thrown when something bad happens, so then the '0' is not returned but a '1'. Can be improved further. case List(cli.test) => val config = readConfig(cli.test, applyPlatform = false) - ViashTest(config, keepFiles = cli.test.keep.toOption.map(_.toBoolean)) + ViashTest( + config, + keepFiles = cli.test.keep.toOption.map(_.toBoolean), + cpus = cli.test.cpus.toOption, + memory = cli.test.memory.toOption + ) 0 // Exceptions are thrown when a test fails, so then the '0' is not returned but a '1'. Can be improved further. case List(cli.namespace, cli.namespace.build) => val configs = readConfigs(cli.namespace.build) @@ -103,7 +114,9 @@ object Main { parallel = cli.namespace.test.parallel(), keepFiles = cli.namespace.test.keep.toOption.map(_.toBoolean), tsv = cli.namespace.test.tsv.toOption, - append = cli.namespace.test.append() + append = cli.namespace.test.append(), + cpus = cli.namespace.test.cpus.toOption, + memory = cli.namespace.test.memory.toOption ) val errors = testResults.flatMap(_.right.toOption).count(_.isError) if (errors > 0) 1 else 0 diff --git a/src/main/scala/io/viash/ViashNamespace.scala b/src/main/scala/io/viash/ViashNamespace.scala index 51a915dce..c41d3d49b 100644 --- a/src/main/scala/io/viash/ViashNamespace.scala +++ b/src/main/scala/io/viash/ViashNamespace.scala @@ -76,7 +76,9 @@ object ViashNamespace { parallel: Boolean = false, keepFiles: Option[Boolean] = None, tsv: Option[String] = None, - append: Boolean = false + append: Boolean = false, + cpus: Option[Int], + memory: Option[String] ): List[Either[(Config, ManyTestOutput), Status]] = { // we can't currently test nextflow platforms, so exclude them from the tests val testableConfigs = configs.filter(conf => @@ -151,7 +153,9 @@ object ViashNamespace { config = conf, keepFiles = keepFiles, quiet = true, - parentTempPath = Some(parentTempPath) + parentTempPath = Some(parentTempPath), + cpus = cpus, + memory = memory ) } catch { case e: MissingResourceFileException => diff --git a/src/main/scala/io/viash/ViashRun.scala b/src/main/scala/io/viash/ViashRun.scala index e8b9c2a6c..9f875f028 100644 --- a/src/main/scala/io/viash/ViashRun.scala +++ b/src/main/scala/io/viash/ViashRun.scala @@ -27,7 +27,7 @@ import io.viash.helpers.Circe.{OneOrMore, One, More} import scala.sys.process.{Process, ProcessLogger} object ViashRun { - def apply(config: Config, args: Seq[String], keepFiles: Option[Boolean]): Int = { + def apply(config: Config, args: Seq[String], keepFiles: Option[Boolean], cpus: Option[Int], memory: Option[String]): Int = { val fun = config.functionality val dir = IO.makeTemp("viash_" + fun.name) @@ -39,7 +39,9 @@ object ViashRun { // determine command val cmd = - Array(Paths.get(dir.toString, fun.name).toString) ++ args + Array(Paths.get(dir.toString, fun.name).toString) ++ + args ++ + Array(cpus.map("---cpus=" + _), memory.map("---memory="+_)).flatMap(a => a) // execute command, print everything to console code = Process(cmd).!(ProcessLogger(println, println)) diff --git a/src/main/scala/io/viash/ViashTest.scala b/src/main/scala/io/viash/ViashTest.scala index dcdc86943..7cf37e6dc 100644 --- a/src/main/scala/io/viash/ViashTest.scala +++ b/src/main/scala/io/viash/ViashTest.scala @@ -44,7 +44,9 @@ object ViashTest { setupStrategy: String = "cachedbuild", tempVersion: Boolean = true, verbosityLevel: Int = 6, - parentTempPath: Option[Path] = None + parentTempPath: Option[Path] = None, + cpus: Option[Int], + memory: Option[String] ): ManyTestOutput = { // create temporary directory val dir = IO.makeTemp("viash_test_" + config.functionality.name, parentTempPath) @@ -67,7 +69,9 @@ object ViashTest { dir = dir, verbose = !quiet, setupStrategy = setupStrategy, - verbosityLevel = verbosityLevel + verbosityLevel = verbosityLevel, + cpus = cpus, + memory = memory ) val count = results.count(_.exitValue == 0) val anyErrors = setupRes.exists(_.exitValue > 0) || count < results.length @@ -106,14 +110,22 @@ object ViashTest { ManyTestOutput(setupRes, results) } - def runTests(config: Config, dir: Path, verbose: Boolean = true, setupStrategy: String, verbosityLevel: Int): ManyTestOutput = { + def runTests(config: Config, dir: Path, verbose: Boolean = true, setupStrategy: String, verbosityLevel: Int, cpus: Option[Int], memory: Option[String]): ManyTestOutput = { val fun = config.functionality val platform = config.platform.get val consoleLine = "====================================================================" // build regular executable - val buildFun = platform.modifyFunctionality(config, true) + val configWithReqs = config.copy( + config.functionality.copy( + requirements = config.functionality.requirements.copy( + cpus = if (cpus.isDefined) cpus else config.functionality.requirements.cpus, + memory = if (memory.isDefined) memory else config.functionality.requirements.memory + ) + ) + ) + val buildFun = platform.modifyFunctionality(configWithReqs, true) val buildDir = dir.resolve("build_executable") Files.createDirectories(buildDir) try { @@ -245,8 +257,12 @@ object ViashTest { try { // run command, collect output val executable = Paths.get(newDir.toString, testBash.filename).toString - logger(s"+$executable") - val exitValue = Process(Seq(executable), cwd = newDir.toFile).!(ProcessLogger(logger, logger)) + val cmd = Seq(executable) ++ Seq(cpus.map("---cpus=" + _), memory.map("---memory="+_)).flatMap(a => a) + logger(s"+${cmd.mkString(" ")}") + val exitValue = Process( + cmd, + cwd = newDir.toFile + ).!(ProcessLogger(logger, logger)) printWriter.flush() diff --git a/src/main/scala/io/viash/cli/CLIConf.scala b/src/main/scala/io/viash/cli/CLIConf.scala index 6dee747f6..77b74167a 100644 --- a/src/main/scala/io/viash/cli/CLIConf.scala +++ b/src/main/scala/io/viash/cli/CLIConf.scala @@ -48,6 +48,21 @@ trait ViashCommand { descr = "Modify a viash config at runtime using @[config_mod](dynamic config modding)." ) } +trait ViashRunner { + _: DocumentedSubcommand => + val cpus = registerOpt[Int]( + name = "cpus", + default = None, + descr = "The maximum number of (logical) cpus a component is allowed to use.", + required = false + ) + val memory = registerOpt[String]( + name = "memory", + descr = "The maximum amount of memory a component is allowed to allocate. Unit must be one of B, KB, MB, GB, TB or PB.", + default = None, + required = false + ) +} trait ViashNs { _: DocumentedSubcommand => val query = registerOpt[String]( @@ -143,7 +158,7 @@ class CLIConf(arguments: Seq[String]) extends ScallopConf(arguments) { | |Arguments:""".stripMargin) - val run = new DocumentedSubcommand("run") with ViashCommand with WithTemporary { + val run = new DocumentedSubcommand("run") with ViashCommand with WithTemporary with ViashRunner { banner( "viash run", "Executes a viash component from the provided viash config file. viash generates a temporary executable and immediately executes it with the given parameters.", @@ -197,7 +212,7 @@ class CLIConf(arguments: Seq[String]) extends ScallopConf(arguments) { ) } - val test = new DocumentedSubcommand("test") with ViashCommand with WithTemporary { + val test = new DocumentedSubcommand("test") with ViashCommand with WithTemporary with ViashRunner { banner( "viash test", "Test the component using the tests defined in the viash config file.", @@ -282,7 +297,7 @@ class CLIConf(arguments: Seq[String]) extends ScallopConf(arguments) { ) } - val test = new DocumentedSubcommand("test") with ViashNs with WithTemporary { + val test = new DocumentedSubcommand("test") with ViashNs with WithTemporary with ViashRunner { banner( "viash ns test", "Test a namespace containing many viash config files.",