From b0af9b86080d3e6d62647c2ad4f267bad92a89e9 Mon Sep 17 00:00:00 2001
From: Rodrigo Arias Mallo <rodrigo.arias@bsc.es>
Date: Mon, 12 Apr 2021 17:41:59 +0200
Subject: [PATCH] srun: add postSrun hook

---
 garlic/stages/srun.nix | 6 +++++-
 garlic/stdexp.nix      | 4 ++--
 2 files changed, 7 insertions(+), 3 deletions(-)

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