From e9b908b0cdc22a7b43301e63a23c7911aa371721 Mon Sep 17 00:00:00 2001
From: Justin Bedo <cu@cua0.org>
Date: Thu, 18 Apr 2019 11:33:52 +1000
Subject: sbatch: initial slurm implementation

---
 lib/sbatch.nix | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 lib/sbatch.nix

(limited to 'lib')

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 ];
+  })
-- 
cgit v1.2.3