Files
sglang/test/registered/unit/disaggregation/test_prefill_inflight_liveness.py
T
leaveletandClaude Fable 5 58f2350738 Give prefill inflight transfers a liveness bound
E2e caught one request wedged FOREVER in disagg_prefill_inflight_queue
(probes 21s apart with zero traffic both showed #inflight-req: 1, no
reap/timeout warnings ever logged).  Mechanism, established by reading
the full state machine: prefill Success is set locally by the transfer
worker on the LAST chunk; if the decode peer is torn down between the
handshake and the prefill's final send(), add_transfer_request silently
drops the chunk (no transfer destinations) — Success becomes
unreachable.  The only external rescue, the decode ABORT notification,
is best-effort (silently swallowed on send error, no-op if it races the
room registration), there is no prefill-side heartbeat of decode
sessions, and the sender's only timeout covers Bootstrapping — the
inflight queue itself has no liveness bound.  The orphan pins the
request's KV pages and rides every poll collective.

Two fixes, both reaped through the existing Failed branch via the
CP/TP MIN-reduce poll consensus (Failed=0 wins, so one rank concluding
flips every rank together — rank-uniform by construction):

- add_transfer_request: a room with no transfer destinations that is
  NOT already Success (the dummy-rank handshake marking) now concludes
  Failed loudly instead of dropping the chunk silently.
- Inflight residency timeout: entries stuck in a non-terminal poll
  state past SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT (default 300s,
  matching the sibling BOOTSTRAP/WAITING timeouts) get sender.abort()
  and reap on the next poll.  Covers what the hardening cannot: lost
  ABORT datagrams, decode crashes.

Known sibling gaps left for follow-up: the decode transfer queue has
no Transferring liveness bound, and an abort that matches no queue is
still a silent no-op (much narrower race than first thought — work
requests are ordered before control requests within a tick).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-12 04:55:35 +00:00

86 lines
3.2 KiB
Python

"""Unit tests for the PD prefill inflight liveness fixes.
Background (2026-06-12 e2e incident): a request whose decode peer was torn
down between the handshake and the prefill's final ``send()`` wedged FOREVER
in ``disagg_prefill_inflight_queue`` — ``add_transfer_request`` silently
dropped the last chunk (no transfer destinations), so Success was never set,
and the lost decode ABORT notification meant Failed was never injected.
Registered: CPU CI.
"""
import threading
import unittest
import numpy as np
from sglang.srt.disaggregation.base.conn import KVPoll
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
ROOM = 4242
def _make_manager(status: KVPoll) -> MooncakeKVManager:
m = MooncakeKVManager.__new__(MooncakeKVManager) # skip heavy init
m.disaggregation_mode = DisaggregationMode.PREFILL
m.request_status = {ROOM: status}
m.transfer_infos = {} # no destinations registered for ROOM
m.failure_records = {}
m.failure_lock = threading.Lock()
m.transfer_queues = [] # must never be reached in these tests
return m
class TestAddTransferRequestWithoutDestinations(unittest.TestCase):
def _send_last_chunk(self, manager: MooncakeKVManager) -> None:
manager.add_transfer_request(
ROOM,
kv_indices=np.array([0, 1], dtype=np.int32),
index_slice=slice(0, 2),
is_last_chunk=True,
aux_index=0,
)
def test_non_success_room_is_failed_loudly(self):
# The wedge case: room bootstrapped (WaitingForInput) but the decode
# peer is gone. The chunk must NOT be dropped silently — the room
# must conclude Failed so the scheduler's poll consensus reaps it.
manager = _make_manager(KVPoll.WaitingForInput)
self._send_last_chunk(manager)
self.assertEqual(manager.check_status(ROOM), KVPoll.Failed)
self.assertIn(ROOM, manager.failure_records)
def test_dummy_rank_success_room_is_untouched(self):
# Dummy ranks are marked Success at handshake time and legitimately
# have no transfer_infos: the early return must stay a no-op.
manager = _make_manager(KVPoll.Success)
self._send_last_chunk(manager)
self.assertEqual(manager.check_status(ROOM), KVPoll.Success)
self.assertNotIn(ROOM, manager.failure_records)
def test_already_failed_room_stays_silent(self):
# Pre-existing behavior: chunks for failed rooms are dropped.
manager = _make_manager(KVPoll.Failed)
self._send_last_chunk(manager)
self.assertEqual(manager.check_status(ROOM), KVPoll.Failed)
self.assertNotIn(ROOM, manager.failure_records)
class TestInflightResidencyTimeoutEnv(unittest.TestCase):
def test_default_matches_sibling_disagg_timeouts(self):
from sglang.srt.environ import envs
self.assertEqual(envs.SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT.get(), 300)
self.assertEqual(
envs.SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT.get(),
envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get(),
)
if __name__ == "__main__":
unittest.main()