diff options
| -rw-r--r-- | lib/qsub.nix | 135 | ||||
| -rw-r--r-- | lib/slurm.nix | 42 |
2 files changed, 90 insertions, 87 deletions
diff --git a/lib/qsub.nix b/lib/qsub.nix index f5a00ec..a8827a8 100644 --- a/lib/qsub.nix +++ b/lib/qsub.nix @@ -1,74 +1,79 @@ -{stdenv, lib, writeScript}: +{ stdenv, lib, writeScript }: with lib; -{ ppn, mem, walltime, queue ? null, qsubFlags ? null, tmpDir, sleepTime, qsubPath ? "/usr/bin" }: +let escape = x: if builtins.typeOf x == "string" then escapeShellArg x else x; + +in { ppn, mem, walltime, queue ? null, qsubFlags ? null, tmpDir, sleepTime +, qsubPath ? "/usr/bin" }: drv: - let ppnReified = if drv.multicore then ppn else 1; - in lib.overrideDerivation drv ({ args, builder, name, ... }: { - builder = "/bin/bash"; - args = let - script = writeScript "qsub-script" '' - #!${stdenv.shell} - while [ ! -e ${tmpDir}/qsub-$PBS_JOBID ] ; do - sleep ${toString sleepTime} - done - set -a - . ${tmpDir}/qsub-$PBS_JOBID/nix-set - set +a - TMPDIR=${tmpDir}/qsub-$PBS_JOBID - TEMP=$TMPDIR - TMP=$TMPDIR - NIX_BUILD_TOP=$TMPDIR - cd $TMPDIR - ${builder} ${lib.escapeShellArgs args} &> qsub-log - echo $? > qsub-exit - ''; +let ppnReified = if drv.multicore then ppn else 1; +in overrideDerivation drv ({ args, builder, name, ... }: { + builder = "/bin/bash"; + args = let + script = writeScript "qsub-script" '' + #!${stdenv.shell} + while [ ! -e ${tmpDir}/qsub-$PBS_JOBID ] ; do + sleep ${toString sleepTime} + done + set -a + . ${tmpDir}/qsub-$PBS_JOBID/nix-set + set +a + TMPDIR=${tmpDir}/qsub-$PBS_JOBID + TEMP=$TMPDIR + TMP=$TMPDIR + NIX_BUILD_TOP=$TMPDIR + cd $TMPDIR + ${builder} ${concatMapStringsSep " " escape args} &> qsub-log + echo $? > qsub-exit + ''; - qsub = writeScript "qsub" '' - #!${stdenv.shell} - PATH=${qsubPath} - SHELL=/bin/sh - NIX_BUILD_CORES=${toString ppnReified} + qsub = writeScript "qsub" '' + #!${stdenv.shell} + PATH=${qsubPath} + SHELL=/bin/sh + NIX_BUILD_CORES=${toString ppnReified} - while : ; do - qsub -l nodes=1:ppn=${toString ppnReified},mem=${toString mem}gb,walltime=${walltime} \ - -N "${name}" \ - ${optionalString (queue != null) "-q ${queue}"} \ - ${optionalString (qsubFlags != null) qsubFlags} \ - ${script} 2>&1 > id - if [ $? -eq 0 ] ; then - break - fi - if ! grep "Please retry" id > /dev/null ; then - cat id >&2 - exit 1 - fi - sleep ${toString sleepTime} - done - id=$(cat id) - echo $id + while : ; do + qsub -l nodes=1:ppn=${toString ppnReified},mem=${ + toString mem + }gb,walltime=${walltime} \ + -N "${name}" \ + ${optionalString (queue != null) "-q ${queue}"} \ + ${optionalString (qsubFlags != null) qsubFlags} \ + ${script} 2>&1 > id + if [ $? -eq 0 ] ; then + break + fi + if ! grep "Please retry" id > /dev/null ; then + cat id >&2 + exit 1 + fi + sleep ${toString sleepTime} + done + id=$(cat id) + echo $id - function cleanup { - qdel $id 2>/dev/null || true - sleep ${toString sleepTime} - rm -rf ${tmpDir}/qsub-$id - } - trap cleanup INT TERM EXIT + function cleanup { + qdel $id 2>/dev/null || true + sleep ${toString sleepTime} + rm -rf ${tmpDir}/qsub-$id + } + trap cleanup INT TERM EXIT - cp -r $TMPDIR ${tmpDir}/qsub-$id - set > ${tmpDir}/qsub-$id/nix-set - until qstat -f ''${id%%.} 2>&1 | grep "\(Unknown Job\|job_state = C\)" > /dev/null ; do - sleep ${toString sleepTime} - done - cat ${tmpDir}/qsub-$id/qsub-log - if [ -e ${tmpDir}/qsub-$id/qsub-exit ]; then - exitCode=$(cat ${tmpDir}/qsub-$id/qsub-exit) - else - exitCode=1 - fi - exit $exitCode - ''; + cp -r $TMPDIR ${tmpDir}/qsub-$id + set > ${tmpDir}/qsub-$id/nix-set + until qstat -f ''${id%%.} 2>&1 | grep "\(Unknown Job\|job_state = C\)" > /dev/null ; do + sleep ${toString sleepTime} + done + cat ${tmpDir}/qsub-$id/qsub-log + if [ -e ${tmpDir}/qsub-$id/qsub-exit ]; then + exitCode=$(cat ${tmpDir}/qsub-$id/qsub-exit) + else + exitCode=1 + fi + exit $exitCode + ''; - in [ "-c" qsub ]; - }) + in [ "-c" qsub ]; +}) diff --git a/lib/slurm.nix b/lib/slurm.nix index 30a5f5e..4ad41ab 100644 --- a/lib/slurm.nix +++ b/lib/slurm.nix @@ -1,28 +1,26 @@ -{stdenv, lib, writeScript, coreutils}: +{ stdenv, lib, writeScript, coreutils }: with lib; -{ ppn, mem, walltime, partition ? null, slurmFlags ? null, salloc ? "/usr/bin/salloc", srun ? "/usr/bin/srun" }: -drv: - let ppnReified = if drv.multicore then ppn else 1; - in lib.overrideDerivation drv ({ args, builder, name, ... }: { - builder = stdenv.shell; - args = let - script = writeScript "slurm-script" '' - #!${stdenv.shell} - ${builder} ${lib.escapeShellArgs args} - ''; +let escape = x: if builtins.typeOf x == "string" then escapeShellArg x else x; - slurm = writeScript "slurm" '' - #!${stdenv.shell} - NIX_BUILD_CORES=${toString ppnReified} +in { ppn, mem, walltime, partition ? null, slurmFlags ? null +, salloc ? "/usr/bin/salloc", srun ? "/usr/bin/srun" }: +drv: +let ppnReified = if drv.multicore then ppn else 1; +in overrideDerivation drv ({ args, builder, name, ... }: { + builder = stdenv.shell; + args = let + slurm = writeScript "slurm" '' + #!${stdenv.shell} + NIX_BUILD_CORES=${toString ppnReified} - ${salloc} -c $NIX_BUILD_CORES --mem=${toString mem}G -t ${walltime} \ - -J "${name}" \ - ${optionalString (partition != null) "-p ${partition}"} \ - ${optionalString (slurmFlags != null) slurmFlags} \ - ${srun} ${script} - ''; + ${salloc} -c $NIX_BUILD_CORES --mem=${toString mem}G -t ${walltime} \ + -J "${name}" \ + ${optionalString (partition != null) "-p ${partition}"} \ + ${optionalString (slurmFlags != null) slurmFlags} \ + ${srun} ${builder} ${concatMapStringsSep " " escape args} + ''; - in [ "-c" slurm ]; - }) + in [ "-c" slurm ]; +}) |