# LiteLLM session-affinity callback — published alongside # "Two engines, one Halo: a benchmark confession" (digitalcyberia.com). # # Problem: two identical inference boxes behind one LiteLLM router. The # default least-busy routing shatters each engine's prompt cache — turn one # warms box A, turn two lands on box B, and a 50k-token conversation gets # re-read from scratch (roughly eighteen seconds), which on engines that # batch prefill first ALSO freezes every other conversation on box B. # # Fix: pin each conversation to one box; send only genuinely new # conversations to the least-loaded box. Pinning works by rewriting # data["model"] to the deployment id (model_info.id in config.yaml) — the # router honours that. (metadata["model_id"] is NOT read by the router; a # first attempt using it silently did nothing.) # # Config side (config.yaml): each deployment gets a stable model_info.id, # and this file is wired in with # litellm_settings: # callbacks: custom_callbacks.proxy_handler_instance # Replace the BASES hosts with your own. import asyncio import hashlib import re import sys import time import httpx from litellm.integrations.custom_logger import CustomLogger BASES = { "halo-a": "http://box-a.internal:4000", "halo-b": "http://box-b.internal:4000", } DEP_IDS = tuple(BASES) SESSION_TTL = 3600.0 # forget a pin after this long without a request LOAD_TTL = 1.0 # seconds to cache a /metrics read _sessions = {} # sid -> (dep_id, last_seen) _load_cache = {} # dep_id -> (ts, load or None if unhealthy) _inflight = {} # litellm_call_id -> dep_id (routed by us, not finished) _recent = [] # (dep_id, ts) of our picks; the server gauge lags behind them RECENT_TTL = 3.0 _pick_lock = asyncio.Lock() # concurrent new sessions must see each other's picks # in-flight gauges: llama.cpp-style (halogen, gufo) or vLLM _METRIC = re.compile(r'^(?:llamacpp:requests_(?:processing|deferred)|vllm:num_requests_(?:running|waiting))(?:\{[^}]*\})?\s+([0-9.eE+-]+)', re.M) # Session keys recognised: metadata.session_id, or x-litellm-session-id / # x-claude-code-session-id / x-session-id / session_id + x-session-affinity / # x-opencode-session headers. def _log(msg): print(f"[affinity] {msg}", file=sys.stderr, flush=True) async def _load(dep_id: str): """in-flight + queued requests on that box, or None if it's unreachable.""" now = time.time() cached = _load_cache.get(dep_id) if cached and now - cached[0] < LOAD_TTL: return cached[1] try: async with httpx.AsyncClient(timeout=1.0) as c: r = await c.get(BASES[dep_id] + "/metrics") load = sum(float(v) for v in _METRIC.findall(r.text)) if r.status_code == 200 else None except Exception: load = None _load_cache[dep_id] = (now, load) return load async def _least_busy(exclude=None): best, best_load = None, None for dep in DEP_IDS: if dep == exclude: continue load = await _load(dep) if load is None: continue # requests we just routed may not show up in the server's gauge yet load += sum(1 for d in _inflight.values() if d == dep) * 0.5 load += sum(1 for d, ts in _recent if d == dep and time.time() - ts < RECENT_TTL) * 0.5 # agents are often between turns (running tools), so both replicas can # read 0; break ties toward the one with fewer sessions pinned to it load += sum(1 for d, _ in _sessions.values() if d == dep) * 0.01 if best_load is None or load < best_load: best, best_load = dep, load return best def _find_session_id(data: dict): meta = data.get("metadata") or data.get("litellm_metadata") or {} sid = meta.get("session_id") if sid: return str(sid) # headers may arrive via metadata.headers, extra_headers, or top-level # headers; on /v1/messages only proxy_server_request.headers has them # (metadata there belongs to Anthropic's own plumbing) psr = (data.get("proxy_server_request") or {}).get("headers") lmeta = (data.get("litellm_metadata") or {}).get("headers") for src in (meta.get("headers"), lmeta, psr, data.get("headers"), data.get("extra_headers")): if isinstance(src, dict): for k, v in src.items(): lk = str(k).lower() if lk in ("x-litellm-session-id", "x-claude-code-session-id", "x-session-id", "session-id", "session_id", "x-session-affinity", "x-opencode-session"): return str(v) return None def _forget(kwargs): call_id = (kwargs or {}).get("litellm_call_id") if call_id: _inflight.pop(call_id, None) class SessionAffinity(CustomLogger): async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): try: if call_type not in ("completion", "acompletion", "anthropic_messages") \ or data.get("model") != "my-model": # the model you pool return data now = time.time() for s, (_, seen) in list(_sessions.items()): if now - seen > SESSION_TTL: del _sessions[s] sid = _find_session_id(data) dep, why = None, "least-busy" async with _pick_lock: _recent[:] = [(d, ts) for d, ts in _recent if now - ts < RECENT_TTL] if sid: pinned = _sessions.get(sid) if pinned and await _load(pinned[0]) is not None: dep, why = pinned[0], "sticky" else: dep = await _least_busy(exclude=pinned[0] if pinned else None) why = "repin (pinned replica down)" if pinned else "new session" if dep: _recent.append((dep, now)) if dep: _sessions[sid] = (dep, now) else: dep = await _least_busy() if dep: _recent.append((dep, now)) if dep: # Anthropic-style /v1/messages calls go to a native passthrough # deployment registered on the same box (suffix "-msg" here) data["model"] = dep + "-msg" if call_type == "anthropic_messages" else dep if data.get("litellm_call_id"): _inflight[data["litellm_call_id"]] = dep sid_s = hashlib.md5(sid.encode()).hexdigest()[:8] if sid else "-" _log(f"{why} sid={sid_s} -> {data['model']} loads=" f"{ {d: (_load_cache.get(d) or (0, None))[1] for d in DEP_IDS} }") except Exception as e: _log(f"hook error, falling back to router: {e!r}") return data async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): _forget(kwargs) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): _forget(kwargs) proxy_handler_instance = SessionAffinity()