101 lines
2.9 KiB
Nix
101 lines
2.9 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 mkPersistDir;
|
|
base = "${namespace}.services.apps.slurm";
|
|
cfg = getAttrByNamespace config base;
|
|
networkCfg = getAttrByNamespace config "${namespace}.services.networking";
|
|
in {
|
|
imports = [
|
|
./client.nix
|
|
./server.nix
|
|
];
|
|
|
|
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 = {
|
|
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
|
|
hostStrings =
|
|
cfg.controlHosts
|
|
|> map (host: "SlurmctldHost=${host}(${resolveHostIP networkCfg.devices host})")
|
|
|> concatStringsSep "\n";
|
|
in ''
|
|
${hostStrings}
|
|
GresTypes=gpu,shard
|
|
TaskPlugin=task/cgroup
|
|
DefCpuPerGPU=1
|
|
DefMemPerCPU=1000
|
|
ReturnToService=2
|
|
'';
|
|
|
|
extraConfigPaths = [(inputs.dotfiles + "/slurm/config")];
|
|
};
|
|
|
|
munge = {
|
|
enable = true;
|
|
password = config.sops.secrets."munge-key".path;
|
|
};
|
|
};
|
|
|
|
sops.secrets = let
|
|
inherit (config.users.users) munge;
|
|
in {
|
|
"munge-key" = {
|
|
sopsFile = "${inputs.self}/secrets/sops/munge.key";
|
|
format = "binary";
|
|
owner = munge.name;
|
|
group = munge.group;
|
|
mode = "0400";
|
|
};
|
|
};
|
|
|
|
${namespace}.services.storage.impermanence.folders = [
|
|
(mkPersistDir config "munge" "/var/lib/munge")
|
|
];
|
|
};
|
|
}
|