aboutsummaryrefslogtreecommitdiff
path: root/lib/slurm.nix
blob: 8c9ad5077ca80f5a27b465ceff2bc40ec6fedb5d (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
29
30
31
32
33
34
35
36
37
{ stdenv, lib, writeScript, coreutils }:

with lib;

let escape = x: if builtins.typeOf x == "string" then escapeShellArg x else x;

in
{ ppn
, mem
, walltime
, partition ? null
, slurmFlags ? null
, salloc ? "/usr/bin/salloc"
, srun ? "/usr/bin/srun"
, ...
}:
drv:
let ppnReified = if drv.multicore then ppn else 1;
in
overrideDerivation drv ({ args, builder, name, ... }: {
  builder = stdenv.shell;
  args =
    let
      slurm = writeScript "slurm" ''
        #!${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} \
          ${srun} ${builder} ${concatMapStringsSep " " escape args}
      '';

    in
    [ "-c" slurm ];
})