From b8d8c1f207a8c80f7267920efa70db785e5d441e Mon Sep 17 00:00:00 2001
From: abysssol <abysssol@pm.me>
Date: Thu, 22 Feb 2024 23:41:25 -0500
Subject: [PATCH] nixos/ollama: add option for hardware acceleration

---
 nixos/modules/services/misc/ollama.nix | 33 ++++++++++++++++++--------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/nixos/modules/services/misc/ollama.nix b/nixos/modules/services/misc/ollama.nix
index d9359d2b5cd4..0d3574a2bac0 100644
--- a/nixos/modules/services/misc/ollama.nix
+++ b/nixos/modules/services/misc/ollama.nix
@@ -1,9 +1,14 @@
-{ config, lib, pkgs, ... }: let
+{ config, lib, pkgs, ... }:
+let
+  inherit (lib.types) nullOr enum;
 
   cfg = config.services.ollama;
-
-in {
-
+  ollamaPackage = cfg.package.override {
+    inherit (cfg) acceleration;
+    linuxPackages.nvidia_x11 = config.hardware.nvidia.package;
+  };
+in
+{
   options = {
     services.ollama = {
       enable = lib.mkEnableOption (
@@ -16,12 +21,22 @@ in {
           Specifies the bind address on which the ollama server HTTP interface listens.
         '';
       };
+      acceleration = lib.mkOption {
+        type = nullOr (enum [ "rocm" "cuda" ]);
+        default = null;
+        example = "rocm";
+        description = lib.mdDoc ''
+          Specifies the interface to use for hardware acceleration.
+
+          - `rocm`: supported by modern AMD GPUs
+          - `cuda`: supported by modern NVIDIA GPUs
+        '';
+      };
       package = lib.mkPackageOption pkgs "ollama" { };
     };
   };
 
   config = lib.mkIf cfg.enable {
-
     systemd = {
       services.ollama = {
         wantedBy = [ "multi-user.target" ];
@@ -33,7 +48,7 @@ in {
           OLLAMA_HOST = cfg.listenAddress;
         };
         serviceConfig = {
-          ExecStart = "${lib.getExe cfg.package} serve";
+          ExecStart = "${lib.getExe ollamaPackage} serve";
           WorkingDirectory = "/var/lib/ollama";
           StateDirectory = [ "ollama" ];
           DynamicUser = true;
@@ -41,10 +56,8 @@ in {
       };
     };
 
-    environment.systemPackages = [ cfg.package ];
-
+    environment.systemPackages = [ ollamaPackage ];
   };
 
-  meta.maintainers = with lib.maintainers; [ onny ];
-
+  meta.maintainers = with lib.maintainers; [ abysssol onny ];
 }