93 lines
3.6 KiB
Nix
93 lines
3.6 KiB
Nix
{ config, pkgs, lib, ... }:
|
|
{
|
|
options = {
|
|
nix = {
|
|
gpuSupport = lib.mkEnableOption "Enable GPU support";
|
|
gpuCount = lib.mkOption {
|
|
type = lib.types.ints.unsigned;
|
|
default = 1;
|
|
description = "Number of GPUs to support";
|
|
};
|
|
|
|
buildUser = lib.mkOption {
|
|
type = lib.types.bool;
|
|
default = config.nix.gpuSupport;
|
|
description = "Create a user for building";
|
|
};
|
|
|
|
builderPublicKey = lib.mkOption {
|
|
type = lib.types.nullOr lib.types.str;
|
|
default = null;
|
|
description = "Public key for the builder user";
|
|
};
|
|
};
|
|
};
|
|
|
|
config = lib.mkIf (config.nix.gpuSupport || config.nix.buildUser) {
|
|
nixpkgs.config.allowUnfree = true;
|
|
security.wrappers.curl = {
|
|
setuid = true;
|
|
owner = "root";
|
|
group = "root";
|
|
source = lib.getExe pkgs.curl;
|
|
};
|
|
|
|
services.openssh = lib.mkIf config.nix.buildUser {
|
|
enable = true;
|
|
openFirewall = lib.mkDefault true;
|
|
extraConfig = lib.mkIf config.nix.buildUser ''
|
|
Match User ${config.users.users.builder.name}
|
|
AuthorizedKeysCommand /run/wrappers/bin/${config.security.wrappers.curl.program} https://servers.wavelens.io/authorized-keys
|
|
AuthorizedKeysCommandUser nobody
|
|
'';
|
|
};
|
|
|
|
security.pam.services.sshd.rules.session.register-login = {
|
|
enable = true;
|
|
order = 14000;
|
|
control = "required";
|
|
modulePath = "pam_exec.so";
|
|
args = let
|
|
configFile = pkgs.writeScript "register-login.sh" ''
|
|
#!/bin/sh
|
|
|
|
if [ "$PAM_USER" != "builder" ]; then
|
|
exit 0
|
|
fi
|
|
|
|
if [ "$PAM_TYPE" != "open_session" ]; then
|
|
publickey=$(${pkgs.systemd}/bin/journalctl -u ssh* -n 10000 | ${lib.getExe pkgs.gnugrep} "Accepted publickey for builder from $PAM_RHOST" | tail -n 1 | ${lib.getExe pkgs.gawk} '{print $16}')
|
|
${lib.getExe pkgs.curl} -sX POST -o /dev/null -H "Content-Type: application/json" -d "{\"publickey\":\"$publickey\",\"login\":false}" https://servers.wavelens.io/session
|
|
else
|
|
${lib.getExe pkgs.curl} -sX POST -o /dev/null -H "Content-Type: application/json" -d "{\"login\":true}" https://servers.wavelens.io/session
|
|
fi
|
|
'';
|
|
in [(toString configFile)];
|
|
};
|
|
|
|
users.users."builder" = lib.mkIf config.nix.buildUser {
|
|
group = "nogroup";
|
|
name = "builder";
|
|
isNormalUser = true;
|
|
openssh.authorizedKeys.keys = lib.optional (config.nix.builderPublicKey != null) config.nix.builderPublicKey;
|
|
};
|
|
|
|
nix.settings = {
|
|
max-jobs = lib.mkIf config.nix.gpuSupport (lib.mkDefault 1);
|
|
builders-use-substitutes = true;
|
|
experimental-features = [ "nix-command" "flakes" "ca-derivations" ];
|
|
trusted-users = [ "root" "@wheel" ] ++ lib.optional config.nix.buildUser "builder";
|
|
substituters = [ "https://attic.wavelens.io/main?priority=5&want-mass-query=true" ];
|
|
trusted-public-keys = [ "main:3VVGDhOgY/x5hn7XIkVhqjEjHvOnU7o1cPlrWv91Mko=" ];
|
|
system-features = let
|
|
gpu = [ "A4000-1" "L40-1" "A6000-1" "A6000-4" "A100-1" "A100-2" "A100-4" "A100-8" "H100-1" "H100-2" "H100-4" "H100-8" ];
|
|
in lib.optionals config.nix.gpuSupport ([ "cuda" "service" "ca-derivations" ] ++ (map (v: "gpu-${v}") gpu));
|
|
extra-sandbox-paths = lib.mkIf (config.nix.gpuSupport && config.nix.settings.sandbox) ([
|
|
"/dev/nvidiactl"
|
|
"/dev/nvidia-uvm"
|
|
"/run/opengl-driver=${config.boot.kernelPackages.nvidia_x11}"
|
|
] ++ (map (i: "/dev/nvidia${toString i}") (lib.range 0 (config.nix.gpuCount - 1))));
|
|
};
|
|
};
|
|
}
|