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>
This commit is contained in:
@@ -1596,9 +1596,30 @@ class MooncakeKVManager(CommonKVManager):
|
||||
return
|
||||
|
||||
if bootstrap_room not in self.transfer_infos:
|
||||
# This means that the current rank is a dummy rank for this request,
|
||||
# and it has already been marked as success, so there is no need to
|
||||
# add further chunks into the transfer queue.
|
||||
if self.check_status(bootstrap_room) == KVPoll.Success:
|
||||
# Dummy rank for this request: it was already marked Success
|
||||
# at handshake time, so there is nothing left to transfer.
|
||||
return
|
||||
# Non-Success room with no transfer destinations is an anomaly
|
||||
# (decode peer torn down between handshake and this send, abort
|
||||
# notification lost). Silently dropping the chunk used to wedge
|
||||
# the request in WaitingForInput FOREVER — Success can never be
|
||||
# set without the last chunk, and nothing else fails the room.
|
||||
# Conclude it loudly instead; the scheduler's poll consensus
|
||||
# (MIN-reduce, Failed=0) reaps it on every rank.
|
||||
logger.warning(
|
||||
"KV chunk for bootstrap_room=%s has no transfer destinations "
|
||||
"(status=%s); failing the request instead of dropping the "
|
||||
"chunk silently.",
|
||||
bootstrap_room,
|
||||
self.check_status(bootstrap_room),
|
||||
)
|
||||
self.record_failure(
|
||||
bootstrap_room,
|
||||
f"Request {bootstrap_room} lost its transfer destinations "
|
||||
"before the KV chunk was submitted (decode peer gone?)",
|
||||
)
|
||||
self.update_status(bootstrap_room, KVPoll.Failed)
|
||||
return
|
||||
|
||||
# NOTE(shangming): sharding according to the dst_infos to make sure
|
||||
|
||||
@@ -972,6 +972,11 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
req.output_ids.append(next_token_id)
|
||||
self.tree_cache.cache_unfinished_req(req) # update the tree and lock
|
||||
self.disagg_prefill_inflight_queue.append(req)
|
||||
# Residency clock for the inflight liveness bound (every rank
|
||||
# stamps at the same logical point in the same tick, so the
|
||||
# timeout decision below stays effectively rank-uniform; the
|
||||
# MIN-reduced poll consensus absorbs any microsecond skew).
|
||||
req.disagg_inflight_enter_time = time.perf_counter()
|
||||
if spec_cpu_row is not None:
|
||||
k = spec_cpu_row[i]
|
||||
req.output_topk_p = spec_topk_p_cpu[k]
|
||||
@@ -1102,6 +1107,34 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
continue
|
||||
|
||||
if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]:
|
||||
# Liveness bound: an inflight entry whose transfer never
|
||||
# concludes (decode peer torn down with its abort
|
||||
# notification lost, last chunk silently dropped, decode
|
||||
# crash) would otherwise be re-queued FOREVER — it pins the
|
||||
# request's KV pages and rides every poll collective. On
|
||||
# expiry, conclude the sender locally and keep the entry
|
||||
# undone: Failed(0) wins the CP/TP MIN-reduce, so every rank
|
||||
# reaps it together through the Failed branch on the NEXT
|
||||
# poll — one removal path, rank-uniform by construction.
|
||||
enter_time = getattr(req, "disagg_inflight_enter_time", None)
|
||||
if (
|
||||
enter_time is not None
|
||||
and time.perf_counter() - enter_time
|
||||
>= envs.SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT.get()
|
||||
and hasattr(req.disagg_kv_sender, "abort")
|
||||
):
|
||||
logger.warning(
|
||||
"Prefill inflight transfer timed out; failing the "
|
||||
"request. rid=%s bootstrap_room=%s poll=%s "
|
||||
"elapsed=%.1fs (raise "
|
||||
"SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT if transfers "
|
||||
"can legitimately take longer)",
|
||||
req.rid,
|
||||
req.bootstrap_room,
|
||||
poll,
|
||||
time.perf_counter() - enter_time,
|
||||
)
|
||||
req.disagg_kv_sender.abort()
|
||||
undone_reqs.append(req)
|
||||
elif poll == KVPoll.Success: # transfer done
|
||||
release_kv_cache(req, self.tree_cache) # unlock the tree
|
||||
|
||||
@@ -296,6 +296,10 @@ class Envs:
|
||||
SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
|
||||
SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2)
|
||||
SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
|
||||
# Liveness bound for entries in the prefill inflight queue (transfer
|
||||
# submitted, waiting to conclude). 300 matches the bootstrap/waiting
|
||||
# timeouts; transfers legitimately take seconds, so do not go below ~60.
|
||||
SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT = EnvInt(300)
|
||||
SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX")
|
||||
SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False)
|
||||
SGLANG_DISAGGREGATION_TRANSFER_STATS = EnvBool(False)
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user