Repository: LMCache/LMCache
File: examples/disagg_prefill/disagg_proxy_server.py (Python)

Code before:
    return (tokenization_client, prefill_client, decode_client)

async def wait_decode_kv_ready(req_id: str, num_tp_rank: int):
    while app.state.finished_reqs[req_id] < num_tp_rank:
        await asyncio.sleep(0.0001)  # sleep for 0.1 ms
    logger.debug(f"Prefill node signaled kv ready for req {req_id}")
    app.state.finished_reqs.pop(req_id)

counter_with_bounds = 0
def pick_up_clients(req_data):
    if not global_args.bound_client:
        global counter
        tokenization_client, prefill_client, decode_client = round_robin_pick_clients(counter)
    else:
        client_id_key = os.getenv("CLIENT_BOUND_KEY", "session_id")
        client_id = req_data.get(client_id_key)
        if client_id and client_id in app.state.bound_clients:
            tokenization_client, prefill_client, decode_client = app.state.bound_clients[client_id]
        else:
            # when use bound-client option, there is a need to use a different counter,
            global counter_with_bounds
            counter_with_bounds += 1
            tokenization_client, prefill_client, decode_client = round_robin_pick_clients(counter_with_bounds)
            if client_id:
                app.state.bound_clients[client_id] = (tokenization_client, prefill_client, decode_client)
    return (tokenization_client, prefill_client, decode_client)

@app.post("/v1/completions")
async def handle_completions(request: Request):
    global counter, stats_calculator
    counter += 1
    req_id = str(counter)  # we use counter as req_id

    st = time.time()
    try:
        req_data = await request.json()

        # Pick tokenization, prefill and decode client
        tokenization_client, prefill_client, decode_client = pick_up_clients(req_data)

        tokenize_output = await send_request_to_service(
            tokenization_client.client, "/tokenize", {"prompt": req_data["prompt"]}
        )
        tokenize_output = tokenize_output.json()

        org_max_tokens = req_data["max_tokens"]
        req_data["prompt"] = tokenize_output["tokens"]
        req_data["max_tokens"] = 1


        disagg_spec = {

Review comment (suggestion):
This function's logic seems to be a bit redundant to me. I think it can be simplified as the following:
Check if `client_id` is in `req_data` and in `bound_clients`. If not, Just do round robin.

Code after:
def round_robin_pick_clients() -> tuple[ClientInfo, ClientInfo, ClientInfo]:
    idx = next(round_robin_counter)
    tokenization_client = round_robin_pick_client(app.state.total_clients, idx)
    prefill_client = round_robin_pick_client(app.state.prefill_clients, idx)
    decode_client = round_robin_pick_client(app.state.decode_clients, idx)
    return tokenization_client, prefill_client, decode_client


async def wait_decode_kv_ready(req_id: str, num_tp_rank: int):
    while app.state.finished_reqs[req_id] < num_tp_rank:
        await asyncio.sleep(0.0001)  # sleep for 0.1 ms
    logger.debug(f"Prefill node signaled kv ready for req {req_id}")
    app.state.finished_reqs.pop(req_id)


BOUND_CLIENTS_MAX_NUM = 1024 * 1024


def pick_up_bound_clients(client_id: str) -> tuple[ClientInfo, ClientInfo, ClientInfo]:
    if client_id not in app.state.bound_clients:
        if len(app.state.bound_clients) >= BOUND_CLIENTS_MAX_NUM:
            # Here simply clear the bound_clients if full
            app.state.bound_clients.clear()
        app.state.bound_clients[client_id] = round_robin_pick_clients()
    return app.state.bound_clients[client_id]


BOUND_CLIENT = os.getenv("CLIENT_BOUND", "false").lower() == "true"
# CLIENT_BOUND_KEY, the field name of the client uid in http request
CLIENT_BOUND_KEY = os.getenv("CLIENT_BOUND_KEY", "session-id")


def pick_up_clients(request: Request) -> tuple[ClientInfo, ClientInfo, ClientInfo]:
    bound_client_id = request.headers.get(CLIENT_BOUND_KEY) if BOUND_CLIENT else None
    if bound_client_id:
        # Use or create a persistent set of clients for the session.
        return pick_up_bound_clients(bound_client_id)
    return round_robin_pick_clients()


@app.post("/v1/completions")
async def handle_completions(request: Request):
    global counter, stats_calculator
    counter += 1
    req_id = str(counter)  # we use counter as req_id

    st = time.time()
    try:
        req_data = await request.json()

        # Pick tokenization, prefill and decode client