fa5436e0a7
GitOrigin-RevId: e8057b67ebf307f01bdcc8fba94d94f75039d1f6
56 lines
1.2 KiB
Nix
56 lines
1.2 KiB
Nix
import ./make-test-python.nix ({ pkgs, 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 ];
|
|
};
|
|
|
|
nodes = {
|
|
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;
|
|
};
|
|
};
|
|
|
|
testScript = ''
|
|
vms = [ cpu, rocm, cuda, altAddress ];
|
|
|
|
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)
|
|
'';
|
|
})
|