aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorJustin Bedo <cu@cua0.org>2019-04-18 11:33:52 +1000
committerJustin Bedo <cu@cua0.org>2019-04-18 19:00:53 +1000
commite9b908b0cdc22a7b43301e63a23c7911aa371721 (patch)
treee000b0817ca7cac0e839b98aec738cbcb7fe1649
parent0c09091762accf8716ef7853dded34e1c86497aa (diff)
sbatch: initial slurm implementation
-rw-r--r--default.nix5
-rw-r--r--lib/sbatch.nix28
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 ];
+ })