diff --git a/docs/SESSION_KV_CACHE.md b/docs/SESSION_KV_CACHE.md new file mode 100644 index 0000000..0937e5a --- /dev/null +++ b/docs/SESSION_KV_CACHE.md @@ -0,0 +1,88 @@ +# Session-Managed KV Cache for llama-server + ollama-proxy + +## Goal + +Make the multi-GPU llama.cpp backend (exposed via `ollama-proxy.py`) correctly manage KV cache **per session**, so that: + +1. Different users / agents / OpenWebUI chats do not pollute each other’s context. +2. Large shared system prompts (2–30k tokens) used by Maki, OpenWebUI, and other tools are reused instead of being re-prefilled every time. +3. Session state **survives reboots** (disk-backed). + +Frontends: +- OpenWebUI +- Maki (https://github.com/wmantly/maki) + +Backend stack: +- `llama-server` (port 8080, `--parallel 1`, `--slot-save-path /var/cache/llama-slots`, `--cache-ram 16384`) +- `ollama-proxy.py` (port 11434) – Ollama / OpenAI / Anthropic compatible surface with automatic session affinity + +--- + +## Architecture + +```text +OpenWebUI ─┐ + ├──► ollama-proxy.py (port 11434) ──► llama-server (port 8080) +Maki ─┘ │ + └── session_id → /var/cache/llama-slots/.bin +``` + +- Single slot (`--parallel 1`) for maximum context and dedicated tensor parallelism across all 3 GPUs. +- Proxy owns session affinity and decides when to save / restore / erase the slot. +- Disk persistence via `--slot-save-path /var/cache/llama-slots`. +- Hot prefix reuse via `--cache-ram 16384` + `cache_prompt: true`. + +--- + +## 1. Session Identity Extraction + +The proxy extracts a stable `session_id` on every request according to this priority: + +1. **HTTP Headers**: + - `X-Session-Id` + - `X-Conversation-Id` + - `Session-Id` + - `Conversation-Id` +2. **JSON Body Fields**: + - `session_id` + - `conversation_id` + - `chat_id` + - `id` (when structured as a chat identifier) +3. **Fallback**: + - `sys-` (hash of the system prompt so identical system prompts share a base KV slot) + - `"default"` + +--- + +## 2. Proxy Session Lifecycle + +When a request arrives at `ollama-proxy.py`: + +```text +with session_lock: + if session_id == current_session: + proceed + + # 1. Persist previous session slot + if current_session and current_session != "default": + POST /slots/0?action=save {"filename": f"{current_session}.bin"} + + # 2. Restore new session or start fresh + if os.path.exists(f"/var/cache/llama-slots/{session_id}.bin"): + POST /slots/0?action=restore {"filename": f"{session_id}.bin"} + else: + POST /slots/0?action=erase + + current_session = session_id +``` + +--- + +## 3. Session Management Endpoints + +Exposed on `ollama-proxy` (`port 11434`): + +* `GET /api/sessions/current`: Returns active `session_id` and slot file status. +* `GET /api/sessions`: Lists all saved session `.bin` slot files with file sizes and timestamps. +* `POST /api/sessions/clear`: Forces erase on slot 0 and resets active session to `"default"`. +* `DELETE /api/sessions/`: Deletes the persisted `.bin` cache file for a specific session. diff --git a/scripts/ollama-proxy.py b/scripts/ollama-proxy.py index 0cb1f44..bbca64b 100644 --- a/scripts/ollama-proxy.py +++ b/scripts/ollama-proxy.py @@ -15,6 +15,7 @@ and forwards to a llama.cpp `llama-server` OpenAI-compatible backend Pure stdlib — no third-party deps. Requires Python 3.8+. """ +import os import json import sys import time @@ -25,6 +26,7 @@ import http.server import threading import re import subprocess +import hashlib # ---------------------------------------------------------------------------- # Config @@ -32,6 +34,7 @@ import subprocess HOST = "0.0.0.0" PORT = 11434 BACKEND = "http://127.0.0.1:8080/v1" # llama.cpp OpenAI endpoint +LLAMA_BASE = "http://127.0.0.1:8080" # llama.cpp server root MODEL_NAME = "Qwen3.8-Uncensored" # primary name BACKEND_MODEL = MODEL_NAME # what we send llama.cpp VERSION = "0.5.4" # fake ollama version @@ -39,6 +42,90 @@ CTX_SIZE = 262144 MAX_OUTPUT = 131072 KEEP_ALIVE = 300 +# Session Management & Slot Persistence +SLOT_SAVE_PATH = "/var/cache/llama-slots" +SESSION_LOCK = threading.Lock() +CURRENT_SESSION = "default" + +def _slot_action(action, filename=None): + """Call llama-server /slots/0?action=save|restore|erase.""" + url = f"{LLAMA_BASE}/slots/0?action={action}" + payload = {} + if filename: + payload["filename"] = filename + req = urllib.request.Request( + url, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=30) as r: + return json.loads(r.read().decode()) + except Exception as e: + print(f"[ollama-proxy] Slot action {action} (file: {filename}) error: {e}", flush=True) + return None + +def extract_session_id(headers, payload=None): + """Extract a stable session identifier from headers, body, or system prompt.""" + if headers: + for h in ("X-Session-Id", "X-Conversation-Id", "Session-Id", "Conversation-Id", "x-session-id", "x-conversation-id"): + val = headers.get(h) + if val and str(val).strip(): + return re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(val).strip())[:64] + + if payload and isinstance(payload, dict): + for k in ("session_id", "conversation_id", "chat_id", "id"): + val = payload.get(k) + if val and isinstance(val, str) and val.strip(): + return re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(val).strip())[:64] + + # Fallback: hash of system message content + messages = payload.get("messages", []) + if isinstance(messages, list): + for m in messages: + if isinstance(m, dict) and m.get("role") == "system": + content = m.get("content") or "" + if content and len(str(content)) > 20: + h = hashlib.sha256(str(content).encode("utf-8", errors="replace")).hexdigest()[:16] + return f"sys-{h}" + + return "default" + +def ensure_session(session_id): + """Ensure slot 0 contains the KV cache for session_id, saving/restoring as needed.""" + global CURRENT_SESSION + if not session_id: + session_id = "default" + + clean_id = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(session_id))[:64] + + with SESSION_LOCK: + if clean_id == CURRENT_SESSION: + return + + # 1. Save old session if not default + if CURRENT_SESSION and CURRENT_SESSION != "default": + save_file = f"{CURRENT_SESSION}.bin" + print(f"[ollama-proxy] Saving slot 0 for session '{CURRENT_SESSION}' -> {save_file}", flush=True) + _slot_action("save", save_file) + + # 2. Restore new session or erase + target_file = f"{clean_id}.bin" + target_path = os.path.join(SLOT_SAVE_PATH, target_file) + + if os.path.exists(target_path): + print(f"[ollama-proxy] Restoring slot 0 for session '{clean_id}' <- {target_file}", flush=True) + res = _slot_action("restore", target_file) + if not res: + print(f"[ollama-proxy] Restore failed for '{clean_id}', falling back to erase", flush=True) + _slot_action("erase") + else: + print(f"[ollama-proxy] Starting fresh slot for session '{clean_id}'", flush=True) + _slot_action("erase") + + CURRENT_SESSION = clean_id + MODEL_TAGS = [ "Qwen3.8-Uncensored:latest", "Qwen3.8-Uncensored", @@ -87,7 +174,7 @@ def _post(path, payload, stream=False): def _backend_chat(messages, **kw): """Call llama.cpp /chat/completions, return parsed JSON (non-streaming).""" - body = {"model": BACKEND_MODEL, "messages": messages} + body = {"model": BACKEND_MODEL, "messages": messages, "cache_prompt": True} body.update(kw) body.setdefault("stream", False) with _post("/chat/completions", body) as r: @@ -99,6 +186,7 @@ def _backend_stream(messages, **kw): "model": BACKEND_MODEL, "messages": messages, "stream": True, + "cache_prompt": True, "stream_options": {"include_usage": True}, } body.update(kw) @@ -1133,6 +1221,31 @@ class Handler(http.server.BaseHTTPRequestHandler): self._send(200, ollama_tags()) elif path == "/api/ps": self._send(200, ollama_ps()) + elif path == "/api/sessions/current": + active_file = os.path.join(SLOT_SAVE_PATH, f"{CURRENT_SESSION}.bin") + self._send(200, { + "session_id": CURRENT_SESSION, + "slot_save_path": SLOT_SAVE_PATH, + "filename": f"{CURRENT_SESSION}.bin", + "exists_on_disk": os.path.exists(active_file), + "file_size_bytes": os.path.getsize(active_file) if os.path.exists(active_file) else 0, + }) + elif path == "/api/sessions": + sessions = [] + if os.path.exists(SLOT_SAVE_PATH): + for fname in os.listdir(SLOT_SAVE_PATH): + if fname.endswith(".bin"): + fpath = os.path.join(SLOT_SAVE_PATH, fname) + stat = os.stat(fpath) + sessions.append({ + "session_id": fname[:-4], + "filename": fname, + "size_bytes": stat.st_size, + "size_mb": round(stat.st_size / (1024 * 1024), 2), + "modified_iso": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(stat.st_mtime)), + "is_active": (fname[:-4] == CURRENT_SESSION), + }) + self._send(200, {"sessions": sorted(sessions, key=lambda s: s["size_bytes"], reverse=True)}) elif path == "/v1/models": self._send(200, { "object": "list", @@ -1164,6 +1277,27 @@ class Handler(http.server.BaseHTTPRequestHandler): payload = self._read_json() try: + # Session-Aware Management Routes + if path == "/api/sessions/clear": + with SESSION_LOCK: + _slot_action("erase") + global CURRENT_SESSION + CURRENT_SESSION = "default" + self._send(200, {"status": "cleared", "current_session": "default"}) + return + elif path == "/api/sessions/save": + sid = payload.get("session_id") or CURRENT_SESSION + clean_id = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', str(sid))[:64] + with SESSION_LOCK: + _slot_action("save", f"{clean_id}.bin") + self._send(200, {"status": "saved", "session_id": clean_id, "filename": f"{clean_id}.bin"}) + return + + # Extract & enforce session affinity for all generation and chat endpoints + if path in ("/api/chat", "/api/generate", "/v1/messages", "/v1/chat/completions", "/v1/completions", "/v1/responses"): + session_id = extract_session_id(self.headers, payload) + ensure_session(session_id) + if path == "/api/chat": result, is_stream = ollama_chat(payload) if is_stream: @@ -1204,6 +1338,7 @@ class Handler(http.server.BaseHTTPRequestHandler): result, _ = anthropic_count_tokens(payload) self._send(200, result) elif path in ("/v1/chat/completions", "/v1/completions", "/v1/embeddings", "/v1/responses"): + payload.setdefault("cache_prompt", True) body = json.dumps(payload).encode() backend_path = BACKEND + path[len("/v1"):] req = urllib.request.Request( @@ -1248,8 +1383,24 @@ class Handler(http.server.BaseHTTPRequestHandler): except Exception as e: self._send(500, {"error": f"{type(e).__name__}: {e}"}) + def do_DELETE(self): + path = self.path.split("?")[0] + if path.startswith("/api/sessions/"): + sid = path[len("/api/sessions/"):].strip() + clean_id = re.sub(r'[^a-zA-Z0-9_\-\.]', '_', sid)[:64] + target = os.path.join(SLOT_SAVE_PATH, f"{clean_id}.bin") + if os.path.exists(target): + try: + os.remove(target) + self._send(200, {"status": "deleted", "session_id": clean_id}) + except Exception as e: + self._send(500, {"error": f"Failed to delete session file: {e}"}) + else: + self._send(404, {"error": f"Session '{clean_id}' not found on disk"}) + else: + self.do_POST() + do_PUT = do_POST - do_DELETE = do_POST def main(): diff --git a/scripts/start-server.sh b/scripts/start-server.sh index 65b5d38..3e11218 100755 --- a/scripts/start-server.sh +++ b/scripts/start-server.sh @@ -19,6 +19,8 @@ exec /opt/llama.cpp-xrip/build-nccl/bin/llama-server \ -ngl 99 \ -c 262144 \ --parallel 1 \ + --slot-save-path /var/cache/llama-slots \ + --cache-ram 16384 \ --split-mode tensor \ --flash-attn on \ --batch-size 1024 \ diff --git a/systemd/llama-server.service b/systemd/llama-server.service index 66204eb..fac5312 100644 --- a/systemd/llama-server.service +++ b/systemd/llama-server.service @@ -19,7 +19,7 @@ TimeoutStopSec=60 # Basic hardening (server only listens on :8080, reads models read-only) ProtectSystem=strict ProtectHome=read-only -ReadWritePaths= +ReadWritePaths=/var/cache/llama-slots /tmp NoNewPrivileges=true PrivateTmp=true # Do not cap memory/CPU; inference needs all of it