diff --git a/microbench/fresh_setup/mb5_launch.sh b/microbench/fresh_setup/mb5_launch.sh index d530d2c..deb96ac 100755 --- a/microbench/fresh_setup/mb5_launch.sh +++ b/microbench/fresh_setup/mb5_launch.sh @@ -190,7 +190,9 @@ for port in "${all_ports[@]}"; do done if [ "${ROLES}" = "pd" ]; then - echo "[mb5] launching mooncake_connector_proxy on ${PROXY_PORT}" + P_ROUTING="${MB5_P_ROUTING:-rr}" + echo "[mb5] launching mooncake_connector_proxy on ${PROXY_PORT} (P routing=${P_ROUTING})" + MB5_P_ROUTING="${P_ROUTING}" \ nohup python "${PROXY_SRC}" "${proxy_args[@]}" --port "${PROXY_PORT}" --host 0.0.0.0 \ > "${LOGS_DIR}/proxy.log" 2>&1 & disown diff --git a/microbench/fresh_setup/mb5_pd_proxy.py b/microbench/fresh_setup/mb5_pd_proxy.py index d28e3fa..498e99f 100644 --- a/microbench/fresh_setup/mb5_pd_proxy.py +++ b/microbench/fresh_setup/mb5_pd_proxy.py @@ -3,6 +3,7 @@ import argparse import asyncio +import hashlib import ipaddress import itertools import os @@ -227,6 +228,28 @@ def _parse_decode_urls(decode_list): return [url[0] for url in decode_list] +# MB5: routing mode for the prefill (producer) side. +# "rr" — round-robin (official upstream behavior) +# "session" — consistent hash on X-Session-Id, so all turns of a session +# land on the same producer and reuse its prefix cache. +# Decode side stays round-robin (load balance) regardless. +MB5_P_ROUTING = os.environ.get("MB5_P_ROUTING", "rr").lower() + + +def get_prefill_by_session(app, session_id: str): + """Pick a (prefill_client, dp_rank) deterministically from session_id. + + Uses a stable (non-PYTHONHASHSEED-dependent) hash so the mapping is + reproducible across processes. dp_size is usually 1 here (TP=1, no DP), + but we hash into the flat (client, dp_rank) slot space to stay correct + if a producer ever reports dp_size > 1. + """ + clients = app.state.prefill_clients + slots = [(c, r) for c in clients for r in range(max(1, c.get("dp_size", 1)))] + h = int(hashlib.md5(session_id.encode()).hexdigest()[:8], 16) + return slots[h % len(slots)] + + def get_next_client(app, service_type: str): """ Get the next client in round-robin fashion. @@ -325,8 +348,17 @@ async def _handle_completions(api: str, request: Request): req_data = await request.json() request_id = str(uuid.uuid4()) - # Get the next prefill client in round-robin fashion - prefill_client_info, prefill_dp_rank = get_next_client(request.app, "prefill") + # Select the prefill (producer) client. + if MB5_P_ROUTING == "session": + session_id = request.headers.get("X-Session-Id") or request_id + prefill_client_info, prefill_dp_rank = get_prefill_by_session( + request.app, session_id + ) + else: + # Round-robin (official upstream behavior). + prefill_client_info, prefill_dp_rank = get_next_client( + request.app, "prefill" + ) # Send request to prefill service asyncio.create_task(