diff options
author | Justin Bedo <cu@cua0.org> | 2019-04-18 11:33:52 +1000 |
---|---|---|
committer | Justin Bedo <cu@cua0.org> | 2019-04-18 19:00:53 +1000 |
commit | e9b908b0cdc22a7b43301e63a23c7911aa371721 (patch) | |
tree | e000b0817ca7cac0e839b98aec738cbcb7fe1649 /lib | |
parent | 0c09091762accf8716ef7853dded34e1c86497aa (diff) |
sbatch: initial slurm implementation
Diffstat (limited to 'lib')
-rw-r--r-- | lib/sbatch.nix | 28 |
1 files changed, 28 insertions, 0 deletions
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 ]; + }) |