srun: add postSrun hook

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
   '';
 }

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";