MB5 driver updates: PD-proxy + snapshot instrument + launcher tweaks

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

microbench/fresh_setup/mb5_pd_proxy.py

@@ -72,8 +72,32 @@ async def lifespan(app: FastAPI):
     # Startup: Initialize client pools for prefiller and decoder services
     app.state.prefill_clients = []
     app.state.decode_clients = []
+    app.state.colo_clients = []
     app.state.ready = asyncio.Event()
 
+    # Colo (PD-combined) passthrough mode: no bootstrap handshake needed.
+    if global_args.colo:
+        for url in global_args.colo:
+            app.state.colo_clients.append({
+                "client": httpx.AsyncClient(
+                    timeout=None,
+                    base_url=url,
+                    trust_env=False,  # ignore http_proxy env: backends are localhost
+                    limits=httpx.Limits(
+                        max_connections=None,
+                        max_keepalive_connections=None,
+                    ),
+                ),
+                "url": url,
+            })
+        app.state.colo_iterator = itertools.cycle(range(len(app.state.colo_clients)))
+        app.state.ready.set()
+        print(f"Colo passthrough mode: {len(app.state.colo_clients)} kv_both clients.")
+        yield
+        for client_info in app.state.colo_clients:
+            await client_info["client"].aclose()
+        return
+
     # Create prefill clients
     for i, (url, bootstrap_port) in enumerate(global_args.prefill):
         parsed_url = urllib.parse.urlparse(url)
@@ -169,9 +193,25 @@ def parse_args():
         help="Decode server URL. Can be specified multiple times.",
     )
 
+    # MB5: colocated (PD-combined) instances. When given, the proxy runs in
+    # "colo" mode — it round-robins /v1/completions to these kv_both instances
+    # with a plain streaming passthrough (no P->D split, no kv_transfer_params).
+    # This exists so the 8C baseline pays the SAME proxy hop as PD configs,
+    # removing the "8C bypasses the proxy" confound from the comparison.
+    parser.add_argument(
+        "--colo",
+        nargs=1,
+        action="append",
+        dest="colo_raw",
+        metavar=("URL",),
+        help="Colocated (kv_both) server URL. Can be specified multiple times. "
+             "Enables colo passthrough mode.",
+    )
+
     args = parser.parse_args()
     args.prefill = _parse_prefill_urls(args.prefill_raw)
     args.decode = _parse_decode_urls(args.decode_raw)
+    args.colo = [u[0] for u in args.colo_raw] if args.colo_raw else []
 
     return args
 
@@ -235,6 +275,14 @@ def _parse_decode_urls(decode_list):
 # Decode side stays round-robin (load balance) regardless.
 MB5_P_ROUTING = os.environ.get("MB5_P_ROUTING", "rr").lower()
 
+# MB5: routing mode for the COLO (kv_both) passthrough proxy.
+#   "rr"      — round-robin (loses session-local prefix cache)
+#   "session" — consistent hash on X-Session-Id, so all turns of a session land
+#               on the same kv_both instance and reuse its prefix cache. This is
+#               the cache-aware colo baseline (the fair strong baseline for the
+#               agentic reuse regime — D4).
+MB5_COLO_ROUTING = os.environ.get("MB5_COLO_ROUTING", "rr").lower()
+
 
 def get_prefill_by_session(app, session_id: str):
     """Pick a (prefill_client, dp_rank) deterministically from session_id.
@@ -340,7 +388,58 @@ async def stream_service_response(
             yield chunk
 
 
+async def stream_colo_response(
+    colo_client_info: dict, endpoint: str, req_data: dict, headers: dict
+):
+    """Plain streaming passthrough to one colocated (kv_both) instance.
+
+    The request body is forwarded unchanged (stream/min_tokens/stream_options
+    all preserved) so the replayer's streaming + usage parsing works exactly
+    as it does when it talks to a colo instance directly.
+    """
+    async with colo_client_info["client"].stream(
+        "POST", endpoint, json=req_data, headers=headers
+    ) as response:
+        response.raise_for_status()
+        async for chunk in response.aiter_bytes():
+            yield chunk
+
+
+async def _handle_colo(api: str, request: Request):
+    if not app.state.ready.is_set():
+        raise HTTPException(status_code=503, detail="Service Unavailable")
+
+    req_data = await request.json()
+    request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
+    headers = {"X-Request-Id": request_id}
+    session_id = request.headers.get("X-Session-Id")
+    if session_id:
+        headers["X-Session-Id"] = session_id
+    key = os.environ.get("OPENAI_API_KEY")
+    if key:
+        headers["Authorization"] = f"Bearer {key}"
+
+    if MB5_COLO_ROUTING == "session" and session_id:
+        # consistent hash -> same kv_both instance reuses its prefix cache
+        h = int(hashlib.md5(session_id.encode()).hexdigest()[:8], 16)
+        idx = h % len(app.state.colo_clients)
+    else:
+        idx = next(app.state.colo_iterator)
+    colo_client_info = app.state.colo_clients[idx]
+
+    async def generate_stream():
+        async for chunk in stream_colo_response(
+            colo_client_info, api, req_data, headers
+        ):
+            yield chunk
+
+    return StreamingResponse(generate_stream(), media_type="text/event-stream")
+
+
 async def _handle_completions(api: str, request: Request):
+    if getattr(global_args, "colo", None):
+        return await _handle_colo(api, request)
+
     if not app.state.ready.is_set():
         raise HTTPException(status_code=503, detail="Service Unavailable")