[PD] Share the prefill->decode failure notification across backends (#36612)

Co-authored-by: inkcherry <mingzhi.liu@amd.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
jambow0320
2026-09-11 23:16:17 +08:00
committed by GitHub
co-authored by inkcherry Shangming Cai
parent e016de462c
commit 2a46cf2ca0
6 changed files with 453 additions and 368 deletions
@@ -165,7 +165,7 @@ class TestNixlTransferInfo(CustomTestCase):
]
)
self.assertFalse(info.is_dummy())
self.assertFalse(info.is_dummy)
def test_empty_indices_without_decode_prefix_is_dummy(self):
info = TransferInfo.from_zmq(
@@ -182,7 +182,7 @@ class TestNixlTransferInfo(CustomTestCase):
]
)
self.assertTrue(info.is_dummy())
self.assertTrue(info.is_dummy)
def test_explicit_dummy_frame_true_is_dummy(self):
# msg[9] is the explicit is_dummy frame the sender writes
@@ -848,10 +848,9 @@ class TestNixlReceiverPoll(CustomTestCase):
mgr.update_transfer_status.assert_called_once_with()
mgr.record_failure.assert_not_called()
mgr.update_status.assert_not_called()
self.assertNotIn(11, mgr.transfer_statuses)
@patch("sglang.srt.disaggregation.nixl.conn.time.time")
def test_transfer_done_returns_success_and_cleans_room_state(self, mock_time):
def test_transfer_done_returns_success_and_clear_drops_room_state(self, mock_time):
mock_time.return_value = 12.0
receiver, mgr = self._make_receiver(status=KVPoll.WaitingForInput)
receiver.started_transfer = True
@@ -864,9 +863,16 @@ class TestNixlReceiverPoll(CustomTestCase):
mgr.check_transfer_done.return_value = True
self.assertEqual(receiver.poll(), KVPoll.Success)
self.assertEqual(receiver.conclude_state, KVPoll.Success)
# poll() only concludes now; dropping room state is left to clear(), the
# way mooncake and mori already do it. The scheduler calls clear() as
# soon as poll() reports Success or Failed, so both terminal paths clean
# up -- the cleanup that used to live in poll() ran on Success only.
receiver.clear()
self.assertNotIn(11, mgr.transfer_statuses)
self.assertNotIn(11, mgr.addr_to_rooms_tracker["prefill:8998"])
self.assertEqual(receiver.conclude_state, KVPoll.Success)
class TestNixlNodeFailure(CustomTestCase):