diff --git a/modules/modules.nix b/modules/modules.nix index ed2e177f2..3c9150e0e 100644 --- a/modules/modules.nix +++ b/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 diff --git a/modules/services/ollama.nix b/modules/services/ollama.nix new file mode 100644 index 000000000..71255b68d --- /dev/null +++ b/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 ]; + }; +} diff --git a/tests/default.nix b/tests/default.nix index 4117ea8ea..30557839b 100644 --- a/tests/default.nix +++ b/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 diff --git a/tests/modules/services/ollama/basic.nix b/tests/modules/services/ollama/basic.nix new file mode 100644 index 000000000..fecee950f --- /dev/null +++ b/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' + ''; + }; +} diff --git a/tests/modules/services/ollama/default.nix b/tests/modules/services/ollama/default.nix new file mode 100644 index 000000000..4887e720b --- /dev/null +++ b/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; +} diff --git a/tests/modules/services/ollama/load-models.nix b/tests/modules/services/ollama/load-models.nix new file mode 100644 index 000000000..1d19cf658 --- /dev/null +++ b/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" + ''; + }; +} diff --git a/tests/modules/services/ollama/set-environment-variables.nix b/tests/modules/services/ollama/set-environment-variables.nix new file mode 100644 index 000000000..b9ad280c9 --- /dev/null +++ b/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' + ''; + }; +}