Allow changing tmpdir for qsub

1 files changed, 15 insertions, 17 deletions

diff --git a/lib/qsub.nix b/lib/qsub.nix
index ddff8fe..7c02144 100644
--- a/lib/qsub.nix
+++ b/lib/qsub.nix
@@ -1,26 +1,23 @@
 {stdenv, lib, writeScript}:
 
-{ ppn ? 1, mem ? 1, walltime ? "24:00:00" }: drv: lib.overrideDerivation drv ({ args, builder, ... }: {
+{ ppn ? 1, mem ? 1, walltime ? "24:00:00", tmpDir ? "/tmp" }: drv: lib.overrideDerivation drv ({ args, builder, ... }: {
   builder = "/bin/bash";
   args = let
     script = writeScript "qsub-script" ''
       #!${stdenv.shell}
-      while [ ! -e /stornext/HPCScratch/$PBS_JOBID ] ; do
+      while [ ! -e ${tmpDir}/$PBS_JOBID ] ; do
         sleep 5
       done
-      set -a
-      . /stornext/HPCScratch/$PBS_JOBID
-      set +a
-      TMPDIR=/tmp/$PBS_JOBID
+      TMPDIR=${tmpDir}/$PBS_JOBID
       TEMP=$TMPDIR
       TMP=$TMPDIR
       NIX_BUILD_TOP=$TMPDIR
-      mkdir $TMPDIR
       cd $TMPDIR
-      rm /stornext/HPCScratch/$PBS_JOBID
-      ${builder} ${lib.escapeShellArgs args} && touch /stornext/HPCScratch/$PBS_JOBID
-      cd /
-      rm -rf $TMPDIR
+      set -a
+      . nix-set
+      set +a
+      ${builder} ${lib.escapeShellArgs args} > qsub-stdout 2> qsub-stderr
+      echo $? > qsub-exit
     '';
 
     qsub = writeScript "qsub" ''
@@ -29,15 +26,16 @@
       SHELL=/bin/sh
       NIX_BUILD_CORES=${toString ppn}
       id=$(qsub -l nodes=1:ppn=${toString ppn},mem=${toString mem}gb,walltime=${walltime} ${script})
-      set > /stornext/HPCScratch/$id
+      cp -r $TMPDIR ${tmpDir}/$id
+      set > ${tmpDir}/$id/nix-set
       while qstat ''${id%%.} 2> /dev/null > /dev/null ; do
         sleep 5
       done
-      if [[ -e /stornext/HPCScratch/$id ]] ; then
-        rm /stornext/HPCScratch/$id
-        exit 0
-      fi
-      exit 1
+      cat ${tmpDir}/$id/qsub-stderr >&2
+      cat ${tmpDir}/$id/qsub-stdout
+      exitCode=$(cat ${tmpDir}/$id/qsub-exit)
+      rm -rf ${tmpDir}/$id
+      exit $exitCode
     '';
 
     in [ "-c" qsub ];