Signed-off-by: inkcherry <mingzhi.liu@amd.com>
This commit is contained in:
inkcherry 2025-11-20 10:25:06 +00:00
parent 4034937733
commit bba4c89ca4

View File

@ -9,15 +9,13 @@ import socket
import threading import threading
import uuid import uuid
import aiohttp
import msgpack import msgpack
import zmq import zmq
from quart import Quart, make_response, request from quart import Quart, make_response, request
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG) logger.setLevel(logging.DEBUG)
import aiohttp
prefill_instances = [] prefill_instances = []
decode_instances = [] decode_instances = []
request_nums = 0 request_nums = 0
@ -69,15 +67,21 @@ def _listen_for_register(hostname, port):
data = msgpack.loads(msg) data = msgpack.loads(msg)
if data["type"] == "HELLO": if data["type"] == "HELLO":
pass pass
elif data["type"] == "register" and data["role"] == "P": elif (
if data["request_address"] not in prefill_instances: data["type"] == "register"
with _list_lock: and data["role"] == "P"
_append_whole_dict_unique(prefill_instances, data) and data["request_address"] not in prefill_instances
):
with _list_lock:
_append_whole_dict_unique(prefill_instances, data)
elif data["type"] == "register" and data["role"] == "D": elif (
if data["request_address"] not in decode_instances: data["type"] == "register"
with _list_lock: and data["role"] == "D"
_append_whole_dict_unique(decode_instances, data) and data["request_address"] not in decode_instances
):
with _list_lock:
_append_whole_dict_unique(decode_instances, data)
def start_service_discovery(hostname, port): def start_service_discovery(hostname, port):
@ -133,7 +137,7 @@ async def send_request_to_prefill(
else: else:
raise RuntimeError( raise RuntimeError(
"send_request_to_prefill response.status != 200,response.statuus = ", "send_request_to_prefill response.status != 200response.status = ",
response.status, response.status,
) )