diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix index 29eb36ab1f28..c22c3e961d08 100644 --- a/nixos/tests/all-tests.nix +++ b/nixos/tests/all-tests.nix @@ -687,7 +687,9 @@ in { ocis = handleTest ./ocis.nix {}; oddjobd = handleTestOn [ "x86_64-linux" "aarch64-linux" ] ./oddjobd.nix {}; oh-my-zsh = handleTest ./oh-my-zsh.nix {}; - ollama = handleTest ./ollama.nix {}; + ollama = runTest ./ollama.nix; + ollama-cuda = runTestOn ["x86_64-linux" "aarch64-linux"] ./ollama-cuda.nix; + ollama-rocm = runTestOn ["x86_64-linux" "aarch64-linux"] ./ollama-rocm.nix; ombi = handleTest ./ombi.nix {}; openarena = handleTest ./openarena.nix {}; openldap = handleTest ./openldap.nix {}; diff --git a/nixos/tests/ollama-cuda.nix b/nixos/tests/ollama-cuda.nix new file mode 100644 index 000000000000..bbab7e24d35c --- /dev/null +++ b/nixos/tests/ollama-cuda.nix @@ -0,0 +1,17 @@ +{ lib, ... }: +{ + name = "ollama-cuda"; + meta.maintainers = with lib.maintainers; [ abysssol ]; + + nodes.cuda = + { ... }: + { + services.ollama.enable = true; + services.ollama.acceleration = "cuda"; + }; + + testScript = '' + cuda.wait_for_unit("multi-user.target") + cuda.wait_for_open_port(11434) + ''; +} diff --git a/nixos/tests/ollama-rocm.nix b/nixos/tests/ollama-rocm.nix new file mode 100644 index 000000000000..81915630d950 --- /dev/null +++ b/nixos/tests/ollama-rocm.nix @@ -0,0 +1,17 @@ +{ lib, ... }: +{ + name = "ollama-rocm"; + meta.maintainers = with lib.maintainers; [ abysssol ]; + + nodes.rocm = + { ... }: + { + services.ollama.enable = true; + services.ollama.acceleration = "rocm"; + }; + + testScript = '' + rocm.wait_for_unit("multi-user.target") + rocm.wait_for_open_port(11434) + ''; +} diff --git a/nixos/tests/ollama.nix b/nixos/tests/ollama.nix index 30e475553eb1..34347716af72 100644 --- a/nixos/tests/ollama.nix +++ b/nixos/tests/ollama.nix @@ -1,56 +1,53 @@ -import ./make-test-python.nix ({ pkgs, lib, ... }: +{ lib, ... }: let mainPort = 11434; altPort = 11435; - - curlRequest = port: request: - "curl http://127.0.0.1:${toString port}/api/generate -d '${builtins.toJSON request}'"; - - prompt = { - model = "tinydolphin"; - prompt = "lorem ipsum"; - options = { - seed = 69; - temperature = 0; - }; - }; in { name = "ollama"; - meta = with lib.maintainers; { - maintainers = [ abysssol ]; - }; + meta.maintainers = with lib.maintainers; [ abysssol ]; nodes = { - cpu = { ... }: { - services.ollama.enable = true; - }; + cpu = + { ... }: + { + services.ollama.enable = true; + }; - rocm = { ... }: { - services.ollama.enable = true; - services.ollama.acceleration = "rocm"; - }; - - cuda = { ... }: { - services.ollama.enable = true; - services.ollama.acceleration = "cuda"; - }; - - altAddress = { ... }: { - services.ollama.enable = true; - services.ollama.port = altPort; - }; + altAddress = + { ... }: + { + services.ollama.enable = true; + services.ollama.port = altPort; + }; }; testScript = '' - vms = [ cpu, rocm, cuda, altAddress ]; + import json + + def curl_request_ollama(prompt, port): + json_prompt = json.dumps(prompt) + return f"""curl http://127.0.0.1:{port}/api/generate -d '{json_prompt}'""" + + prompt = { + "model": "tinydolphin", + "prompt": "lorem ipsum", + "options": { + "seed": 69, + "temperature": 0, + }, + } + + + vms = [ + (cpu, ${toString mainPort}), + (altAddress, ${toString altPort}), + ] start_all() - for vm in vms: - vm.wait_for_unit("multi-user.target") - - stdout = cpu.succeed("""${curlRequest mainPort prompt}""", timeout=100) - - stdout = altAddress.succeed("""${curlRequest altPort prompt}""", timeout=100) + for (vm, port) in vms: + vm.wait_for_unit("multi-user.target") + vm.wait_for_open_port(port) + stdout = vm.succeed(curl_request_ollama(prompt, port), timeout = 100) ''; -}) +} diff --git a/pkgs/by-name/ol/ollama/package.nix b/pkgs/by-name/ol/ollama/package.nix index 31edc28dc8fa..a59fe29a15dc 100644 --- a/pkgs/by-name/ol/ollama/package.nix +++ b/pkgs/by-name/ol/ollama/package.nix @@ -214,6 +214,8 @@ goBuild ((lib.optionalAttrs enableRocm { }; } // lib.optionalAttrs stdenv.isLinux { inherit ollama-rocm ollama-cuda; + service-cuda = nixosTests.ollama-cuda; + service-rocm = nixosTests.ollama-rocm; }; meta = {