Skip to content

[disagg][mooncake] Keep decode status thread alive on bad messages #32909

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
guptaishaan wants to merge 1 commit into from
Closed

[disagg][mooncake] Keep decode status thread alive on bad messages #32909

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 69 additions & 60 deletions python/sglang/srt/disaggregation/mooncake/conn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1968,68 +1968,77 @@ def start_decode_thread(self):
def decode_thread():
while True:
msg = self.server_socket.recv_multipart()
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
self._handle_aux_data(msg)
continue

# Staging: prefill notifies a chunk written to staging buffer
if msg[0] == b"CHUNK_READY":
room = int(msg[1].decode("ascii"))
chunk_idx = int(msg[2].decode("ascii"))
page_start = int(msg[3].decode("ascii"))
num_pages = int(msg[4].decode("ascii"))
session_id = msg[5].decode("ascii")
handler = self._staging_handler
assert (
handler is not None
), "CHUNK_READY received before staging handler initialized"
handler.handle_chunk_arrived(
room,
chunk_idx,
page_start,
num_pages,
session_id,
)
continue

# Staging: prefill pre-requests staging allocation before forward
if msg[0] == b"STAGING_REQ":
self._handle_staging_req(msg)
continue

# Prefill acknowledges abort notification
if msg[0] == b"ABORT_ACK":
# TODO(shangming): use this info to implement the deferred release mechanism if needed
ack_aborted_room = int(msg[1].decode("ascii"))
logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}")
continue

bootstrap_room, status, prefill_rank = msg
status = int(status.decode("ascii"))
bootstrap_room = int(bootstrap_room.decode("ascii"))
prefill_rank = int(prefill_rank.decode("ascii"))

if status == KVPoll.Success:
if bootstrap_room in self.request_status:
self.prefill_response_tracker[bootstrap_room].add(prefill_rank)
expected_response_num = (
self.required_prefill_response_num_table[bootstrap_room]
try:
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
self._handle_aux_data(msg)
continue

# Staging: prefill notifies a chunk written to staging buffer
if msg[0] == b"CHUNK_READY":
room = int(msg[1].decode("ascii"))
chunk_idx = int(msg[2].decode("ascii"))
page_start = int(msg[3].decode("ascii"))
num_pages = int(msg[4].decode("ascii"))
session_id = msg[5].decode("ascii")
handler = self._staging_handler
assert (
handler is not None
), "CHUNK_READY received before staging handler initialized"
handler.handle_chunk_arrived(
room,
chunk_idx,
page_start,
num_pages,
session_id,
)
arrived_response_num = len(
self.prefill_response_tracker[bootstrap_room]
continue

# Staging: prefill pre-requests staging allocation before forward
if msg[0] == b"STAGING_REQ":
self._handle_staging_req(msg)
continue

# Prefill acknowledges abort notification
if msg[0] == b"ABORT_ACK":
# TODO(shangming): use this info to implement the deferred release mechanism if needed
ack_aborted_room = int(msg[1].decode("ascii"))
logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}")
continue

bootstrap_room, status, prefill_rank = msg
status = int(status.decode("ascii"))
bootstrap_room = int(bootstrap_room.decode("ascii"))
prefill_rank = int(prefill_rank.decode("ascii"))

if status == KVPoll.Success:
if bootstrap_room in self.request_status:
self.prefill_response_tracker[bootstrap_room].add(
prefill_rank
)
expected_response_num = (
self.required_prefill_response_num_table[bootstrap_room]
)
arrived_response_num = len(
self.prefill_response_tracker[bootstrap_room]
)
if arrived_response_num == expected_response_num:
if self.enable_staging:
handler = self._staging_handler
if handler.is_staging_room(bootstrap_room):
handler.submit_last_scatter_async(
bootstrap_room
)
self.update_status(bootstrap_room, KVPoll.Success)
elif status == KVPoll.Failed:
self.record_failure(
bootstrap_room,
"Failed to get kvcache from prefill instance, it might be dead",
)
if arrived_response_num == expected_response_num:
if self.enable_staging:
handler = self._staging_handler
if handler.is_staging_room(bootstrap_room):
handler.submit_last_scatter_async(bootstrap_room)
self.update_status(bootstrap_room, KVPoll.Success)
elif status == KVPoll.Failed:
self.record_failure(
bootstrap_room,
"Failed to get kvcache from prefill instance, it might be dead",
)
self.update_status(bootstrap_room, status)
self.update_status(bootstrap_room, status)
except Exception:
# The socket is reachable by any peer, a single malformed
# message must not take down the decode status thread.
logger.exception("Decode status thread failed to handle a message")

threading.Thread(target=decode_thread).start()
self._start_heartbeat_checker_thread()
Expand Down
75 changes: 75 additions & 0 deletions test/registered/unit/disaggregation/test_mooncake_decode_thread.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import threading
import time
import unittest
from collections import defaultdict

import zmq

from sglang.srt.disaggregation.base.conn import KVPoll
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=5, suite="base-a-test-cpu")

ROOM = 42


class _ExitingSocket:
"""Stops the decode thread at teardown: threading swallows SystemExit."""

def recv_multipart(self):
raise SystemExit


class TestMooncakeDecodeThread(CustomTestCase):
"""The decode KV manager binds a plain zmq.PULL socket, so any peer that can
reach it may send a message that does not match the status wire format. Such
a message used to raise out of the decode status thread and kill it; the
process stayed up and healthy while every later KV-transfer completion went
unprocessed and requests hung until they timed out.
"""

def test_decode_thread_survives_malformed_message(self):
ctx = zmq.Context()
server_socket = ctx.socket(zmq.PULL)
port = server_socket.bind_to_random_port("tcp://127.0.0.1")

kv_mgr = MooncakeKVManager.__new__(MooncakeKVManager)
kv_mgr.server_socket = server_socket
kv_mgr.request_status = {ROOM: KVPoll.WaitingForInput}
kv_mgr.prefill_response_tracker = defaultdict(set)
kv_mgr.required_prefill_response_num_table = {ROOM: 1}
kv_mgr.enable_staging = False
kv_mgr._staging_handler = None
kv_mgr.update_status = kv_mgr.request_status.__setitem__
kv_mgr._start_heartbeat_checker_thread = lambda: None
threads_before = set(threading.enumerate())
MooncakeKVManager.start_decode_thread(kv_mgr)
decode_threads = set(threading.enumerate()) - threads_before

sender = ctx.socket(zmq.PUSH)
sender.connect(f"tcp://127.0.0.1:{port}")
# A single frame that is neither a known header nor a status triple.
sender.send_multipart([b"\x00"])
# A legitimate transfer-done notification from prefill.
sender.send_multipart([str(ROOM).encode(), str(KVPoll.Success).encode(), b"0"])

deadline = time.time() + 10
while kv_mgr.request_status[ROOM] != KVPoll.Success and time.time() < deadline:
time.sleep(0.05)

try:
self.assertEqual(kv_mgr.request_status[ROOM], KVPoll.Success)
finally:
# SystemExit on the next recv ends the thread without a traceback.
kv_mgr.server_socket = _ExitingSocket()
sender.send_multipart([b"ABORT_ACK", b"0"])
for thread in decode_threads:
thread.join(timeout=5)
sender.close()
ctx.destroy(linger=0)


if __name__ == "__main__":
unittest.main()
Loading