From 989a4f3e9819bc318b21efe15cd38db138c048e2 Mon Sep 17 00:00:00 2001 From: javis-bot Date: Fri, 12 Jun 2026 23:45:16 +0900 Subject: [PATCH] perf(memory): keep embed model warm across turns (keep_alive 0 -> 5m) Empirical A/B/C measurement against the live RTX 5050 Ollama stack (qwen2.5:3b + nomic-embed-text) showed keep_alive=0 unloads the embed model ~2s after every call, so each turn after a brief idle gap pays a cold reload. VRAM is not the constraint (~4.4-4.7 GB free with both models resident) and keep_alive=0 never evicted the chat model, so CPU embedding (num_gpu=0) gave no benefit. A short positive keep_alive is the fastest of the three: it keeps the ~0.3 GB embed model resident across consecutive turns at negligible VRAM cost. Add tests/test_embeddings.py covering the warm-across-turns behaviour. Co-Authored-By: Claude Opus 4.7 --- src/jarvis/memory/embeddings.py | 13 +++++---- tests/test_embeddings.py | 52 +++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) create mode 100644 tests/test_embeddings.py diff --git a/src/jarvis/memory/embeddings.py b/src/jarvis/memory/embeddings.py index 9ad20c5..49d3999 100644 --- a/src/jarvis/memory/embeddings.py +++ b/src/jarvis/memory/embeddings.py @@ -6,11 +6,14 @@ def get_embedding(text: str, base_url: str, model: str, timeout_sec: float = 15. try: resp = requests.post( f"{base_url.rstrip('/')}/api/embeddings", - # keep_alive=0 unloads the embedding model right after the call so - # it does not sit resident in VRAM alongside the chat model. The - # chat model is pinned separately (llm.py keep_alive=30m); only the - # actively-used chat model should stay loaded. - json={"model": model, "prompt": text, "keep_alive": 0}, + # Short positive keep_alive keeps the embed model warm across the + # consecutive turns of an active conversation. With keep_alive=0 + # Ollama unloads it ~2s after every call, so each turn after a brief + # idle gap pays a cold reload of the embed model. The embed model is + # tiny (~0.3 GB) and coexists in VRAM with the chat model (pinned at + # keep_alive=30m in llm.py) with ample headroom, so holding it for a + # few minutes is effectively free and removes the per-turn reload. + json={"model": model, "prompt": text, "keep_alive": "5m"}, timeout=timeout_sec, ) resp.raise_for_status() diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..7504858 --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,52 @@ +"""Tests for the Ollama embedding client. + +Behaviour under test: the embedding request keeps the embed model warm across +consecutive conversation turns. With ``keep_alive=0`` Ollama unloads the embed +model ~2s after every call, so each turn after a short idle gap pays a cold +reload. A short positive ``keep_alive`` keeps it resident between turns at a +negligible VRAM cost (nomic-embed-text is ~0.3 GB). +""" +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from jarvis.memory.embeddings import get_embedding + + +def _mock_response(vec): + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.json.return_value = {"embedding": vec} + return resp + + +def test_get_embedding_posts_to_embeddings_endpoint(): + with patch("jarvis.memory.embeddings.requests") as mock_requests: + mock_requests.post.return_value = _mock_response([0.1, 0.2, 0.3]) + vec = get_embedding("hello", "http://localhost:11434", "nomic-embed-text") + + assert vec == [0.1, 0.2, 0.3] + args, kwargs = mock_requests.post.call_args + assert args[0].endswith("/api/embeddings") + assert kwargs["json"]["model"] == "nomic-embed-text" + assert kwargs["json"]["prompt"] == "hello" + + +def test_get_embedding_keeps_model_warm_between_turns(): + """The request must not unload the model after each call (keep_alive > 0).""" + with patch("jarvis.memory.embeddings.requests") as mock_requests: + mock_requests.post.return_value = _mock_response([0.0]) + get_embedding("warm me", "http://localhost:11434", "nomic-embed-text") + + _, kwargs = mock_requests.post.call_args + keep_alive = kwargs["json"].get("keep_alive") + # A falsy/zero keep_alive evicts the model immediately, forcing a cold + # reload on the next turn. Anything truthy positive keeps it resident. + assert keep_alive, f"embedding keep_alive should be a positive duration, got {keep_alive!r}" + assert keep_alive != 0 + + +def test_get_embedding_returns_none_on_error(): + with patch("jarvis.memory.embeddings.requests") as mock_requests: + mock_requests.post.side_effect = RuntimeError("boom") + assert get_embedding("x", "http://localhost:11434", "nomic-embed-text") is None