srun: add postSrun hook

This commit is contained in:
Rodrigo Arias Mallo 2021-04-12 17:41:59 +02:00
parent 4afda7dbfb
commit b0af9b8608
2 changed files with 7 additions and 3 deletions

View File

@ -8,6 +8,7 @@
, cpuBind , cpuBind
, nixPrefix , nixPrefix
, preSrun ? "" , preSrun ? ""
, postSrun ? ""
, srunOptions ? "" , srunOptions ? ""
, output ? "stdout.log" , output ? "stdout.log"
, error ? "stderr.log" , error ? "stderr.log"
@ -26,14 +27,17 @@ stdenv.mkDerivation rec {
${preSrun} ${preSrun}
exec ${slurm}/bin/srun \ ${slurm}/bin/srun \
--mpi=pmi2 \ --mpi=pmi2 \
--cpu-bind=${cpuBind} \ --cpu-bind=${cpuBind} \
--output=${output} \ --output=${output} \
--error=${error} \ --error=${error} \
${srunOptions} \ ${srunOptions} \
${nixPrefix}${stageProgram nextStage} ${nixPrefix}${stageProgram nextStage}
${postSrun}
EOF EOF
chmod +x $out chmod +x $out
''; '';
} }

View File

@ -76,12 +76,12 @@ rec {
inherit nextStage; inherit nextStage;
}; };
srun = {nextStage, conf, preSrun ? "", ...}: ( srun = {nextStage, conf, preSrun ? "", postSrun ? "", ...}: (
assert (assertMsg (!(conf ? cpuBind)) assert (assertMsg (!(conf ? cpuBind))
"cpuBind is no longer available in the standard srun stage"); "cpuBind is no longer available in the standard srun stage");
stages.srun { stages.srun {
inherit (conf) nixPrefix; inherit (conf) nixPrefix;
inherit nextStage preSrun; inherit nextStage preSrun postSrun;
# Binding is set to cores always # Binding is set to cores always
cpuBind = "cores,verbose"; cpuBind = "cores,verbose";