ollama: add module

modules/modules.nix

@@ -334,6 +334,7 @@ let
     ./services/nextcloud-client.nix
     ./services/nix-gc.nix
     ./services/notify-osd.nix
+    ./services/ollama.nix
     ./services/opensnitch-ui.nix
     ./services/osmscout-server.nix
     ./services/owncloud-client.nix

modules/services/ollama.nix

@@ -0,0 +1,116 @@
+{ config, lib, pkgs, ... }:
+
+with lib;
+
+let
+
+  cfg = config.services.ollama;
+
+  ollamaPackage = if cfg.acceleration == null then
+    cfg.package
+  else
+    cfg.package.override { inherit (cfg) acceleration; };
+
+  postStartScript = pkgs.writeShellScript "ollama-post-start" ''
+    set -x
+    export OLLAMA_HOST=${escapeShellArg cfg.host}:${builtins.toString cfg.port}
+    for model in ${escapeShellArgs cfg.loadModels}
+    do
+      ${escapeShellArg (getExe ollamaPackage)} pull "$model"
+    done
+  '';
+
+in {
+  meta.maintainers = [ maintainers.terlar ];
+
+  options = {
+    services.ollama = {
+      enable = mkEnableOption "ollama server for local large language models";
+
+      package = mkPackageOption pkgs "ollama" { };
+
+      host = mkOption {
+        type = types.str;
+        default = "127.0.0.1";
+        example = "[::]";
+        description = ''
+          The host address which the ollama server HTTP interface listens to.
+        '';
+      };
+
+      port = mkOption {
+        type = types.port;
+        default = 11434;
+        example = 11111;
+        description = ''
+          Which port the ollama server listens to.
+        '';
+      };
+
+      acceleration = mkOption {
+        type = types.nullOr (types.enum [ false "rocm" "cuda" ]);
+        default = null;
+        example = "rocm";
+        description = ''
+          What interface to use for hardware acceleration.
+
+          - `null`: default behavior
+            - if `nixpkgs.config.rocmSupport` is enabled, uses `"rocm"`
+            - if `nixpkgs.config.cudaSupport` is enabled, uses `"cuda"`
+            - otherwise defaults to `false`
+          - `false`: disable GPU, only use CPU
+          - `"rocm"`: supported by most modern AMD GPUs
+            - may require overriding gpu type with `services.ollama.rocmOverrideGfx`
+              if rocm doesn't detect your AMD gpu
+          - `"cuda"`: supported by most modern NVIDIA GPUs
+        '';
+      };
+
+      environmentVariables = mkOption {
+        type = types.attrsOf types.str;
+        default = { };
+        example = {
+          OLLAMA_LLM_LIBRARY = "cpu";
+          HIP_VISIBLE_DEVICES = "0,1";
+        };
+        description = ''
+          Set arbitrary environment variables for the ollama service.
+
+          Be aware that these are only seen by the ollama server (systemd service),
+          not normal invocations like `ollama run`.
+          Since `ollama run` is mostly a shell around the ollama server, this is usually sufficient.
+        '';
+      };
+
+      loadModels = mkOption {
+        type = types.listOf types.str;
+        default = [ ];
+        description = ''
+          The models to download as soon as the service starts.
+          Search for models of your choice from: https://ollama.com/library
+        '';
+      };
+    };
+  };
+
+  config = mkIf cfg.enable {
+    systemd.user.services.ollama = {
+      Unit = {
+        Description = "Server for local large language models";
+        After = [ "network.target" ];
+      };
+
+      Service = {
+        ExecStart = "${getExe ollamaPackage} serve";
+        ExecStartPost = mkIf (cfg.loadModels != [ ]) (toString postStartScript);
+        Environment =
+          (mapAttrsToList (n: v: "${n}=${v}") cfg.environmentVariables)
+          ++ [ "OLLAMA_HOST=${cfg.host}:${toString cfg.port}" ];
+      };
+
+      Install = { WantedBy = [ "default.target" ]; };
+    };
+
+    home.packages = [ ollamaPackage ];
+  };
+}

tests/default.nix

@@ -256,6 +256,7 @@ in import nmtSrc {
     ./modules/services/mpd-mpris
     ./modules/services/mpdris2
     ./modules/services/nix-gc
+    ./modules/services/ollama
     ./modules/services/osmscout-server
     ./modules/services/pantalaimon
     ./modules/services/parcellite

tests/modules/services/ollama/basic.nix

@@ -0,0 +1,13 @@
+{
+  config = {
+    services.ollama.enable = true;
+
+    test.stubs.ollama = { };
+
+    nmt.script = ''
+      serviceFile="home-files/.config/systemd/user/ollama.service"
+      assertFileRegex "$serviceFile" 'After=network\.target'
+      assertFileRegex "$serviceFile" 'Environment=OLLAMA_HOST=127.0.0.1:11434'
+    '';
+  };
+}

tests/modules/services/ollama/default.nix

@@ -0,0 +1,5 @@
+{
+  ollama-basic = ./basic.nix;
+  ollama-load-models = ./load-models.nix;
+  ollama-set-environment-variables = ./set-environment-variables.nix;
+}

tests/modules/services/ollama/load-models.nix

@@ -0,0 +1,17 @@
+{
+  config = {
+    services.ollama = {
+      enable = true;
+      loadModels = [ "llama2" ];
+    };
+
+    test.stubs.ollama = { };
+
+    nmt.script = ''
+      serviceFile="home-files/.config/systemd/user/ollama.service"
+      assertFileRegex "$serviceFile" 'ExecStartPost=/nix/store/.*-ollama-post-start'
+      generated="$(grep -o '/nix/store/.*-ollama-post-start' "$TESTED/home-files/.config/systemd/user/ollama.service")"
+      assertFileContains "$generated" "for model in llama2"
+    '';
+  };
+}

tests/modules/services/ollama/set-environment-variables.nix

@@ -0,0 +1,22 @@
+{
+  config = {
+    services.ollama = {
+      enable = true;
+      host = "localhost";
+      port = 11111;
+      environmentVariables = {
+        OLLAMA_LLM_LIBRARY = "cpu";
+        HIP_VISIBLE_DEVICES = "0,1";
+      };
+    };
+
+    test.stubs.ollama = { };
+
+    nmt.script = ''
+      serviceFile="home-files/.config/systemd/user/ollama.service"
+      assertFileRegex "$serviceFile" 'Environment=OLLAMA_HOST=localhost:11111'
+      assertFileRegex "$serviceFile" 'Environment=OLLAMA_LLM_LIBRARY=cpu'
+      assertFileRegex "$serviceFile" 'Environment=HIP_VISIBLE_DEVICES=0,1'
+    '';
+  };
+}