Skip to content

Commit d84efa4

Browse files
shivaraj-bhshivaraj-bh
committed
feat(ollama): init
ported from https://github.com/shivaraj-bh/ollama-flake/blob/main/services/ollama.nix Also see the nixos module: https://github.com/NixOS/nixpkgs/blob/master/nixos/modules/services/misc/ollama.nix
1 parent 3e68d66 commit d84efa4

File tree

5 files changed

+208
-0
lines changed

5 files changed

+208
-0
lines changed

‎doc/ollama.md

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
# Ollama
2+
3+
[Ollama](https://github.com/ollama/ollama) enables you to get up and running with Llama 3, Mistral, Gemma, and other large language models.
4+
5+
## Getting Started
6+
7+
```nix
8+
# In `perSystem.process-compose.<name>`
9+
{
10+
services.ollama."ollama1".enable = true;
11+
}
12+
```
13+
14+
## Acceleration
15+
16+
By default Ollama uses the CPU for inference. To enable GPU acceleration:
17+
18+
### Cuda
19+
20+
```nix
21+
# In `perSystem.process-compose.<name>`
22+
{
23+
services.ollama."ollama1" = {
24+
enable = true;
25+
acceleration = "cuda";
26+
};
27+
}
28+
```
29+
30+
### ROCm
31+
32+
```nix
33+
# In `perSystem.process-compose.<name>`
34+
{
35+
services.ollama."ollama1" = {
36+
enable = true;
37+
acceleration = "rocm";
38+
};
39+
}
40+
```

‎nix/default.nix

+1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ in
99
./elasticsearch.nix
1010
./mysql
1111
./nginx
12+
./ollama.nix
1213
./postgres
1314
./redis-cluster.nix
1415
./redis.nix

‎nix/ollama.nix

+150
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
# Based on https://github.com/shivaraj-bh/ollama-flake/blob/main/services/ollama.nix
2+
{ pkgs, lib, name, config, ... }:
3+
let
4+
inherit (lib) types;
5+
ollamaPackage = pkgs.ollama.override {
6+
inherit (config) acceleration;
7+
linuxPackages = config.boot.kernelPackages // {
8+
nvidia_x11 = config.hardware.nvidia.package;
9+
};
10+
};
11+
in
12+
{
13+
options = {
14+
enable = lib.mkEnableOption "Enable the Ollama service";
15+
package = lib.mkOption {
16+
type = types.package;
17+
default = ollamaPackage;
18+
description = "The Ollama package to use";
19+
};
20+
port = lib.mkOption {
21+
type = types.port;
22+
default = 11434;
23+
description = "The port on which the Ollama service's REST API will listen";
24+
};
25+
host = lib.mkOption {
26+
type = types.str;
27+
default = "127.0.0.1";
28+
example = "0.0.0.0";
29+
description = "The host on which the Ollama service's REST API will listen";
30+
};
31+
dataDir = lib.mkOption {
32+
type = types.str;
33+
default = "./data/ollama";
34+
description = ''
35+
The directory containing the Ollama models.
36+
Sets the `OLLAMA_MODELS` environment variable.
37+
'';
38+
};
39+
keepAlive = lib.mkOption {
40+
type = types.str;
41+
default = "5m";
42+
description = ''
43+
The duration that models stay loaded in memory.
44+
Sets the `OLLAMA_KEEP_ALIVE` environment variable.
45+
46+
Note: Use a duration string like "5m" for 5 minutes. Or "70" for 70 seconds.
47+
'';
48+
example = "70";
49+
};
50+
models = lib.mkOption {
51+
type = types.listOf types.str;
52+
default = [ ];
53+
description = ''
54+
The models to load post start.
55+
Search for models of your choice from: https://ollama.com/library
56+
'';
57+
};
58+
acceleration = lib.mkOption {
59+
type = types.nullOr (types.enum [ false "rocm" "cuda" ]);
60+
default = null;
61+
example = "rocm";
62+
description = ''
63+
What interface to use for hardware acceleration.
64+
65+
- `null`: default behavior
66+
- if `nixpkgs.config.rocmSupport` is enabled, uses `"rocm"`
67+
- if `nixpkgs.config.cudaSupport` is enabled, uses `"cuda"`
68+
- otherwise defaults to `false`
69+
- `false`: disable GPU, only use CPU
70+
- `"rocm"`: supported by most modern AMD GPUs
71+
- `"cuda"`: supported by most modern NVIDIA GPUs
72+
'';
73+
};
74+
environment = lib.mkOption {
75+
type = types.attrsOf types.str;
76+
default = { };
77+
example = {
78+
OLLAMA_DEBUG = "1";
79+
};
80+
description = ''
81+
Extra environment variables passed to the `ollama-server` process.
82+
'';
83+
};
84+
85+
outputs.settings = lib.mkOption {
86+
type = types.deferredModule;
87+
internal = true;
88+
readOnly = true;
89+
default = {
90+
processes = {
91+
"${name}" =
92+
let
93+
startScript = pkgs.writeShellApplication {
94+
name = "ollama-server";
95+
text = ''
96+
if [ ! -d ${config.dataDir} ]; then
97+
echo "Creating directory ${config.dataDir}"
98+
mkdir -p ${config.dataDir}
99+
fi
100+
101+
${lib.getExe config.package} serve
102+
'';
103+
};
104+
in
105+
{
106+
command = startScript;
107+
108+
environment = {
109+
OLLAMA_MODELS = config.dataDir;
110+
OLLAMA_HOST = "${config.host}:${toString config.port}";
111+
OLLAMA_KEEP_ALIVE = config.keepAlive;
112+
} // config.environment;
113+
114+
readiness_probe = {
115+
http_get = {
116+
host = config.host;
117+
port = config.port;
118+
};
119+
initial_delay_seconds = 2;
120+
period_seconds = 10;
121+
timeout_seconds = 4;
122+
success_threshold = 1;
123+
failure_threshold = 5;
124+
};
125+
namespace = name;
126+
availability.restart = "on_failure";
127+
};
128+
129+
"${name}-models" = {
130+
command = pkgs.writeShellApplication {
131+
name = "ollama-models";
132+
text = ''
133+
set -x
134+
OLLAMA_HOST=${config.host}:${toString config.port}
135+
export OLLAMA_HOST
136+
models="${lib.concatStringsSep " " config.models}"
137+
for model in $models
138+
do
139+
${lib.getExe config.package} pull "$model"
140+
done
141+
'';
142+
};
143+
namespace = name;
144+
depends_on."${name}".condition = "process_healthy";
145+
};
146+
};
147+
};
148+
};
149+
};
150+
}

‎nix/ollama_test.nix

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
{ pkgs, ... }: {
2+
services.ollama."ollama1".enable = true;
3+
4+
# Cannot test auto-loading models yet because that requires internet connection.
5+
settings.processes.test =
6+
{
7+
command = pkgs.writeShellApplication {
8+
runtimeInputs = [ pkgs.curl ];
9+
text = ''
10+
curl http://127.0.0.1:11434
11+
'';
12+
name = "ollama-test";
13+
};
14+
depends_on."ollama1".condition = "process_healthy";
15+
};
16+
}

‎test/flake.nix

+1
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
"${inputs.services-flake}/nix/elasticsearch_test.nix"
3939
"${inputs.services-flake}/nix/mysql/mysql_test.nix"
4040
"${inputs.services-flake}/nix/nginx/nginx_test.nix"
41+
"${inputs.services-flake}/nix/ollama_test.nix"
4142
"${inputs.services-flake}/nix/postgres/postgres_test.nix"
4243
"${inputs.services-flake}/nix/redis_test.nix"
4344
"${inputs.services-flake}/nix/redis-cluster_test.nix"

0 commit comments

Comments
 (0)