pyrox.dev
1import ./make-test-python.nix ({ pkgs, lib, ... }:
2let
3 mainPort = "11434";
4 altPort = "11435";
5
6 curlRequest = port: request:
7 "curl http://127.0.0.1:${port}/api/generate -d '${builtins.toJSON request}'";
8
9 prompt = {
10 model = "tinydolphin";
11 prompt = "lorem ipsum";
12 options = {
13 seed = 69;
14 temperature = 0;
15 };
16 };
17in
18{
19 name = "ollama";
20 meta = with lib.maintainers; {
21 maintainers = [ abysssol ];
22 };
23
24 nodes = {
25 cpu = { ... }: {
26 services.ollama.enable = true;
27 };
28
29 rocm = { ... }: {
30 services.ollama.enable = true;
31 services.ollama.acceleration = "rocm";
32 };
33
34 cuda = { ... }: {
35 services.ollama.enable = true;
36 services.ollama.acceleration = "cuda";
37 };
38
39 altAddress = { ... }: {
40 services.ollama.enable = true;
41 services.ollama.listenAddress = "127.0.0.1:${altPort}";
42 };
43 };
44
45 testScript = ''
46 vms = [ cpu, rocm, cuda, altAddress ];
47
48 start_all()
49 for vm in vms:
50 vm.wait_for_unit("multi-user.target")
51
52 stdout = cpu.succeed("""${curlRequest mainPort prompt}""", timeout=100)
53
54 stdout = altAddress.succeed("""${curlRequest altPort prompt}""", timeout=100)
55 '';
56})