From b8d139d8b1718c44437374a96044d00b1cf318bd Mon Sep 17 00:00:00 2001 From: shivaraj-bh Date: Sun, 24 Mar 2024 01:56:24 +0530 Subject: [PATCH 1/3] feat(ollama): add service --- nix/default.nix | 1 + nix/ollama.nix | 111 ++++++++++++++++++++++++++++++++++++++++++++ nix/ollama_test.nix | 18 +++++++ test/flake.nix | 1 + 4 files changed, 131 insertions(+) create mode 100644 nix/ollama.nix create mode 100644 nix/ollama_test.nix diff --git a/nix/default.nix b/nix/default.nix index bb86f813..8e4e17f9 100644 --- a/nix/default.nix +++ b/nix/default.nix @@ -7,6 +7,7 @@ in ./apache-kafka.nix ./clickhouse ./elasticsearch.nix + ./ollama.nix ./mysql ./nginx.nix ./postgres diff --git a/nix/ollama.nix b/nix/ollama.nix new file mode 100644 index 00000000..5b4f8e3f --- /dev/null +++ b/nix/ollama.nix @@ -0,0 +1,111 @@ +{ pkgs, lib, name, config, ... }: +let + inherit (lib) types; +in +{ + options = { + enable = lib.mkEnableOption "Enable the Ollama service"; + package = lib.mkOption { + type = types.package; + default = pkgs.ollama; + description = "The Ollama package to use"; + }; + port = lib.mkOption { + type = types.int; + default = 11434; + description = "The port on which the Ollama service's REST API will listen"; + }; + host = lib.mkOption { + type = types.str; + default = "localhost"; + description = "The host on which the Ollama service's REST API will listen"; + }; + dataDir = lib.mkOption { + type = types.str; + default = "./data/ollama"; + description = '' + The directory containing the Ollama models. + Sets the `OLLAMA_MODELS` environment variable. + ''; + }; + keepAlive = lib.mkOption { + type = types.str; + default = "5m"; + description = '' + The duration that models stay loaded in memory. + Sets the `OLLAMA_KEEP_ALIVE` environment variable. + + Note: Use a duration string like "5m" for 5 minutes. Or "70" for 70 seconds. + ''; + example = "70"; + }; + models = lib.mkOption { + type = types.listOf types.str; + default = [ ]; + description = '' + The models to load post start. + Search for models of your choice from: https://ollama.com/library + ''; + }; + + outputs.settings = lib.mkOption { + type = types.deferredModule; + internal = true; + readOnly = true; + default = { + processes = { + "${name}" = + let + startScript = pkgs.writeShellApplication { + name = "ollama-server"; + text = '' + if [ ! -d ${config.dataDir} ]; then + echo "Creating directory ${config.dataDir}" + mkdir -p ${config.dataDir} + fi + OLLAMA_HOST=${config.host}:${toString config.port} + OLLAMA_MODELS=${config.dataDir} + OLLAMA_KEEP_ALIVE=${config.keepAlive} + + export OLLAMA_HOST OLLAMA_MODELS OLLAMA_KEEP_ALIVE + + ${lib.getExe config.package} serve + ''; + }; + in + { + command = startScript; + readiness_probe = { + http_get = { + host = config.host; + port = config.port; + }; + initial_delay_seconds = 2; + period_seconds = 10; + timeout_seconds = 4; + success_threshold = 1; + failure_threshold = 5; + }; + namespace = "ollama"; + availability.restart = "on_failure"; + }; + "${name}-models" = { + command = pkgs.writeShellApplication { + name = "${name}-models"; + text = '' + set -x + models="${lib.concatStringsSep " " config.models}" + for model in $models + do + ${lib.getExe config.package} pull "$model" + done + ''; + }; + namespace = name; + depends_on."${name}".condition = "process_healthy"; + }; + }; + }; + }; + }; +} diff --git a/nix/ollama_test.nix b/nix/ollama_test.nix new file mode 100644 index 00000000..02682a32 --- /dev/null +++ b/nix/ollama_test.nix @@ -0,0 +1,18 @@ +{ pkgs, config, ... }: { + services.ollama."lama1".enable = true; + + settings.processes.test = + let + cfg = config.services.ollama."lama1"; + in + { + command = pkgs.writeShellApplication { + runtimeInputs = [ pkgs.curl ]; + text = '' + curl ${cfg.host}:${builtins.toString cfg.port} + ''; + name = "ollama-test"; + }; + depends_on."lama1-models".condition = "process_completed_successfully"; + }; +} diff --git a/test/flake.nix b/test/flake.nix index 0c5b2d87..0d892fe5 100644 --- a/test/flake.nix +++ b/test/flake.nix @@ -36,6 +36,7 @@ "${inputs.services-flake}/nix/apache-kafka_test.nix" "${inputs.services-flake}/nix/clickhouse/clickhouse_test.nix" "${inputs.services-flake}/nix/elasticsearch_test.nix" + "${inputs.services-flake}/nix/ollama_test.nix" "${inputs.services-flake}/nix/mysql/mysql_test.nix" "${inputs.services-flake}/nix/nginx_test.nix" "${inputs.services-flake}/nix/postgres/postgres_test.nix" From 009377e7d43fec85531ac5968d78647ac648055f Mon Sep 17 00:00:00 2001 From: shivaraj-bh Date: Mon, 25 Mar 2024 00:51:42 +0530 Subject: [PATCH 2/3] fix: export OLLAMA_HOST in ollama-models --- nix/ollama.nix | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nix/ollama.nix b/nix/ollama.nix index 5b4f8e3f..29ba2acf 100644 --- a/nix/ollama.nix +++ b/nix/ollama.nix @@ -94,6 +94,8 @@ in name = "${name}-models"; text = '' set -x + OLLAMA_HOST=${config.host}:${toString config.port} + export OLLAMA_HOST models="${lib.concatStringsSep " " config.models}" for model in $models do From d1fb893b83378a9cdcdecf26bdf460e49d8c2fc0 Mon Sep 17 00:00:00 2001 From: shivaraj-bh Date: Thu, 4 Apr 2024 00:22:05 +0530 Subject: [PATCH 3/3] chore: add default and extra Envs --- nix/ollama.nix | 31 ++++++++++++++++++++++++++----- nix/ollama_test.nix | 5 ++++- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/nix/ollama.nix b/nix/ollama.nix index 29ba2acf..9285b10e 100644 --- a/nix/ollama.nix +++ b/nix/ollama.nix @@ -48,6 +48,31 @@ in ''; }; + defaultEnvs = lib.mkOption { + type = types.attrsOf types.str; + internal = true; + readOnly = true; + default = { + OLLAMA_MODELS = config.dataDir; + OLLAMA_HOST = "${config.host}:${toString config.port}"; + OLLAMA_KEEP_ALIVE = config.keepAlive; + }; + description = '' + Default environment variables passed to the `ollama-server` process. + ''; + }; + + extraEnvs = lib.mkOption { + type = types.attrsOf types.str; + default = { }; + example = { + OLLAMA_DEBUG = "1"; + }; + description = '' + Extra environment variables passed to the `ollama-server` process. This is prioritized over `defaultEnvs`. + ''; + }; + outputs.settings = lib.mkOption { type = types.deferredModule; internal = true; @@ -63,11 +88,6 @@ in echo "Creating directory ${config.dataDir}" mkdir -p ${config.dataDir} fi - OLLAMA_HOST=${config.host}:${toString config.port} - OLLAMA_MODELS=${config.dataDir} - OLLAMA_KEEP_ALIVE=${config.keepAlive} - - export OLLAMA_HOST OLLAMA_MODELS OLLAMA_KEEP_ALIVE ${lib.getExe config.package} serve ''; @@ -75,6 +95,7 @@ in in { command = startScript; + environment = lib.recursiveUpdate config.defaultEnvs config.extraEnvs; readiness_probe = { http_get = { host = config.host; diff --git a/nix/ollama_test.nix b/nix/ollama_test.nix index 02682a32..657bd8bc 100644 --- a/nix/ollama_test.nix +++ b/nix/ollama_test.nix @@ -1,5 +1,8 @@ { pkgs, config, ... }: { - services.ollama."lama1".enable = true; + services.ollama."lama1" = { + enable = true; + port = 14544; + }; settings.processes.test = let