Commit b61914ea authored by abysssol's avatar abysssol
Browse files

nixos/ollama: add `syncModels` to remove models not defined in `loadModels`

parent 73ca95cb

nixos/modules/services/misc/ollama.nix

+41 −1
Original line number Diff line number Diff line
@@ -165,10 +165,19 @@ in 
          Download these models using `ollama pull` as soon as `ollama.service` has started.



          This creates a systemd unit `ollama-model-loader.service`.

          Use `services.ollama.syncModels` to automatically remove any models not currently declared here.



          Search for models of your choice from: <https://ollama.com/library>

        '';

      };

      syncModels = lib.mkOption {

        type = types.bool;

        default = false;

        description = ''

          Synchronize all currently installed models with those declared in `services.ollama.loadModels`,

          removing any models that are installed but not currently declared there.

        '';

      };

      openFirewall = lib.mkOption {

        type = types.bool;

        default = false;
@@ -272,7 +281,7 @@ in 
        };

    };



    systemd.services.ollama-model-loader = lib.mkIf (cfg.loadModels != [ ]) {

    systemd.services.ollama-model-loader = lib.mkIf (cfg.loadModels != [ ] || cfg.syncModels) {

      description = "Download ollama models in the background";

      wantedBy = [

        "multi-user.target"
@@ -300,13 +309,44 @@ in 
          binaryInputs = lib.mapAttrs (_: lib.getExe) {

            ollama = ollamaPackage;

            parallel = pkgs.parallel;

            awk = pkgs.gawk;

            sed = pkgs.gnused;

          };

          inherit (binaryInputs)

            ollama

            parallel

            awk

            sed

            ;



          declaredModelsRegex = lib.pipe cfg.loadModels [

            (map lib.escapeRegex)

            (lib.concatStringsSep "|")

            (lib.escape [ "/" ])

            lib.escapeShellArg

          ];

        in

        ''

          ${lib.optionalString cfg.syncModels ''

            installed=$('${ollama}' list | '${awk}' 'NR > 1 {print $1}')

            ${

              # if `declaredModelsRegex` is empty, sed will err

              if (cfg.loadModels != [ ]) then

                ''

                  echo declared models regex: ${declaredModelsRegex}

                  undeclared=$(echo "$installed" | '${sed}' -E /${declaredModelsRegex}/d)

                ''

              else

                ''

                  undeclared="$installed"

                ''

            }

            if [ -n "$undeclared" ]; then

              echo removing: $undeclared

              '${ollama}' rm $undeclared

            fi

          ''}



          '${parallel}' --tag '${ollama}' pull ::: ${lib.escapeShellArgs cfg.loadModels}

        '';

    };