nixos/ollama: add option for hardware acceleration

(cherry picked from commits:
b8d8c1f207a8c80f7267920efa70db785e5d441e
5c143f03663eb59a7a1eac4b24b7c034abc4f483
87203977204d1c3a7c7ccd39147b17dadf3156e8)

1 files changed, 26 insertions, 11 deletions

diff --git a/nixos/modules/services/misc/ollama.nix b/nixos/modules/services/misc/ollama.nix
index d9359d2b5cd44..3ac3beb4de078 100644
--- a/nixos/modules/services/misc/ollama.nix
+++ b/nixos/modules/services/misc/ollama.nix
@@ -1,27 +1,44 @@
-{ config, lib, pkgs, ... }: let
+{ config, lib, pkgs, ... }:
+let
+  inherit (lib) types;
 
   cfg = config.services.ollama;
-
-in {
-
+  ollamaPackage = cfg.package.override {
+    inherit (cfg) acceleration;
+    linuxPackages = config.boot.kernelPackages // {
+      nvidia_x11 = config.hardware.nvidia.package;
+    };
+  };
+in
+{
   options = {
     services.ollama = {
       enable = lib.mkEnableOption (
         lib.mdDoc "Server for local large language models"
       );
       listenAddress = lib.mkOption {
-        type = lib.types.str;
+        type = types.str;
         default = "127.0.0.1:11434";
         description = lib.mdDoc ''
           Specifies the bind address on which the ollama server HTTP interface listens.
         '';
       };
+      acceleration = lib.mkOption {
+        type = types.nullOr (types.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 +50,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 +58,8 @@ in {
       };
     };
 
-    environment.systemPackages = [ cfg.package ];
-
+    environment.systemPackages = [ ollamaPackage ];
   };
 
-  meta.maintainers = with lib.maintainers; [ onny ];
-
+  meta.maintainers = with lib.maintainers; [ abysssol onny ];
 }