mirror of
https://github.com/CHN-beta/nixos.git
synced 2024-10-23 03:18:45 +08:00
206 lines
8.2 KiB
Nix
206 lines
8.2 KiB
Nix
inputs:
|
|
{
|
|
options.nixos.services.slurm = let inherit (inputs.lib) mkOption types; in
|
|
{
|
|
enable = mkOption { type = types.bool; default = false; };
|
|
# 本机是否为控制节点,如果不是,填写控制节点的主机名
|
|
master = mkOption { type = types.nonEmptyStr; default = inputs.config.nixos.system.networking.hostname; };
|
|
node = mkOption { type = types.attrsOf (types.submodule (submoduleInputs: { options =
|
|
{
|
|
# slurm 中使用的节点名称
|
|
name = mkOption { type = types.nonEmptyStr; };
|
|
address = mkOption { type = types.nonEmptyStr; };
|
|
cpu =
|
|
{
|
|
sockets = mkOption { type = types.ints.unsigned; default = 1; };
|
|
cores = mkOption { type = types.ints.unsigned; default = 1; };
|
|
threads = mkOption { type = types.ints.unsigned; default = 1; };
|
|
mpiThreads = mkOption { type = types.ints.unsigned; default = 1; };
|
|
openmpThreads = mkOption { type = types.ints.unsigned; default = 1; };
|
|
};
|
|
memoryMB = mkOption { type = types.ints.unsigned; default = 1024; };
|
|
gpus = mkOption { type = types.nullOr (types.attrsOf types.ints.unsigned); default = null; };
|
|
};}));};
|
|
partitions = mkOption { type = types.attrsOf (types.listOf types.nonEmptyStr); default = {}; };
|
|
defaultPartition = mkOption { type = types.nonEmptyStr; default = "localhost"; };
|
|
tui =
|
|
{
|
|
cpuMpiThreads = mkOption { type = types.ints.unsigned; default = 1; };
|
|
cpuOpenmpThreads = mkOption { type = types.ints.unsigned; default = 1; };
|
|
gpus = mkOption { type = types.nullOr (types.listOf types.nonEmptyStr); default = null; };
|
|
};
|
|
# 是否打开防火墙相应端口,对于多节点部署需要打开
|
|
setupFirewall = mkOption { type = types.bool; default = false; };
|
|
};
|
|
config = let inherit (inputs.config.nixos.services) slurm; in inputs.lib.mkIf slurm.enable (inputs.lib.mkMerge
|
|
[
|
|
# worker 配置
|
|
{
|
|
services =
|
|
{
|
|
slurm =
|
|
{
|
|
package = (inputs.pkgs.slurm.override { enableGtk2 = true; }).overrideAttrs
|
|
(prev:
|
|
let
|
|
inherit (inputs.config.nixos.system.nixpkgs) cuda;
|
|
inherit (inputs.pkgs.cudaPackages) cuda_nvml_dev;
|
|
additionalInputs = inputs.lib.optionals cuda.enable [ cuda_nvml_dev cuda_nvml_dev.lib ];
|
|
additionalFlags = inputs.lib.optional cuda.enable "-L${cuda_nvml_dev.lib}/lib/stubs";
|
|
in
|
|
{
|
|
buildInputs = prev.buildInputs or [] ++ additionalInputs;
|
|
LDFLAGS = prev.LDFLAGS or [] ++ additionalFlags;
|
|
nativeBuildInputs = prev.nativeBuildInputs ++ [ inputs.pkgs.wrapGAppsHook ];
|
|
postInstall =
|
|
''
|
|
pushd contribs/pmi2
|
|
make install
|
|
popd
|
|
pushd contribs/pmi
|
|
make install
|
|
popd
|
|
'' + prev.postInstall;
|
|
}
|
|
);
|
|
client.enable = true;
|
|
nodeName = builtins.map
|
|
(node:
|
|
let gpuString =
|
|
if node.value.gpus == null then ""
|
|
else "Gres=" + builtins.concatStringsSep "," (builtins.map
|
|
(gpu: "gpu:${gpu.name}:${builtins.toString gpu.value}")
|
|
(inputs.lib.attrsToList node.value.gpus));
|
|
in builtins.concatStringsSep " "
|
|
[
|
|
node.value.name
|
|
"NodeHostname=${node.name}"
|
|
"NodeAddr=${node.value.address}"
|
|
"RealMemory=${builtins.toString node.value.memoryMB}"
|
|
"Sockets=${builtins.toString node.value.cpu.sockets}"
|
|
"CoresPerSocket=${builtins.toString node.value.cpu.cores}"
|
|
"ThreadsPerCore=${builtins.toString node.value.cpu.threads}"
|
|
"${gpuString}"
|
|
"State=UNKNOWN"
|
|
])
|
|
(inputs.localLib.attrsToList slurm.node);
|
|
partitionName = builtins.map
|
|
(partition:
|
|
let nodes = builtins.concatStringsSep "," partition.value;
|
|
in builtins.concatStringsSep " "
|
|
[
|
|
partition.name
|
|
"Nodes=${builtins.concatStringsSep "," (builtins.map (n: slurm.node.${n}.name) partition.value)}"
|
|
"Default=${if partition.name == slurm.defaultPartition then "YES" else "NO"}"
|
|
"MaxTime=INFINITE"
|
|
"State=UP"
|
|
])
|
|
(inputs.localLib.attrsToList slurm.partitions);
|
|
procTrackType = "proctrack/cgroup";
|
|
controlMachine = slurm.master;
|
|
controlAddr = slurm.node.${slurm.master}.address;
|
|
extraConfig =
|
|
''
|
|
SelectType=select/cons_tres
|
|
SelectTypeParameters=CR_Core
|
|
GresTypes=gpu
|
|
DefCpuPerGPU=1
|
|
|
|
TaskProlog=${inputs.pkgs.writeShellScript "set_env" "echo export CUDA_DEVICE_ORDER=PCI_BUS_ID"}
|
|
|
|
AccountingStorageType=accounting_storage/slurmdbd
|
|
AccountingStorageHost=localhost
|
|
AccountingStoreFlags=job_comment,job_env,job_extra,job_script
|
|
|
|
JobCompType=jobcomp/filetxt
|
|
JobCompLoc=/var/log/slurmctld/jobcomp.log
|
|
|
|
SchedulerParameters=enable_user_top
|
|
|
|
SlurmdDebug=debug2
|
|
DebugFlags=NO_CONF_HASH
|
|
|
|
# automatically resume node after drain
|
|
ReturnToService=2
|
|
'';
|
|
extraConfigPaths =
|
|
let gpus = slurm.node.${inputs.config.nixos.system.networking.hostname}.gpus or null;
|
|
in inputs.lib.mkIf (gpus != null)
|
|
(
|
|
let gpuString = builtins.concatStringsSep "\n" (builtins.map
|
|
(gpu: "Name=gpu Type=${gpu.name} Count=${builtins.toString gpu.value}")
|
|
(inputs.localLib.attrsToList gpus));
|
|
in [(inputs.pkgs.writeTextDir "gres.conf" "AutoDetect=nvml\n${gpuString}")]
|
|
);
|
|
};
|
|
munge = { enable = true; password = inputs.config.sops.secrets."munge.key".path; };
|
|
};
|
|
systemd =
|
|
{
|
|
services.slurmd.environment =
|
|
let gpus = slurm.node.${inputs.config.nixos.system.networking.hostname}.gpus or null;
|
|
in inputs.lib.mkIf (gpus != null)
|
|
{
|
|
CUDA_PATH = "${inputs.pkgs.cudatoolkit}";
|
|
LD_LIBRARY_PATH = "${inputs.config.hardware.nvidia.package}/lib";
|
|
};
|
|
};
|
|
sops.secrets."munge.key" =
|
|
{
|
|
format = "binary";
|
|
sopsFile = "${builtins.dirOf inputs.config.sops.defaultSopsFile}/munge.key";
|
|
owner = inputs.config.systemd.services.munged.serviceConfig.User;
|
|
};
|
|
networking.firewall =
|
|
let config = inputs.lib.mkIf slurm.setupFirewall [ 6818 ];
|
|
in { allowedTCPPorts = config; allowedUDPPorts = config; };
|
|
}
|
|
# master 配置
|
|
(inputs.lib.mkIf (slurm.master == inputs.config.nixos.system.networking.hostname)
|
|
{
|
|
services.slurm =
|
|
{
|
|
server.enable = true;
|
|
dbdserver =
|
|
{
|
|
enable = true;
|
|
dbdHost = "localhost";
|
|
storagePassFile = inputs.config.sops.secrets."slurm/db".path;
|
|
extraConfig =
|
|
''
|
|
StorageHost=*
|
|
StorageLoc=slurm
|
|
'';
|
|
};
|
|
};
|
|
systemd.tmpfiles.rules = [ "d /var/log/slurmctld 700 slurm slurm" ];
|
|
sops.secrets."slurm/db" = { owner = "slurm"; key = "mariadb/slurm"; };
|
|
nixos =
|
|
{
|
|
packages.packages._packages = [ inputs.pkgs.localPackages.sbatch-tui ];
|
|
user.sharedModules = [{ home.packages =
|
|
[
|
|
(inputs.pkgs.writeShellScriptBin "sbatch"
|
|
''
|
|
if [ "$#" -eq 0 ]; then
|
|
sbatch-tui
|
|
else
|
|
/run/current-system/sw/bin/sbatch "$@"
|
|
fi
|
|
'')
|
|
];}];
|
|
services.mariadb = { enable = true; instances.slurm = {}; };
|
|
};
|
|
environment.etc."sbatch-tui.yaml".text = builtins.toJSON
|
|
{
|
|
CpuMpiThreads = slurm.tui.cpuMpiThreads;
|
|
CpuOpenmpThreads = slurm.tui.cpuOpenmpThreads;
|
|
GpuIds = slurm.tui.gpus;
|
|
};
|
|
networking.firewall =
|
|
let config = inputs.lib.mkIf slurm.setupFirewall [ 6817 ];
|
|
in { allowedTCPPorts = config; allowedUDPPorts = config; };
|
|
})
|
|
]);
|
|
}
|