diff options
| -rw-r--r-- | default.nix | 5 | ||||
| -rw-r--r-- | lib/sbatch.nix | 28 |
2 files changed, 33 insertions, 0 deletions
diff --git a/default.nix b/default.nix index 604badb..608fa9b 100644 --- a/default.nix +++ b/default.nix @@ -34,6 +34,11 @@ let strelka = callBionix ./tools/strelka.nix {}; ascat = callBionix ./tools/ascat.nix {}; + sbatch = attrs: bionix.extend (self: super: with self; rec { + sbatchDefs = { ppn = 1; mem = 1; walltime = "24:00:00"; partition = null; slurmFlags = null; } // attrs; + sbatch = attrs: (callPackage ./lib/sbatch.nix {}) (sbatchDefs // attrs); + exec = f: x: y: sbatch (builtins.intersectAttrs sbatchDefs x) (f (builtins.removeAttrs x (builtins.attrNames sbatchDefs)) y); + }); qsub = attrs: bionix.extend (self: super: with self; rec { qsubDefs = { ppn = 1; mem = 1; walltime = "24:00:00"; tmpDir = "/tmp"; sleepTime = 60; queue = null; qsubFlags = null; } // attrs; qsub = attrs: (callPackage ./lib/qsub.nix {}) (qsubDefs // attrs); diff --git a/lib/sbatch.nix b/lib/sbatch.nix new file mode 100644 index 0000000..efb7654 --- /dev/null +++ b/lib/sbatch.nix @@ -0,0 +1,28 @@ +{stdenv, lib, writeScript, coreutils}: + +with lib; + +{ ppn, mem, walltime, partition ? null, slurmFlags ? null}: +drv: + let ppnReified = if drv.multicore then ppn else 1; + in lib.overrideDerivation drv ({ args, builder, name, ... }: { + builder = stdenv.shell; + args = let + script = writeScript "sbatch-script" '' + #!${stdenv.shell} + ${builder} ${lib.escapeShellArgs args} + ''; + + sbatch = writeScript "sbatch" '' + #!${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} \ + ${script} + ''; + + in [ "-c" sbatch ]; + }) |