[PD] Move the prefill parallel-info fetch off the decode scheduler thread

[zkyue]

Motivation

CommonKVManager.try_ensure_parallel_info() is documented as a "single
non-blocking attempt", but it performs a synchronous
requests.get(url, timeout=5). It is called from the decode scheduler thread:

SchedulerDisaggregationDecodeMixin.process_decode_queue()
  -> DecodePreallocQueue.pop_preallocated()
    -> _resolve_pending_reqs()
      -> _ensure_prefill_info()
        -> CommonKVManager.try_ensure_parallel_info()   # requests.get(..., timeout=5)

So a bootstrap server that stops answering rather than refusing stalls the
entire decode loop — every running request on that instance — for the full
timeout.

The 1 s retry interval in _ensure_prefill_info() does not bound this.
now = time.monotonic() is captured once before the loop and
_ensure_last_attempt_time[bootstrap_addr] is stamped before the call, so a
5 s timeout fully absorbs the 1 s interval and attempts land essentially back to
back. With _max_ensure_retries = 15 that is up to ~75 s of near-continuous
stall for one unreachable addr, and addr_to_reqs is iterated sequentially, so
K such addrs cost K × 5 s within a single scheduling cycle.

This is not a hypothetical state. The decode-side heartbeat checker evicts an
addr from the cache after max_failures consecutive /health failures:

def _handle_node_failure(self, failed_bootstrap_addr: str):
    with self.connection_lock:
        ...
        self.prefill_info_table.pop(failed_bootstrap_addr, None)

Requests still routed to that addr then miss the cache, land in pending_reqs,
and the scheduler thread itself starts dialing the unreachable prefill server.

/health_generate passes as long as last_receive_tstamp advances within
HEALTH_CHECK_TIMEOUT (default 20 s), so a single 5 s stall only inflates ITL —
but a sustained stall exceeds the window, the decode server reports unhealthy,
and a router in front of it will pull it out of rotation while its requests are
still in flight. #20252 reports exactly that shape at scale: a prefill restart
leaves decode servers "stuck continuously retrying to connect to the downed
prefill server", their health checks time out, sglang-router removes them, and
the remaining decoders take the concentrated traffic.

#20785 already set out to fix this — its summary says add() becomes "a pure
cache lookup — no network calls, no blocking", and its background notes the
old ensure_parallel_info() "blocks up to 30s (5 retries × 1s sleep + 5s HTTP
timeout)". It did remove the retry loop and the time.sleep(). But the
remaining single 5 s GET only moved from add() to the pending_reqs
resolution path; both run on the scheduler thread.

Modifications

Hand the HTTP fetch to a small ThreadPoolExecutor, mirroring the existing
_prefill_recompute_executor idiom in the same class:

• try_ensure_parallel_info() is now a cache lookup plus at most one fetch
  submission per addr, and returns False while a fetch is outstanding. Its
  callers already treat False as "not available yet, retry on a later
  scheduling cycle", so no caller contract changes.
• An in-flight set keeps repeated scheduling cycles from queueing duplicate
  fetches for the same addr. It is mutated under connection_lock, and the cache
  is re-checked inside that lock, so a fetch that publishes between the unlocked
  check and the locked section does not cause a redundant fetch. The
  parallel_info_fetch_in_flight() predicate reads the set without the lock; a
  stale answer only costs a skipped cycle or one extra call that the locked
  re-check rejects.
• The old body is split into _fetch_parallel_info() (the blocking GET, HTTP
  error handling unchanged — log and return no info), _check_parallel_info()
  (the page-size / kv-cache-dtype / DCP validations, unchanged), and
  _publish_parallel_info() (the executor-thread body).
• Config mismatches stay fatal. They currently raise on the scheduler thread and
  take the engine down, which is deliberate; an executor would swallow them, so
  the fetch thread logs the error and records the exception, and
  raise_parallel_info_error() re-raises it on the scheduler thread — both from
  try_ensure_parallel_info() and once per scheduling cycle in
  _ensure_prefill_info(), so a mismatch found on the last allowed attempt is
  not masked by the abort path taking over on that cycle.
• _ensure_prefill_info() gets four adjustments so its retry bookkeeping keeps
  the meaning it has today, now that a fetch can outlive a scheduling cycle:
  • it checks the cache before the retry-interval gate, so a result that lands
    asynchronously is consumed on the next cycle instead of waiting out the
    interval (otherwise the first request to a new addr would pay 1 s);
  • a cycle that finds a fetch still in flight does not spend a retry attempt.
    Without this, _max_ensure_retries would silently become "15 scheduling
    cycles" instead of "15 fetch attempts", cutting the tolerance for a slow
    bootstrap server by ~5× and making the abort message ("after N attempts")
    untrue;
  • the abort is decided before spending an attempt rather than after. When
    the call was blocking, the outcome of the last attempt was known in the same
    cycle; with an asynchronous fetch, aborting right after submitting the last
    attempt would discard its result unseen, so a fetch that succeeds on attempt
    15 would still lose its requests. Checking the budget first means the abort
    can only happen on a cycle where no fetch is outstanding, i.e. after all 15
    have finished. The abort therefore waits out the last attempt — up to one
    timeout=5 later than before, and one retry interval later if the attempt
    finishes sooner than that — and the message is unchanged;
  • the abort branch re-checks the cache before discarding anything. A fetch can
    publish between the check at the top of the cycle and the abort below it, and
    those requests are fine. On main this could not happen — prefill_info_table
    was written only by this thread — so the window is one this change opens.

Behaviour

• Steady state is unchanged: parallel info is fetched once per (decode process,
  prefill addr) — a failed fetch is retried as before — and every later call is
  the same dict lookup as today.
• First contact with a new addr no longer blocks the loop; the request is
  admitted on a later scheduling cycle instead of after a round trip.
• The retry budget is unchanged: 15 fetch attempts paced at ≥1 s apart, and all
  of them complete before the same abort message is logged. For an unresponsive
  server the wall-clock time to abort is therefore comparable to today's; the
  difference is that the scheduler is not blocked while those attempts run.
• ThreadPoolExecutor worker threads are joined at interpreter exit, so
  outstanding fetches can delay shutdown until they finish. There is at most one
  fetch per bootstrap addr, and requests' timeout=5 bounds connect and read
  individually rather than total wall clock. This matches the existing
  _prefill_recompute_executor in this class.
• Not addressed here: a fetch in flight when the heartbeat checker evicts the
  same addr can still republish it afterwards. That race predates this change
  (the inline fetch raced _handle_node_failure on the heartbeat thread in
  exactly the same window), so it is left out to keep this change focused.

Tests

test/registered/unit/disaggregation/test_parallel_info_nonblocking.py (new,
CPU-only, no sleeps — a fake fetch blocks on one threading.Event and signals
that it has been entered with another, a single-worker executor barrier waits for
a submission to drain, and one custom lock context covers the publish-during-lock
case):

• a cache hit returns True without touching the network;
• try_ensure_parallel_info() returns while the GET is still outstanding, and
  the info is published once the fetch completes — this is the regression guard;
• repeated calls while a fetch is in flight issue exactly one GET;
• a result published exactly at the locked section is not fetched again;
• a failed fetch (non-200) publishes nothing, clears the in-flight marker, and is
  retried by the next call;
• a page-size mismatch found on the fetch thread is re-raised as RuntimeError
  on the caller thread;
• _ensure_prefill_info() admits an addr as soon as the result lands, even
  inside the retry interval, while the interval still paces fetch attempts;
• cycles spent waiting on an outstanding fetch do not consume the retry budget;
• a fetch that succeeds on the last allowed attempt still admits its requests,
  instead of being aborted in the cycle that submitted it;
• a result that publishes between the cache check at the top of a cycle and the
  abort below it admits its requests instead of aborting them;
• a config mismatch found on the last allowed attempt still raises, rather than
  being reported as retry exhaustion.

I checked that these fail without the code they cover: making the submission
blocking again (.result()) fails the two non-blocking cases and stretches the
file from 12 s to 83 s as the fake fetches serialize; dropping the in-flight
skip in _ensure_prefill_info() fails the retry-budget case; dropping the
re-check under the lock fails the duplicate-fetch case; deciding the abort after
submitting rather than before fails the last-attempt case; dropping the cache
re-check in the abort branch fails the publish-before-abort case; dropping the
per-cycle raise_parallel_info_error() fails the final-attempt-mismatch case.

test/registered/unit/disaggregation/ passes in full (163 tests).
test_decode_queue_cleanup.py needed two lines and one value: it mocks the KV
manager with a bare MagicMock, so the new predicates returned truthy mocks,
and it pre-seeds _ensure_retry_count at 0 with _max_ensure_retries = 1 to
reach the abort path in a single call. Since the budget is now checked before an
attempt is spent, that seed is 1; what the test covers (a request whose
kv_receiver was already cleared must not crash on .abort()) is unchanged.

Not run by me: a multi-node PD failure smoke test (kill a prefill instance,
confirm the decode instance stays healthy and keeps serving). I do not have a
multi-node rig available; happy to add one if a maintainer prefers it in-tree.

Accuracy Tests

Not applicable: this changes which thread an HTTP GET runs on. No model forward,
kernel, or numerics change.

Speed Tests and Profiling

No steady-state performance change to measure: after the first fetch per
(decode process, prefill addr) the call is a dict lookup on both main and this
branch, and nothing is added to the per-token path. The change is to the failure
path, where the scheduler loop no longer blocks for 5 s per fetch attempt; the
quantities involved (5 s timeout, 1 s interval, 15 attempts,
HEALTH_CHECK_TIMEOUT = 20 s) are all in-tree constants cited above.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations — no user-facing docs cover this internal path.
• Provide accuracy and speed benchmark results — not applicable, see above.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #32809376569
Latest PR Test (Extra): ❌ Run #32809376468
Latest PR Test (AMD ROCm 7.2): ❌ Run #32809376582

[zkyue]

Re-verified against current main (f2ef826f0c): the branch merges cleanly and the underlying issue is still present — CommonKVManager.try_ensure_parallel_info() still performs a synchronous requests.get(url, timeout=5) on the decode scheduler thread. This PR moves that fetch onto a background executor so an unresponsive bootstrap server cannot stall the decode loop.

Relationship to the recently merged #35071: complementary, not overlapping — #35071 parallelizes the per-DP-rank bootstrap queries in the decode prealloc path, while the blocking parallel-info fetch addressed here is a different call site and remains synchronous on main.

Local verification (head 4e91af3af, merge-tested against f2ef826f0c, re-run today):

• test/registered/unit/disaggregation/test_parallel_info_nonblocking.py — 12 passed
• test/registered/unit/disaggregation/test_decode_queue_cleanup.py — 8 passed

This is in the PD-disaggregation area. Could a maintainer add the run-ci label when convenient?