Function _pd_handle_task

lightllm/server/httpserver/pd_loop.py:62–126

pd_handle_loop 主要负责与 pd master 进行注册连接，然后接收pd master发来的请求，然后 将推理结果转发给 pd master进行处理。

(manager: HttpServerManager, pd_master_obj: PD_Master_Obj)

async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_Obj):
    """
    pd_handle_loop 主要负责与 pd master 进行注册连接，然后接收pd master发来的请求，然后
    将推理结果转发给 pd master进行处理。
    """
    # 创建转发队列
    forwarding_queue = AsyncQueue()

    while True:
        forwarding_tokens_task = None
        try:
            uri = f"ws://{pd_master_obj.host_ip_port}/pd_register"
            async with websockets.connect(
                uri, max_size=get_lightllm_websocket_max_message_size(), max_queue=(2048 * 1024, 2048 * 1023)  # 关键修改
            ) as websocket:

                sock = websocket.transport.get_extra_info("socket")
                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

                args_dict = vars(manager.args)
                args_dict["host"] = manager.host_ip
                # 发送注册信息
                regist_json = {
                    "node_id": manager.args.pd_node_id,
                    "client_ip_port": f"{manager.host_ip}:{manager.args.port}",
                    "mode": manager.pd_mode.value,
                    "start_args": args_dict,
                }

                await websocket.send(json.dumps(regist_json))
                logger.info(f"Sent registration JSON: {regist_json}")

                # 转发任务
                forwarding_tokens_task = asyncio.create_task(_up_tokens_to_pd_master(forwarding_queue, websocket))

                # 接收 pd master 发来的请求，并推理后，将生成的token转发回pd master。
                while True:
                    recv_bytes = await websocket.recv()
                    obj = pickle.loads(recv_bytes)
                    if obj[0] == ObjType.REQ:
                        prompt, sampling_params, multimodal_params = obj[1]
                        asyncio.create_task(
                            _pd_process_generate(manager, prompt, sampling_params, multimodal_params, forwarding_queue)
                        )
                    elif obj[0] == ObjType.ABORT:
                        group_req_id = obj[1]
                        await manager.abort(group_req_id)
                    else:
                        logger.error(f"recevie error obj {str(obj)}")

        except asyncio.CancelledError:
            # 如果任务被取消，则退出循环
            logger.warning(f"forwarding_tokens_task {pd_master_obj} cancelled")
            if forwarding_tokens_task is not None:
                forwarding_tokens_task.cancel()
            return

        except Exception as e: