agentic-kvc/scripts/cache_aware_proxy.py

"""Unified cache-aware + token-level load-balanced global scheduler.

Supports two modes:
  --combined URL [URL ...]: PD co-located instances (normal vLLM, no KV transfer)
  --prefill URL BP --decode URL: PD disaggregated instances (Mooncake KV transfer)

Routing policy (same for both modes):
  score = ongoing_tokens / avg_ongoing  -  ALPHA * cache_hit_ratio
  Normalized load prevents "rich get richer"; cache bonus gives affinity.
  Session affinity: multi-turn sessions stick to same instance.
"""

import argparse
import asyncio
import os
import urllib.parse
import uuid
from contextlib import asynccontextmanager

import httpx
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse

BLOCK_SIZE = 512
CACHE_HIT_ALPHA = 1.0  # weight for cache bonus in scoring


class InstanceState:
    def __init__(self, url: str, bootstrap_port: int | None = None):
        self.url = url
        self.bootstrap_port = bootstrap_port
        self.client = httpx.AsyncClient(
            timeout=None, base_url=url,
            limits=httpx.Limits(max_connections=None, max_keepalive_connections=None),
        )
        self.ongoing_tokens = 0
        self.engine_id: dict[int, str] = {}
        self.dp_size = 1
        self.cached_blocks: set[int] = set()

    def estimate_cache_hit(self, token_ids: list[int] | None) -> int:
        if not token_ids or len(token_ids) < BLOCK_SIZE:
            return 0
        hit = 0
        for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):
            bh = hash(tuple(token_ids[i:i + BLOCK_SIZE]))
            if bh in self.cached_blocks:
                hit += BLOCK_SIZE
            else:
                break
        return hit

    def record_prefix(self, token_ids: list[int] | None):
        if not token_ids:
            return
        for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):
            self.cached_blocks.add(hash(tuple(token_ids[i:i + BLOCK_SIZE])))
        if len(self.cached_blocks) > 200000:
            self.cached_blocks = set(list(self.cached_blocks)[-100000:])


def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
                  session_id: str | None, input_length: int,
                  affinity: dict[str, int]) -> tuple[InstanceState, int]:
    """Normalized load - cache bonus scoring."""
    if session_id and session_id in affinity:
        idx = affinity[session_id]
        if idx < len(instances):
            return instances[idx], idx

    avg_load = max(sum(i.ongoing_tokens for i in instances) / len(instances), 1.0)
    best_idx, best_score = 0, float("inf")
    for i, inst in enumerate(instances):
        cache_hit = inst.estimate_cache_hit(token_ids)
        cache_ratio = cache_hit / input_length if input_length > 0 else 0.0
        score = inst.ongoing_tokens / avg_load - CACHE_HIT_ALPHA * cache_ratio
        if score < best_score:
            best_score = score
            best_idx = i

    if session_id:
        affinity[session_id] = best_idx
    return instances[best_idx], best_idx


global_args = None
combined_instances: list[InstanceState] = []
prefill_instances: list[InstanceState] = []
decode_instances: list[InstanceState] = []
session_affinity: dict[str, int] = {}
is_pd_sep = False


async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio.Event):
    for inst in instances:
        if inst.bootstrap_port is None:
            continue
        while True:
            try:
                await inst.client.get("/health")
            except Exception:
                await asyncio.sleep(1)
                continue
            parsed = urllib.parse.urlparse(str(inst.client.base_url))
            url = f"http://{parsed.hostname}:{inst.bootstrap_port}/query"
            resp = await inst.client.get(url)
            resp.raise_for_status()
            data = resp.json()
            for dp_rank, dp_entry in data.items():
                inst.engine_id[int(dp_rank)] = dp_entry["engine_id"]
            inst.dp_size = len(data)
            print(f"Inited {inst.url} engine_ids={inst.engine_id}")
            break
    ready.set()


@asynccontextmanager
async def lifespan(app: FastAPI):
    global is_pd_sep
    app.state.ready = asyncio.Event()

    if global_args.combined:
        is_pd_sep = False
        for url in global_args.combined:
            combined_instances.append(InstanceState(url))
        app.state.ready.set()
        print(f"Combined mode: {len(combined_instances)} instances")
    else:
        is_pd_sep = True
        for url, bp in global_args.prefill:
            prefill_instances.append(InstanceState(url, bp))
        for url in global_args.decode:
            decode_instances.append(InstanceState(url))
        asyncio.create_task(init_prefill_bootstrap(prefill_instances, app.state.ready))
        print(f"PD-Sep mode: {len(prefill_instances)}P + {len(decode_instances)}D")

    yield
    for inst in combined_instances + prefill_instances + decode_instances:
        await inst.client.aclose()


app = FastAPI(lifespan=lifespan)


@app.post("/v1/completions")
async def handle_completions(request: Request):
    return await _handle(request, "/v1/completions")


@app.post("/v1/chat/completions")
async def handle_chat(request: Request):
    return await _handle(request, "/v1/chat/completions")


async def _handle(request: Request, api: str):
    if not app.state.ready.is_set():
        raise HTTPException(status_code=503, detail="Service Unavailable")

    req_data = await request.json()
    request_id = str(uuid.uuid4())
    prompt = req_data.get("prompt")
    token_ids = prompt if isinstance(prompt, list) else None
    input_length = len(token_ids) if token_ids else 0
    session_id = request.headers.get("X-Session-Id")

    headers = {"X-Request-Id": request_id}
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    if is_pd_sep:
        return await _handle_pd_sep(api, req_data, request_id, token_ids,
                                     input_length, session_id, headers)
    else:
        return await _handle_combined(api, req_data, token_ids,
                                       input_length, session_id, headers)


async def _handle_combined(api, req_data, token_ids, input_length, session_id, headers):
    """Combined mode: route to best instance, send normal request."""
    inst, idx = pick_instance(combined_instances, token_ids, session_id,
                               input_length, session_affinity)
    inst.ongoing_tokens += input_length

    async def generate():
        try:
            async with inst.client.stream("POST", api, json=req_data, headers=headers) as resp:
                resp.raise_for_status()
                async for chunk in resp.aiter_bytes():
                    yield chunk
            inst.record_prefix(token_ids)
        finally:
            inst.ongoing_tokens -= input_length

    return StreamingResponse(generate(), media_type="text/event-stream")


async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
                          session_id, headers):
    """PD-Sep mode: await prefill, then stream decode."""
    p_inst, _ = pick_instance(prefill_instances, token_ids, session_id,
                               input_length, session_affinity)
    d_inst = min(decode_instances, key=lambda x: x.ongoing_tokens)

    # Await prefill
    p_inst.ongoing_tokens += input_length
    try:
        prefill_data = req_data.copy()
        prefill_data["kv_transfer_params"] = {
            "do_remote_decode": True, "do_remote_prefill": False,
            "transfer_id": f"xfer-{request_id}",
        }
        prefill_data["stream"] = False
        prefill_data["max_tokens"] = 1
        prefill_data.pop("max_completion_tokens", None)
        prefill_data.pop("stream_options", None)

        p_headers = {**headers, "X-data-parallel-rank": "0"}
        resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers)
        resp.raise_for_status()
        await resp.aclose()
        p_inst.record_prefix(token_ids)
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Prefill failed: {e}")
    finally:
        p_inst.ongoing_tokens -= input_length

    # Stream decode
    d_inst.ongoing_tokens += input_length
    parsed = urllib.parse.urlparse(str(p_inst.client.base_url))
    bootstrap_addr = f"http://{parsed.hostname}:{p_inst.bootstrap_port}"

    decode_data = req_data.copy()
    decode_data["kv_transfer_params"] = {
        "do_remote_decode": False, "do_remote_prefill": True,
        "remote_bootstrap_addr": bootstrap_addr,
        "remote_engine_id": p_inst.engine_id.get(0, ""),
        "transfer_id": f"xfer-{request_id}",
    }

    async def generate():
        try:
            async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
                resp.raise_for_status()
                async for chunk in resp.aiter_bytes():
                    yield chunk
        finally:
            d_inst.ongoing_tokens -= input_length

    return StreamingResponse(generate(), media_type="application/json")


def parse_args():
    p = argparse.ArgumentParser(description="Unified cache-aware global scheduler")
    p.add_argument("--port", type=int, default=8000)
    p.add_argument("--host", type=str, default="0.0.0.0")
    p.add_argument("--combined", nargs="+", help="Combined mode: list of instance URLs")
    p.add_argument("--prefill", nargs="+", action="append", dest="prefill_raw",
                   help="PD-Sep prefill: URL [bootstrap_port]")
    p.add_argument("--decode", nargs=1, action="append", dest="decode_raw",
                   help="PD-Sep decode: URL")
    args = p.parse_args()

    args.prefill = []
    if args.prefill_raw:
        for entry in args.prefill_raw:
            url = entry[0]
            bp = int(entry[1]) if len(entry) > 1 and entry[1].lower() != "none" else None
            args.prefill.append((url, bp))
    args.decode = [e[0] for e in (args.decode_raw or [])]

    if not args.combined and not args.prefill:
        p.error("Must specify either --combined or --prefill/--decode")
    return args


if __name__ == "__main__":
    global_args = parse_args()
    uvicorn.run(app, host=global_args.host, port=global_args.port)