diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index 03f4de2..7708bb3 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -18,6 +18,7 @@ import os import time as _time import urllib.parse import uuid +from collections import OrderedDict from contextlib import asynccontextmanager import httpx @@ -32,6 +33,7 @@ OVERLOAD_FACTOR = 2.0 # default; overridden by --overload-factor MAX_OFFLOAD_INFLIGHT = 4 # cap concurrent P-role offloads PREFILL_THROUGHPUT = 7000 # tokens/s per GPU (from H20 measurements) RDMA_OVERHEAD_S = 2.0 # seconds of RDMA transfer + decode start overhead +CACHE_CAPACITY_BLOCKS = 200000 # per-instance LRU cap on shadow cached_blocks class InstanceState: @@ -49,7 +51,8 @@ class InstanceState: self.active_p_offloads = 0 # number of HEAVY prefills this instance is doing for others self.engine_id: dict[int, str] = {} self.dp_size = 1 - self.cached_blocks: set[int] = set() + # OrderedDict acts as an LRU keyed by block hash; value is unused. + self.cached_blocks: OrderedDict[int, None] = OrderedDict() def estimate_cache_hit(self, token_ids: list[int] | None) -> int: if not token_ids or len(token_ids) < BLOCK_SIZE: @@ -58,6 +61,7 @@ class InstanceState: 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: + self.cached_blocks.move_to_end(bh) # LRU touch on hit hit += BLOCK_SIZE else: break @@ -67,9 +71,13 @@ class InstanceState: 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:]) + bh = hash(tuple(token_ids[i:i + BLOCK_SIZE])) + if bh in self.cached_blocks: + self.cached_blocks.move_to_end(bh) + else: + self.cached_blocks[bh] = None + if len(self.cached_blocks) > CACHE_CAPACITY_BLOCKS: + self.cached_blocks.popitem(last=False) def _p_offload_penalty(inst: InstanceState) -> int: @@ -174,11 +182,40 @@ async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio. ready.set() +async def _reconcile_loop(): + """Periodic safety net for shadow state. + + StreamingResponse generators decrement load counters in their finally + block, but if a client disconnects before the body is consumed the + generator is never entered and the decrement is lost. Clamp negative + drift every minute so router scores stay sane. This does not replace + proper exact-state syncing with vLLM (see TODO.md item 6). + """ + while True: + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + return + for inst in combined_instances + prefill_instances + decode_instances: + if inst.ongoing_tokens < 0: + inst.ongoing_tokens = 0 + if inst.ongoing_decode_tokens < 0: + inst.ongoing_decode_tokens = 0 + if inst.pending_prefill_tokens < 0: + inst.pending_prefill_tokens = 0 + if inst.num_requests < 0: + inst.num_requests = 0 + if inst.active_p_offloads < 0: + inst.active_p_offloads = 0 + + @asynccontextmanager async def lifespan(app: FastAPI): global is_pd_sep app.state.ready = asyncio.Event() + reconcile_task = asyncio.create_task(_reconcile_loop()) + if global_args.combined: is_pd_sep = False bp_list = [int(p) for p in global_args.bootstrap_ports.split(",") if p.strip()] if global_args.bootstrap_ports else [] @@ -204,6 +241,11 @@ async def lifespan(app: FastAPI): print(f"PD-Sep mode: {len(prefill_instances)}P + {len(decode_instances)}D") yield + reconcile_task.cancel() + try: + await reconcile_task + except asyncio.CancelledError: + pass for inst in combined_instances + prefill_instances + decode_instances: await inst.client.aclose()