Commit aa8e0258 authored by abysssol's avatar abysssol
Browse files

nixos/ollama: split cuda and rocm from service test 

Cuda and rocm are platform specific, only working on linux.
Also, ofborg won't run any of the service test at all because
cuda is unfreely licensed.
parent fdd0aa0a

nixos/tests/all-tests.nix

+3 −1
Original line number Diff line number Diff line
@@ -685,7 +685,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 {};

nixos/tests/ollama-cuda.nix

0 → 100644
+17 −0
Original line number Diff line number Diff line 
{ 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)

  '';

}

nixos/tests/ollama-rocm.nix

0 → 100644
+17 −0
Original line number Diff line number Diff line 
{ 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)

  '';

}

nixos/tests/ollama.nix

+38 −41
Original line number Diff line number Diff line 
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;

    };



    rocm = { ... }: {

      services.ollama.enable = true;

      services.ollama.acceleration = "rocm";

    };



    cuda = { ... }: {

    cpu =

      { ... }:

      {

        services.ollama.enable = true;

      services.ollama.acceleration = "cuda";

      };



    altAddress = { ... }: {

    altAddress =

      { ... }:

      {

        services.ollama.enable = true;

        services.ollama.port = altPort;

      };

  };



  testScript = ''

    vms = [ cpu, rocm, cuda, altAddress ];

    import json



    start_all()

    for vm in vms:

        vm.wait_for_unit("multi-user.target")

    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}'"""



    stdout = cpu.succeed("""${curlRequest mainPort prompt}""", timeout=100)

    prompt = {

      "model": "tinydolphin",

      "prompt": "lorem ipsum",

      "options": {

        "seed": 69,

        "temperature": 0,

      },

    }





    vms = [

      (cpu, ${toString mainPort}),

      (altAddress, ${toString altPort}),

    ]



    stdout = altAddress.succeed("""${curlRequest altPort prompt}""", timeout=100)

    start_all()

    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)

  '';

})
 
}

pkgs/by-name/ol/ollama/package.nix

+2 −0
Original line number Diff line number Diff line
@@ -212,6 +212,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 = {