From e291ba6f89be9b064ea4487457df8ee6008b8e86 Mon Sep 17 00:00:00 2001
From: Max <max@privatevoid.net>
Date: Fri, 3 Mar 2023 15:47:07 +0100
Subject: [PATCH] koboldai: add NixOS modules

---
 projects/koboldai/default.nix       | 25 ++++++++-
 projects/koboldai/nixos/amd.nix     | 12 ++++
 projects/koboldai/nixos/default.nix | 87 +++++++++++++++++++++++++++++
 3 files changed, 123 insertions(+), 1 deletion(-)
 create mode 100644 projects/koboldai/nixos/amd.nix
 create mode 100644 projects/koboldai/nixos/default.nix

diff --git a/projects/koboldai/default.nix b/projects/koboldai/default.nix
index 97a683c..eb3f7f6 100644
--- a/projects/koboldai/default.nix
+++ b/projects/koboldai/default.nix
@@ -1,4 +1,4 @@
-{ inputs, lib, ... }:
+{ config, inputs, lib, withSystem, ... }:
 
 {
   perSystem = { config, pkgs, ... }: let
@@ -17,4 +17,27 @@
       };
     };
   };
+
+  flake.nixosModules = let
+    packageModule = pkgAttrName: { pkgs, ... }: {
+      services.koboldai.package = withSystem pkgs.system (
+        { config, ... }: lib.mkOptionDefault config.packages.${pkgAttrName}
+      );
+    };
+  in {
+    koboldai = ./nixos;
+    koboldai-amd = {
+      imports = [
+        config.flake.nixosModules.koboldai
+        ./nixos/amd.nix
+        (packageModule "koboldai-amd")
+      ];
+    };
+    koboldai-nvidia = {
+      imports = [
+        config.flake.nixosModules.koboldai
+        (packageModule "koboldai-nvidia")
+      ];
+    };
+  };
 }
diff --git a/projects/koboldai/nixos/amd.nix b/projects/koboldai/nixos/amd.nix
new file mode 100644
index 0000000..0d3ce7c
--- /dev/null
+++ b/projects/koboldai/nixos/amd.nix
@@ -0,0 +1,12 @@
+{ pkgs, ... }:
+
+{
+  systemd = {
+    # Allow "unsupported" AMD GPUs
+    services.koboldai.environment.HSA_OVERRIDE_GFX_VERSION = "10.3.0";
+    # HACK: The PyTorch build we use on ROCm wants this to exist
+    tmpfiles.rules = [
+      "L+ /opt/amdgpu - - - - ${pkgs.libdrm}"
+    ];
+  };
+}
diff --git a/projects/koboldai/nixos/default.nix b/projects/koboldai/nixos/default.nix
new file mode 100644
index 0000000..a3956d2
--- /dev/null
+++ b/projects/koboldai/nixos/default.nix
@@ -0,0 +1,87 @@
+{ config, lib, ... }:
+
+let
+  inherit (lib)
+    mkIf mkOption mkEnableOption types
+    escapeShellArgs getExe optional
+  ;
+
+  cfg = config.services.koboldai;
+in
+
+{
+  options.services.koboldai= {
+    enable = mkEnableOption "KoboldAI Web UI";
+
+    package = mkOption {
+      description = "Which KoboldAI package to use.";
+      type = types.package;
+    };
+
+    user = mkOption {
+      description = "Which user to run KoboldAI as.";
+      default = "koboldai";
+      type = types.str;
+    };
+
+    group = mkOption {
+      description = "Which group to run KoboldAI as.";
+      default = "koboldai";
+      type = types.str;
+    };
+
+    host = mkOption {
+      description = "Whether to make KoboldAI remotely accessible.";
+      default = false;
+      type = types.bool;
+    };
+
+    port = mkOption {
+      description = "Which port to listen on.";
+      default = 5000;
+      type = types.port;
+    };
+
+    dataDir = mkOption {
+      description = "Where to store KoboldAI's state.";
+      default = "/var/lib/koboldai";
+      type = types.path;
+    };
+
+    extraArgs = mkOption {
+      description = "Extra command line arguments.";
+      default = [];
+      type = with types; listOf str;
+    };
+  };
+
+  config = let
+    cliArgs = (optional cfg.host "--host") ++ [
+      "--port" cfg.port
+    ] ++ cfg.extraArgs;
+  in mkIf cfg.enable {
+    users.users = mkIf (cfg.user == "koboldai") {
+      koboldai = {
+        isSystemUser = true;
+        inherit (cfg) group;
+      };
+    };
+    users.groups = mkIf (cfg.group == "koboldai") {
+      koboldai = {};
+    };
+    systemd.services.koboldai = {
+      after = [ "network.target" ];
+      wantedBy = [ "multi-user.target" ];
+      environment.HOME = cfg.dataDir;
+      serviceConfig = {
+        User = cfg.user;
+        Group = cfg.group;
+        ExecStart = "${getExe cfg.package} ${escapeShellArgs cliArgs}";
+        PrivateTmp = true;
+      };
+    };
+    systemd.tmpfiles.rules = [
+      "d '${cfg.dataDir}' 0755 ${cfg.user} ${cfg.group} - -"
+    ];
+  };
+}