diff --git a/garlic/stages/srun.nix b/garlic/stages/srun.nix index d7e8539..57fc667 100644 --- a/garlic/stages/srun.nix +++ b/garlic/stages/srun.nix @@ -8,6 +8,7 @@ , cpuBind , nixPrefix , preSrun ? "" +, postSrun ? "" , srunOptions ? "" , output ? "stdout.log" , error ? "stderr.log" @@ -26,14 +27,17 @@ stdenv.mkDerivation rec { ${preSrun} - exec ${slurm}/bin/srun \ + ${slurm}/bin/srun \ --mpi=pmi2 \ --cpu-bind=${cpuBind} \ --output=${output} \ --error=${error} \ ${srunOptions} \ ${nixPrefix}${stageProgram nextStage} + + ${postSrun} EOF + chmod +x $out ''; } diff --git a/garlic/stdexp.nix b/garlic/stdexp.nix index 52e8ab9..33a18b1 100644 --- a/garlic/stdexp.nix +++ b/garlic/stdexp.nix @@ -76,12 +76,12 @@ rec { inherit nextStage; }; - srun = {nextStage, conf, preSrun ? "", ...}: ( + srun = {nextStage, conf, preSrun ? "", postSrun ? "", ...}: ( assert (assertMsg (!(conf ? cpuBind)) "cpuBind is no longer available in the standard srun stage"); stages.srun { inherit (conf) nixPrefix; - inherit nextStage preSrun; + inherit nextStage preSrun postSrun; # Binding is set to cores always cpuBind = "cores,verbose";