From 8c4c67f3ad95b6c69ab83217377d67f7dc5e2c15 Mon Sep 17 00:00:00 2001 From: Safi Date: Sat, 9 May 2026 23:42:44 +0100 Subject: [PATCH] fix Ollama num_ctx: derive from actual chunk size instead of hardcoding 131072 (#798) --- graphify/llm.py | 35 +++++++++++++++++++++++------------ tests/test_llm_backends.py | 30 +++++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 13 deletions(-) diff --git a/graphify/llm.py b/graphify/llm.py index b4237c2..9c78fcb 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -281,17 +281,25 @@ def _call_openai_compat( kwargs["extra_body"] = {"thinking": {"type": "disabled"}} # Ollama defaults num_ctx to 2048 and silently truncates prompts larger # than that — the symptom is hollow 200 OK responses after the first few - # chunks (#798). We send num_ctx large enough to fit our default 60k-token - # chunk budget plus system prompt and JSON output headroom. Ollama caps - # gracefully at the model's built-in limit if it's lower, so a large - # default is safe. keep_alive pins the model in VRAM across chunks so it - # isn't unloaded/reloaded mid-run under concurrency pressure. + # chunks (#798). We derive num_ctx from the actual prompt size so we don't + # over-allocate KV-cache VRAM. Over-allocation (e.g. 128k slots for an 8k + # prompt on a 31B model) exhausts VRAM by chunk 4 and produces the same + # hollow-200 symptom — just from a different direction (#798 follow-up). + # Formula: actual input tokens + output cap + system prompt headroom. + # Capped at 131072 (enough for the default 60k token_budget); env var wins. if backend == "ollama": num_ctx_raw = os.environ.get("GRAPHIFY_OLLAMA_NUM_CTX", "").strip() - try: - num_ctx = int(num_ctx_raw) if num_ctx_raw else 131072 - except ValueError: - num_ctx = 131072 + if num_ctx_raw: + try: + num_ctx = int(num_ctx_raw) + except ValueError: + num_ctx = 131072 + else: + # Estimate input tokens: user_message chars / 4 (standard BPE + # heuristic) + 400 for the system prompt, then add output headroom. + estimated_input = len(user_message) // _CHARS_PER_TOKEN + 400 + num_ctx = min(estimated_input + max_completion_tokens + 2000, 131072) + num_ctx = max(num_ctx, 8192) # floor: never under-allocate badly keep_alive = os.environ.get("GRAPHIFY_OLLAMA_KEEP_ALIVE", "30m") kwargs["extra_body"] = {"options": {"num_ctx": num_ctx}, "keep_alive": keep_alive} resp = client.chat.completions.create(**kwargs) @@ -321,9 +329,12 @@ def _call_openai_compat( output_tokens = result["output_tokens"] if output_tokens < 50 and backend == "ollama": print( - "[graphify] warning: ollama returned very few tokens — the model may be " - "too small or not following the JSON instruction format. " - "Try a larger model with --model (e.g. --model qwen2.5-coder:14b).", + "[graphify] warning: ollama returned very few tokens — likely causes: " + "(1) VRAM pressure: check `nvidia-smi` and reduce chunk size with " + "--token-budget (e.g. --token-budget 4096) or set " + "GRAPHIFY_OLLAMA_NUM_CTX to a smaller value; " + "(2) model too small for JSON instruction following — " + "try a larger model with --model (e.g. --model qwen2.5-coder:14b).", file=sys.stderr, ) return result diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py index ec7f7a8..d2ee058 100644 --- a/tests/test_llm_backends.py +++ b/tests/test_llm_backends.py @@ -369,10 +369,38 @@ def test_ollama_extra_body_sets_num_ctx_and_keep_alive(monkeypatch): assert "extra_body" in captured, "extra_body must be sent to Ollama" eb = captured["extra_body"] - assert eb.get("options", {}).get("num_ctx") == 131072, "default num_ctx must be 131072" + # num_ctx is now dynamic: derived from message size, not hardcoded 131072 + assert "num_ctx" in eb.get("options", {}), "num_ctx must be present" + assert eb["options"]["num_ctx"] >= 8192, "num_ctx must be at least the floor value" assert eb.get("keep_alive") == "30m", "default keep_alive must be 30m" +def test_ollama_num_ctx_scales_with_small_token_budget(monkeypatch): + # Regression for #798 follow-up: with --token-budget 8192, the old hardcoded + # 131072 forced Ollama to allocate 128k KV-cache slots on a 31B model, causing + # VRAM exhaustion by chunk 4. num_ctx must now reflect actual chunk size. + captured = _install_capturing_openai(monkeypatch) + monkeypatch.delenv("GRAPHIFY_OLLAMA_NUM_CTX", raising=False) + monkeypatch.delenv("GRAPHIFY_OLLAMA_KEEP_ALIVE", raising=False) + + # Simulate an 8k-token chunk: ~32k chars of content + small_chunk_msg = "x" * 32_000 + + llm._call_openai_compat( + "http://localhost:11434/v1", "ollama", "qwen2.5-coder:7b", + small_chunk_msg, temperature=0, max_completion_tokens=16384, backend="ollama", + ) + + num_ctx = captured["extra_body"]["options"]["num_ctx"] + # Should be far less than 131072 for an 8k input — VRAM-friendly + assert num_ctx < 131072, ( + f"num_ctx={num_ctx} is too large for a small chunk; " + "this wastes VRAM and causes OOM on large models (#798)" + ) + # But still large enough to fit input + output + assert num_ctx >= 8192, "num_ctx must cover at least the output cap" + + def test_ollama_num_ctx_env_override(monkeypatch): captured = _install_capturing_openai(monkeypatch) monkeypatch.setenv("GRAPHIFY_OLLAMA_NUM_CTX", "65536")