nixos/ollama: add `loadModels` config option

Allows users to download model files upon service startup,
instead of at the first use of the model, improving percieved startup latency.

1 files changed, 17 insertions, 1 deletions

diff --git a/nixos/modules/services/misc/ollama.nix b/nixos/modules/services/misc/ollama.nix
index 1467c3f93bc8..c460514783ef 100644
--- a/nixos/modules/services/misc/ollama.nix
+++ b/nixos/modules/services/misc/ollama.nix
@@ -1,6 +1,6 @@
 { config, lib, pkgs, ... }:
 let
-  inherit (lib) types;
+  inherit (lib) types mkBefore;
 
   cfg = config.services.ollama;
   ollamaPackage = cfg.package.override {
@@ -132,6 +132,14 @@ in
           Since `ollama run` is mostly a shell around the ollama server, this is usually sufficient.
         '';
       };
+      loadModels = lib.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
+        '';
+      };
       openFirewall = lib.mkOption {
         type = types.bool;
         default = false;
@@ -161,6 +169,14 @@ in
         DynamicUser = cfg.sandbox;
         ReadWritePaths = cfg.writablePaths;
       };
+      postStart = mkBefore ''
+        set -x
+        export OLLAMA_HOST=${lib.escapeShellArg cfg.host}:${builtins.toString cfg.port}
+        for model in ${lib.escapeShellArgs cfg.loadModels}
+        do
+          ${lib.escapeShellArg (lib.getExe ollamaPackage)} pull "$model"
+        done
+      '';
     };
 
     networking.firewall = lib.mkIf cfg.openFirewall { allowedTCPPorts = [ cfg.port ]; };