srun: add postSrun hook
This commit is contained in:
parent
4afda7dbfb
commit
b0af9b8608
@ -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
|
||||||
'';
|
'';
|
||||||
}
|
}
|
||||||
|
@ -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";
|
||||||
|
Loading…
Reference in New Issue
Block a user