wmantly/turing-multi-gpu-llm-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement session-managed persistent KV cache architecture with slot persistence and management API

Browse Source
This commit is contained in:
wmantly
2026-09-04 02:53:54 +00:00
parent 5daafe24e8
commit 9ab0583cab
4 changed files with 244 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ollama-proxy.py
+153 -2
View File
@@ -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():