"""Prefill/decode-disaggregated serving: a prompt-only engine tray (PE) and a decode-only engine tray (CE).

    # decode tray (user-facing OpenAI API on --port, KV receivers on --kv-port + rank):
    python serve/pd_server.py --role ce --model ... --dp 4 --port 8888 --kv-port 8900 --pe-url http://<pe-host>:8889
    # prefill tray:
    python serve/pd_server.py --role pe --model ... --dp 4 --port 8889

Flow per request: CE frontend renders the prompt, picks a CE rank, POSTs /prefill to the PE with the KV target
(ce-host, kv-port + rank). The PE prefills (host prefix cache on), exports the sequence's KV blocks + full SWA
state (deepinfer.hicache.pd_transfer) and the first sampled token, streams the first token back in the HTTP
response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence straight into decode
(Scheduler.add_prompted_sequence); its tokens stream to the client through the CE frontend.
"""

from __future__ import annotations

import argparse
import asyncio
import collections
import json
import multiprocessing as mp
import os
import queue
import socket
import struct
import sys
import threading
import time
import uuid
from collections.abc import Callable
from pathlib import Path
from types import ModuleType
from typing import TYPE_CHECKING

from fastapi import FastAPI, Request

if TYPE_CHECKING:
    from deepinfer._C.engine import RequestOutput, Sequence
    from deepinfer.engine.llm_engine import LLMEngine
    from deepinfer.hicache.pd_transfer import PDTransfer
    from deepinfer.sampling_params import SamplingParams

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
sys.path.insert(0, str(HERE.parent / "examples"))

import openai_server as oas  # noqa: E402  (the single-engine server: prompt rendering, detokenizer, frontend plumbing)

HDR = struct.Struct("<16sqqqqq")  # rid(16 bytes hex) | first_token | prompt_len | num_pages | page_bytes | state_bytes


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--role", required=True, choices=["pe", "ce"])
    p.add_argument("--model", required=True)
    p.add_argument("--served-model-name", default=None)
    p.add_argument("--dp", type=int, default=4)
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8888)
    p.add_argument("--kv-port", type=int, default=8900, help="ce: TCP receiver base port (+ rank)")
    p.add_argument("--kv-host", default=None, help="ce: address the PE should push KV to (default: this host's name)")
    p.add_argument("--pe-url", default=None, help="ce: prefill tray URL")
    p.add_argument("--max-model-len", type=int, default=1 << 20)
    p.add_argument("--max-num-seqs", type=int, default=128)
    p.add_argument("--max-num-batched-tokens", type=int, default=32768)
    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
    p.add_argument("--yoco-keep-tokens", type=int, default=140)
    p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"])
    p.add_argument("--log-interval-steps", type=int, default=200)
    p.add_argument("--hicache-size-gb", type=float, default=32, help="pe: host prefix cache per rank (0 disables)")
    p.add_argument("--hicache-state-slots", type=int, default=4096)
    p.add_argument("--routing", default="affinity", choices=["affinity", "least-loaded"])
    p.add_argument("--prefill-any-interval", type=int, default=2)
    p.add_argument("--eplb-interval", type=int, default=0, help="see openai_server.py; 0 (default) disables EPLB")
    p.add_argument("--eplb-after-prefill", type=int, default=0, choices=[0, 1], help="see openai_server.py")
    return p.parse_args()


def _recv_exact(sock: socket.socket, buf: memoryview) -> None:
    view = buf
    while len(view):
        n = sock.recv_into(view, min(len(view), 1 << 24))
        if n == 0:
            raise ConnectionError("peer closed during transfer")
        view = view[n:]


# ============================================================================ engine processes


def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None:
    import ctypes
    import signal

    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL)  # PR_SET_PDEATHSIG

    import numpy as np
    from _config import create_llm, setup_common_env

    from deepinfer.engine.engine_type import EngineType

    if args.yoco_keep_tokens > 0:
        os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact"
        os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens)
        os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1"
    os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512")  # identical state layout on PE and CE
    os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval))
    setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False)
    if args.eplb_interval > 0:
        os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval)
        os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill)

    engine_type = EngineType.PROMPT if args.role == "pe" else EngineType.COMPLETION
    llm = create_llm(
        args.model,
        dp_rank,
        ep_size=args.dp,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)},
        data_parallel_size=args.dp,
        max_num_seqs=args.max_num_seqs,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
        engine_type=engine_type,
        stop_preempted_request=True,
    )

    from deepinfer._C.engine import RequestOutput  # pyright: ignore[reportMissingModuleSource]
    from deepinfer.hicache.pd_transfer import PDTransfer
    from deepinfer.sampling_params import SamplingParams

    engine = llm.llm_engine
    transfer = PDTransfer(engine)
    mm = lambda: dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None  # noqa: E731

    def enqueue(rid: str, seq: Sequence, num_cached: int) -> None:
        engine.scheduler.add_sequence(seq)
        if engine.multi_modal_processor is not None:
            engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached)

    if args.role == "pe":
        _pe_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm, enqueue)
    else:
        _ce_loop(args, dp_rank, in_q, out_q, engine, transfer, SamplingParams, RequestOutput, np, mm)


def _pe_loop(
    args: argparse.Namespace,
    dp_rank: int,
    in_q: mp.Queue,
    out_q: mp.Queue,
    engine: LLMEngine,
    transfer: PDTransfer,
    SamplingParams: type[SamplingParams],
    RequestOutput: type[RequestOutput],
    np: ModuleType,
    mm: Callable[[], dict | None],
    enqueue: Callable[[str, Sequence, int], None],
) -> None:
    cache = None
    if args.hicache_size_gb > 0:
        from deepinfer.hicache import HostPrefixCache

        cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots)

    targets: dict[str, tuple[str, int]] = {}  # rid -> (host, port) of the CE rank
    awaiting_token: dict[str, dict] = {}  # rid -> transfer waiting for the step's sampled first token
    send_q: queue.Queue = queue.Queue()

    def on_prefill_step(scheduled: list[Sequence], finished: list[Sequence]) -> None:
        # Engine hook at the start of post-processing for a prompt step: the prefill's KV is complete and the
        # sequence still owns its blocks (with max_tokens=1 it finishes, and is freed, once outputs are processed).
        if cache is not None:
            cache.on_prefill_step(scheduled, finished)
        for seq in scheduled:
            rid = seq.request_id
            if seq.chunked_rest_len != 0 or rid not in targets:
                continue
            prompt_len = seq.get_len()  # no output appended yet
            need = (prompt_len + transfer.block_size - 1) // transfer.block_size
            if len(seq.block_table) < need:
                out_q.put(("error", rid, f"prefilled sequence has {len(seq.block_table)} blocks, needs {need}"))
                targets.pop(rid, None)
                continue
            pages, state, ev = transfer.export(seq, prompt_len)
            item = {
                "pages": pages,
                "state": state,
                "event": ev,
                "prompt_len": prompt_len,
                "target": targets.pop(rid),
                "first_token": None,
                "ready": threading.Event(),
            }
            awaiting_token[rid] = item
            send_q.put((rid, item))

    engine.pe_do_save_kvcache = on_prefill_step

    def sender():
        while True:
            rid, item = send_q.get()
            host, port = item["target"]
            try:
                item["ready"].wait(timeout=600)
                if item["first_token"] is None:
                    raise RuntimeError("no first token produced for the prefilled sequence")
                item["event"].synchronize()
                pages, state = item["pages"], item["state"]
                with socket.create_connection((host, port), timeout=60) as s:
                    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                    s.sendall(
                        HDR.pack(rid.encode()[:16], item["first_token"], item["prompt_len"], pages.size(0), pages.size(1), state.numel())
                    )
                    s.sendall(memoryview(pages.numpy()).cast("B"))
                    s.sendall(memoryview(state.numpy()).cast("B"))
                    ack = s.recv(2)
                if ack != b"ok":
                    out_q.put(("error", rid, f"CE rejected transfer: {ack!r}"))
            except Exception as e:  # noqa: BLE001
                out_q.put(("error", rid, f"KV transfer to {host}:{port} failed: {e!r}"))

    threading.Thread(target=sender, daemon=True).start()
    out_q.put(("ready", dp_rank, None))

    pending = collections.deque()
    inflight = False
    last_stats = time.monotonic()
    while True:
        while True:
            try:
                kind, rid, payload = in_q.get_nowait()
            except queue.Empty:
                break
            if kind == "add":
                try:
                    prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32)
                    sp = SamplingParams(**payload["sampling"])
                    sp.max_tokens = 1  # prefill + first token only; the engine then finishes and frees the sequence
                    seq = engine.add_request(
                        user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=prompt_ids, multi_modal_data=mm(), enqueue=False
                    )
                    if seq is None:
                        out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
                    else:
                        targets[rid] = tuple(payload["kv_target"])
                        pending.append((rid, seq, prompt_ids))
                except Exception as e:  # noqa: BLE001
                    out_q.put(("error", rid, repr(e)))
            elif kind == "abort":
                targets.pop(rid, None)
                if any(p[0] == rid for p in pending):
                    pending = collections.deque(p for p in pending if p[0] != rid)
                else:
                    engine.abort_request(rid)
                if cache is not None:
                    cache.forget(rid)
        while pending:
            rid, seq, prompt_ids = pending[0]
            hit = cache.try_resume(seq, rid, prompt_ids) if cache is not None else 0
            if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
                break
            pending.popleft()
            enqueue(rid, seq, max(hit, 0))
        if cache is not None:
            cache.poll()
            if time.monotonic() - last_stats > 5:
                last_stats = time.monotonic()
                out_q.put(("stats", dp_rank, cache.stats_dict()))

        if engine.has_unfinished_requests():
            outputs = engine.step()
            inflight = True
        elif inflight:
            outputs = engine.flush_inflight_reqs()
            inflight = False
        else:
            time.sleep(0.001)
            continue
        if len(outputs):
            batch = [
                (
                    RequestOutput.get_request_id(p),
                    RequestOutput.get_token_id(p),
                    RequestOutput.get_finished(p),
                    RequestOutput.get_finish_reason(p),
                    RequestOutput.get_prompt_len(p),
                )
                for p in outputs
            ]
            for rid, token_id, finished, _, _ in batch:
                item = awaiting_token.pop(rid, None)
                if item is not None:
                    item["first_token"] = int(token_id)
                    item["ready"].set()
                if finished and cache is not None:
                    cache.forget(rid)
            out_q.put(("tokens", None, batch))
        # A prompt-only engine never reclaims finished sequences by itself (that happens in the decode
        # scheduler): drop them from `running` and free their blocks/state. The export and cache copies were
        # enqueued on the stream before any kernel that could reuse those blocks.
        for seq in engine.scheduler.free_finished_seqs():
            engine.scheduler.free_seq(seq)


def _ce_loop(
    args: argparse.Namespace,
    dp_rank: int,
    in_q: mp.Queue,
    out_q: mp.Queue,
    engine: LLMEngine,
    transfer: PDTransfer,
    SamplingParams: type[SamplingParams],
    RequestOutput: type[RequestOutput],
    np: ModuleType,
    mm: Callable[[], dict | None],
) -> None:
    import torch

    recv_q: queue.Queue = queue.Queue()  # (rid, first_token, prompt_len, pages, state)

    def receiver():
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind((args.host, args.kv_port + dp_rank))
        srv.listen(64)
        while True:
            conn, _ = srv.accept()
            threading.Thread(target=receive_one, args=(conn,), daemon=True).start()

    def receive_one(conn: socket.socket):
        try:
            with conn:
                hdr = bytearray(HDR.size)
                _recv_exact(conn, memoryview(hdr))
                rid_b, first_token, prompt_len, num_pages, page_bytes, state_bytes = HDR.unpack(hdr)
                rid = rid_b.decode()
                assert page_bytes == transfer.page_bytes and state_bytes == transfer.state_bytes, "layout mismatch between PE and CE"
                pages = torch.empty((num_pages, page_bytes), dtype=torch.uint8, pin_memory=True)
                state = torch.empty(state_bytes, dtype=torch.uint8, pin_memory=True)
                _recv_exact(conn, memoryview(pages.numpy()).cast("B"))
                _recv_exact(conn, memoryview(state.numpy()).cast("B"))
                recv_q.put((rid, first_token, prompt_len, pages, state))
                conn.sendall(b"ok")
        except Exception as e:  # noqa: BLE001
            print(f"[ce rank {dp_rank}] transfer failed: {e!r}", flush=True)

    threading.Thread(target=receiver, daemon=True).start()
    out_q.put(("ready", dp_rank, None))

    expected: dict[str, dict] = {}  # rid -> {"prompt_token_ids", "sampling"} from the frontend
    arrived: dict[str, tuple] = {}  # rid -> transfer
    pending: collections.deque = collections.deque()  # rids with both parts, waiting for GPU allocation
    keep: dict[str, torch.cuda.Event] = {}
    inflight = False
    while True:
        while True:
            try:
                kind, rid, payload = in_q.get_nowait()
            except queue.Empty:
                break
            if kind == "expect":
                expected[rid] = payload
                if rid in arrived:
                    pending.append(rid)
            elif kind == "abort":
                expected.pop(rid, None)
                arrived.pop(rid, None)
                if rid in pending:
                    pending.remove(rid)
                else:
                    engine.abort_request(rid)
        while True:
            try:
                rid, first_token, prompt_len, pages, state = recv_q.get_nowait()
            except queue.Empty:
                break
            arrived[rid] = (first_token, prompt_len, pages, state)
            if rid in expected:
                pending.append(rid)
        while pending:
            rid = pending[0]
            spec, (first_token, prompt_len, pages, state) = expected[rid], arrived[rid]
            tokens = np.asarray(list(spec["prompt_token_ids"]) + [first_token], dtype=np.int32)
            assert len(tokens) - 1 == prompt_len, (len(tokens), prompt_len)
            sp = SamplingParams(**spec["sampling"])
            sp.max_tokens = max(1, sp.max_tokens - 1)  # the first token was produced on the PE
            seq = engine.add_request(
                user_id="0", request_id=rid, sampling_params=sp, prompt_token_ids=tokens, multi_modal_data=mm(), enqueue=False
            )
            if seq is None:
                out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
                pending.popleft()
                expected.pop(rid, None)
                arrived.pop(rid, None)
                continue
            try:
                allocated = engine.scheduler.try_allocate_seq(seq)
            except RuntimeError as e:
                out_q.put(("error", rid, f"cannot allocate: {e!r}"))
                allocated = None
            if allocated is False and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
                break  # retry once something finishes
            pending.popleft()
            expected.pop(rid, None)
            arrived.pop(rid, None)
            if not allocated:
                if allocated is False:
                    out_q.put(("error", rid, "cannot allocate KV blocks for the transferred sequence"))
                continue
            # Keep the pinned receive buffers alive until the H2D page-copy kernels have consumed them: the kernel
            # reads them by raw pointer, so the caching host allocator would otherwise hand the memory to the next
            # transfer (or unmap it) while the copy is still queued -> corrupted KV / cudaErrorLaunchFailure.
            keep[rid] = (transfer.admit(seq, pages, state, prompt_len), pages, state)
            if engine.multi_modal_processor is not None:
                engine.multi_modal_processor.add(rid, seq.multi_modal_data, prompt_len)
        for rid in [r for r, (ev, _, _) in keep.items() if ev.query()]:
            del keep[rid]

        if engine.has_unfinished_requests():
            outputs = engine.step()
            inflight = True
        elif inflight:
            outputs = engine.flush_inflight_reqs()
            inflight = False
        else:
            time.sleep(0.001)
            continue
        if len(outputs):
            batch = [
                (
                    RequestOutput.get_request_id(p),
                    RequestOutput.get_token_id(p),
                    RequestOutput.get_finished(p),
                    RequestOutput.get_finish_reason(p),
                    RequestOutput.get_prompt_len(p),
                )
                for p in outputs
            ]
            out_q.put(("tokens", None, batch))


# ============================================================================ frontends


class PDFrontend(oas.Frontend):
    def __init__(self, args: argparse.Namespace) -> None:
        super().__init__(args)
        ctx = mp.get_context("spawn")
        self.out_q = ctx.Queue()
        self.in_qs = [ctx.Queue() for _ in range(args.dp)]
        self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)]


def build_pe_app(fe: PDFrontend) -> FastAPI:
    """Prefill tray: POST /prefill {rid, prompt_token_ids, sampling, kv_target:[host,port]} -> {first_token}."""
    from fastapi import FastAPI
    from fastapi.responses import JSONResponse

    app = FastAPI()

    @app.on_event("startup")
    async def _startup() -> None:
        fe.loop = asyncio.get_running_loop()
        threading.Thread(target=fe.dispatch_forever, daemon=True).start()

    @app.get("/health")
    async def health():
        ok = fe.ready == fe.args.dp
        return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready, "role": "pe"}, status_code=200 if ok else 503)

    @app.get("/deepinfer/metrics")
    async def metrics():
        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
        for r, st in sorted(fe.cache_stats.items()):
            lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()]
        return "\n".join(lines) + "\n"

    @app.post("/prefill")
    async def prefill(request: Request):
        if fe.ready != fe.args.dp:
            return JSONResponse({"error": "engine loading"}, status_code=503)
        body = await request.json()
        rid = body["rid"]
        rank = min(range(fe.args.dp), key=lambda r: fe.load[r])
        key = bytes.fromhex(body["affinity"]) if body.get("affinity") else None
        if fe.args.routing == "affinity" and key is not None:
            rank = int.from_bytes(key[:4], "little") % fe.args.dp
        fe.load[rank] += 1
        fe.owner[rid] = rank
        q: asyncio.Queue = asyncio.Queue()
        fe.streams[rid] = q
        fe.stats["requests"] += 1
        fe.in_qs[rank].put(
            ("add", rid, {"prompt_token_ids": body["prompt_token_ids"], "sampling": body["sampling"], "kv_target": body["kv_target"]})
        )
        try:
            item = await asyncio.wait_for(q.get(), timeout=1800)
        finally:
            fe.release(rid, aborted=False)
        if item[0] == "error":
            return JSONResponse({"error": item[1]}, status_code=500)
        return {"first_token": item[1], "prompt_len": item[4]}

    return app


def build_ce_app(fe: PDFrontend, pe_url: str, kv_host: str) -> FastAPI:
    """Decode tray: the user-facing OpenAI API."""
    import httpx
    from fastapi import FastAPI
    from fastapi.responses import JSONResponse, StreamingResponse

    app = FastAPI()
    client = httpx.AsyncClient(timeout=1800)

    @app.on_event("startup")
    async def _startup() -> None:
        fe.loop = asyncio.get_running_loop()
        threading.Thread(target=fe.dispatch_forever, daemon=True).start()

    @app.get("/health")
    async def health():
        ok = fe.ready == fe.args.dp
        return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready, "role": "ce"}, status_code=200 if ok else 503)

    @app.get("/v1/models")
    async def models():
        return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]}

    @app.get("/deepinfer/metrics")
    async def metrics():
        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
        lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)]
        return "\n".join(lines) + "\n"

    @app.post("/v1/chat/completions")
    async def chat(request: Request):
        if fe.ready != fe.args.dp:
            return JSONResponse({"error": {"message": "engine loading"}}, status_code=503)
        body = await request.json()
        prompt_ids = await asyncio.to_thread(fe.render_prompt, body)
        max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096
        if len(prompt_ids) + 1 > fe.args.max_model_len:
            return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400)
        max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids)))
        sampling = dict(
            temperature=float(body.get("temperature", 1.0)),
            top_p=float(body.get("top_p", 1.0)),
            ignore_eos=bool(body.get("ignore_eos", False)),
            max_tokens=max_tokens,
        )
        if body.get("seed") is not None:
            sampling["seq_seed"] = int(body["seed"])
        fe.stats["prompt_tokens"] += len(prompt_ids)
        include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
        key = fe.affinity_key(body)

        # Register with the CE rank first so the transfer can be admitted as soon as it lands.
        rid = uuid.uuid4().hex[:16]
        rank = (
            int.from_bytes(key[:4], "little") % fe.args.dp
            if (fe.args.routing == "affinity" and key)
            else min(range(fe.args.dp), key=lambda r: fe.load[r])
        )
        fe.load[rank] += 1
        fe.owner[rid] = rank
        q: asyncio.Queue = asyncio.Queue()
        fe.streams[rid] = q
        fe.stats["requests"] += 1
        fe.in_qs[rank].put(("expect", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling}))

        created = int(time.time())
        cid = f"chatcmpl-{rid}"
        detok = oas.IncrementalDetokenizer(fe.tokenizer)

        def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str:
            obj = {
                "id": cid,
                "object": "chat.completion.chunk",
                "created": created,
                "model": fe.model_name,
                "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}],
            }
            if usage:
                obj["usage"] = usage
            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"

        async def prefill_first():
            r = await client.post(
                f"{pe_url}/prefill",
                json={
                    "rid": rid,
                    "prompt_token_ids": prompt_ids,
                    "sampling": sampling,
                    "kv_target": [kv_host, fe.args.kv_port + rank],
                    "affinity": key.hex() if key else None,
                },
            )
            if r.status_code != 200:
                raise RuntimeError(f"PE: {r.text[:200]}")
            return int(r.json()["first_token"])

        async def tokens():
            first = await prefill_first()
            text = detok.add(first)
            fe.stats["completion_tokens"] += 1
            if max_tokens == 1 or first == fe.tokenizer.eos_token_id and not sampling["ignore_eos"]:
                yield text, ("length" if max_tokens == 1 else "stop")
                return
            yield text, None
            while True:
                item = await q.get()
                if item[0] == "error":
                    raise RuntimeError(item[1])
                _, token_id, finished, reason, _ = item
                text = detok.add(token_id)
                fe.stats["completion_tokens"] += 1
                if finished:
                    yield text, ("length" if detok.num_tokens >= max_tokens else "stop")
                    return
                yield text, None

        usage = lambda: {
            "prompt_tokens": len(prompt_ids),
            "completion_tokens": detok.num_tokens,  # noqa: E731
            "total_tokens": len(prompt_ids) + detok.num_tokens,
        }

        if body.get("stream"):

            async def sse():
                done = False
                try:
                    yield chunk({"role": "assistant", "content": ""})
                    async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
                        if text or finish:
                            yield chunk({"content": text} if text else {}, finish)
                    done = True
                    if include_usage:
                        yield chunk({}, usage=usage())
                    yield "data: [DONE]\n\n"
                finally:
                    fe.release(rid, aborted=not done)

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

        done = False
        try:
            parts, finish = [], None
            async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
                parts.append(text)
            done = True
        finally:
            fe.release(rid, aborted=not done)
        return {
            "id": cid,
            "object": "chat.completion",
            "created": created,
            "model": fe.model_name,
            "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}],
            "usage": usage(),
        }

    return app


def main() -> None:
    args = parse_args()
    args.model = str(Path(args.model).expanduser().resolve())
    os.environ.setdefault("WORLD_SIZE", "1")
    os.environ.setdefault("RANK", "0")
    fe = PDFrontend(args)
    fe.start()
    import uvicorn

    if args.role == "pe":
        app = build_pe_app(fe)
    else:
        assert args.pe_url, "--pe-url is required for the decode tray"
        app = build_ce_app(fe, args.pe_url.rstrip("/"), args.kv_host or socket.gethostname())
    uvicorn.run(app, host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900)


if __name__ == "__main__":
    main()