Commit 73ca95cb authored by abysssol's avatar abysssol
Browse files

nixos/ollama: switch `loadModels` script to use `parallel --tag` 

This should theoretically be (slightly) faster, and will prepend the
relevant argument/model before every stdout/stderr message, allowing easy
identification of which model is causing an error (eg due to a typo).
parent 81fcca10

nixos/modules/services/misc/ollama.nix

+22 −24
Original line number Diff line number Diff line
@@ -152,7 +152,15 @@ in 
      };

      loadModels = lib.mkOption {

        type = types.listOf types.str;

        apply = builtins.filter (model: model != "");

        default = [ ];

        example = [

          "dolphin3"

          "gemma3"

          "gemma3:27b"

          "deepseek-r1:latest"

          "deepseek-r1:1.5b"

        ];

        description = ''

          Download these models using `ollama pull` as soon as `ollama.service` has started.
@@ -287,29 +295,19 @@ in 
        RestartSteps = "10";

      };



      script = ''

        total=${toString (builtins.length cfg.loadModels)}

        failed=0



        for model in ${lib.escapeShellArgs cfg.loadModels}; do

          '${lib.getExe ollamaPackage}' pull "$model" &

        done



        for job in $(jobs -p); do

          set +e

          wait $job

          exit_code=$?

          set -e



          if [ $exit_code != 0 ]; then

            failed=$((failed + 1))

          fi

        done



        if [ $failed != 0 ]; then

          echo "error: $failed out of $total attempted model downloads failed" >&2

          exit 1

        fi

      script =

        let

          binaryInputs = lib.mapAttrs (_: lib.getExe) {

            ollama = ollamaPackage;

            parallel = pkgs.parallel;

          };

          inherit (binaryInputs)

            ollama

            parallel

            ;

        in

        ''

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

        '';

    };