131 lines
3.8 KiB
Nix
131 lines
3.8 KiB
Nix
{
|
|
config,
|
|
inputs,
|
|
lib,
|
|
namespace,
|
|
...
|
|
}: let
|
|
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.apps.slurm";
|
|
cfg = getAttrByNamespace config base;
|
|
networkCfg = getAttrByNamespace config "${namespace}.services.networking";
|
|
in {
|
|
options = with types;
|
|
mkOptionsWithNamespace base {
|
|
enable = mkEnableOption "SLURM";
|
|
controlHosts = mkListOpt str [] "Device to use for control hosts";
|
|
nodeMap = mkOptAttrset (submodule {
|
|
options = {
|
|
partitions = mkListOpt str [] "List of partitions for the node";
|
|
configString = mkOpt str "" "Configuration string for the node capabilities";
|
|
};
|
|
}) {} "Mapping of node device defintitions to IPs and device configurations";
|
|
};
|
|
|
|
config = mkIf cfg.enable {
|
|
services = {
|
|
slurm = {
|
|
client.enable = builtins.hasAttr hostName cfg.nodeMap;
|
|
server.enable = builtins.elem hostName cfg.controlHosts;
|
|
|
|
stateSaveLocation = "/mnt/nfs/slurm";
|
|
|
|
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));
|
|
|
|
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
|
|
|
|
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
|
|
''
|
|
else ''''
|
|
);
|
|
|
|
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"
|
|
];
|
|
};
|
|
|
|
sops.secrets = let
|
|
inherit (config.users.users) munge;
|
|
in {
|
|
"munge-key" = {
|
|
sopsFile = "${inputs.self}/secrets/sops/munge.key";
|
|
format = "binary";
|
|
path = "/etc/munge/munge.key";
|
|
owner = munge.name;
|
|
group = munge.group;
|
|
mode = "0400";
|
|
};
|
|
};
|
|
|
|
${namespace}.services.storage.impermanence.folders = [
|
|
"/var/spool/slurmctld"
|
|
"/var/spool/slurmd"
|
|
];
|
|
};
|
|
}
|