diff --git a/modules/nixos/services/ci/slurm/default.nix b/modules/nixos/services/ci/slurm/default.nix index f57ff6e..aaaecef 100644 --- a/modules/nixos/services/ci/slurm/default.nix +++ b/modules/nixos/services/ci/slurm/default.nix @@ -5,7 +5,7 @@ namespace, ... }: let - inherit (lib) types mkEnableOption mkOption mkIf groupBy mapAttrsToList attrNames mapAttrs concatStringsSep flatten; + inherit (lib) types mkEnableOption mkIf groupBy mapAttrsToList attrNames mapAttrs concatStringsSep flatten; inherit (lib.${namespace}) getAttrByNamespace mkOptionsWithNamespace resolveHostIP mkOpt mkOptAttrset mkListOpt; inherit (config.networking) hostName; base = "${namespace}.services.ci.slurm"; @@ -25,78 +25,91 @@ in { }; config = mkIf cfg.enable { - services.slurm = { - client.enable = builtins.hasAttr hostName cfg.nodeMap; - server.enable = builtins.elem hostName cfg.controlHosts; + services = { + slurm = { + client.enable = builtins.hasAttr hostName cfg.nodeMap; + server.enable = builtins.elem hostName cfg.controlHosts; - stateSaveLocation = "/mnt/nfs/slurm"; + stateSaveLocation = "/mnt/nfs/slurm"; - nodeName = let - mkNodeConfig = node: info: "${node} NodeAddr=${resolveHostIP networkCfg.devices node} ${info.configString} State=UNKNOWN"; - in - mapAttrsToList mkNodeConfig cfg.nodeMap; + nodeName = let + mkNodeConfig = node: info: "${node} NodeAddr=${resolveHostIP networkCfg.devices node} ${info.configString} State=UNKNOWN"; + in + mapAttrsToList mkNodeConfig cfg.nodeMap; - partitionName = let - mkPartitions = nodeMap: - nodeMap - |> attrNames - |> map ( - node: - (cfg.nodeMap.${node}.partitions) - |> map (partition: {inherit partition node;}) - ) - |> flatten - |> groupBy (x: x.partition) - |> mapAttrs (name: entries: entries |> map (e: e.node)); + partitionName = let + mkPartitions = nodeMap: + nodeMap + |> attrNames + |> map ( + node: + (cfg.nodeMap.${node}.partitions) + |> map (partition: {inherit partition node;}) + ) + |> flatten + |> groupBy (x: x.partition) + |> mapAttrs (name: entries: entries |> map (e: e.node)); - formatPartition = name: nodes: "${name} Nodes=${nodes |> concatStringsSep ","} Default=${ - if name == "main" - then "YES" - else "NO" - } MaxTime=INFINITE State=UP"; - in - cfg.nodeMap - |> mkPartitions - |> mapAttrsToList formatPartition; + formatPartition = name: nodes: "${name} Nodes=${nodes |> concatStringsSep ","} Default=${ + if name == "main" + then "YES" + else "NO" + } MaxTime=INFINITE State=UP"; + in + cfg.nodeMap + |> mkPartitions + |> mapAttrsToList formatPartition; - extraConfig = let - mkHostString = host: "SlurmctldHost=${host}(${resolveHostIP networkCfg.devices host})"; - hostStrings = - cfg.controlHosts - |> map mkHostString - |> concatStringsSep "\n"; - in '' - ${hostStrings} - GresTypes=gpu,shard - TaskPlugin=task/cgroup - SlurmdParameters=allow_ecores - DefCpuPerGPU=1 - DefMemPerCPU=1000 - ReturnToService=2 + extraConfig = let + mkHostString = host: "SlurmctldHost=${host}(${resolveHostIP networkCfg.devices host})"; + hostStrings = + cfg.controlHosts + |> map mkHostString + |> concatStringsSep "\n"; + in '' + ${hostStrings} + GresTypes=gpu,shard + TaskPlugin=task/cgroup + SlurmdParameters=allow_ecores + DefCpuPerGPU=1 + DefMemPerCPU=1000 + ReturnToService=2 - TaskProlog=${inputs.dotfiles + "/slurm/prolog.sh"} - TaskEpilog=${inputs.dotfiles + "/slurm/epilog.sh"} - ''; + TaskProlog=${inputs.dotfiles + "/slurm/prolog.sh"} + TaskEpilog=${inputs.dotfiles + "/slurm/epilog.sh"} + ''; - extraCgroupConfig = - '' - ConstrainCores=yes - ConstrainDevices=yes - ConstrainRAMSpace=yes - '' - + ( - if hostName != "chibi" - then '' - ConstrainSwapSpace=yes - AllowedSwapSpace=0 + extraCgroupConfig = '' - else '''' - ); + ConstrainCores=yes + ConstrainDevices=yes + ConstrainRAMSpace=yes + '' + + ( + if hostName != "chibi" + then '' + ConstrainSwapSpace=yes + AllowedSwapSpace=0 + '' + else '''' + ); - extraConfigPaths = [(inputs.dotfiles + "/slurm/config")]; + extraConfigPaths = [(inputs.dotfiles + "/slurm/config")]; + }; + + munge.enable = true; + }; + + systemd.services.slurmctld = mkIf (builtins.elem hostName cfg.controlHosts) { + requires = [ + "mnt-nfs-slurm.mount" + ]; + + after = [ + "mnt-nfs-slurm.mount" + ]; }; - services.munge.enable = true; environment.etc."munge/munge.key" = { source = "${inputs.self}/secrets/crypt/munge.key"; user = "munge"; @@ -104,6 +117,9 @@ in { mode = "0400"; }; - ${namespace}.services.storage.impermanence.folders = ["/var/spool/slurmctld" "/var/spool/slurmd"]; + ${namespace}.services.storage.impermanence.folders = [ + "/var/spool/slurmctld" + "/var/spool/slurmd" + ]; }; }