From 7e5542eadad1db9271fe1d18b5682eacebadbf48 Mon Sep 17 00:00:00 2001
From: Justin Bedo <cu@cua0.org>
Date: Tue, 23 Jun 2020 17:28:00 +1000
Subject: qsub,slurm: fix string escaping bug
escapeShellArgs calls toString on the passed item, which for paths
converts it to a string containing the full path outside the nix
store. Avoid this by calling escapeShellArg on strings only.
---
lib/qsub.nix | 135 ++++++++++++++++++++++++++++++----------------------------
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 ];
+})
--
cgit v1.2.3