From aa8e02582dddb739a02b15175d14f42e13a08335 Mon Sep 17 00:00:00 2001
From: abysssol <abysssol@pm.me>
Date: Sat, 13 Jul 2024 17:26:29 -0400
Subject: [PATCH] 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.
---
 nixos/tests/all-tests.nix          |  4 +-
 nixos/tests/ollama-cuda.nix        | 17 +++++++
 nixos/tests/ollama-rocm.nix        | 17 +++++++
 nixos/tests/ollama.nix             | 79 ++++++++++++++----------------
 pkgs/by-name/ol/ollama/package.nix |  2 +
 5 files changed, 77 insertions(+), 42 deletions(-)
 create mode 100644 nixos/tests/ollama-cuda.nix
 create mode 100644 nixos/tests/ollama-rocm.nix

diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index b9a6311eaf35..186071d4d68f 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -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 {};
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 e726239a2386..4ac08628c168 100644
--- a/pkgs/by-name/ol/ollama/package.nix
+++ b/pkgs/by-name/ol/ollama/package.nix
@@ -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 = {