blob: efb765431c8d962e482fcb519b007bf78acecab6 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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 ];
})
|
