Repository navigation
[KVConnector][P2P] Register peers off the scheduler thread - #57139
liranschour wants to merge 9 commits into
Conversation
P2PSession._on_connect called NixlTransport.add_remote_peer inline before replying with ConnectAckMsg. Registration is add_remote_agent plus a prep_xfer_dlist over every block of the peer's region — 19,599 blocks per rank for a 64 GiB CPU tier — so a cold handshake stalled a scheduling iteration on both peers and delayed the lookup queued behind the acknowledgement. Add DataTransport.add_remote_peer_async, defaulting to the existing inline call so transports with cheap registration are unaffected. NixlTransport overrides it with a single-worker executor, mirroring the NIXL connector's _handshake_initiation_executor. The session starts the registration in _on_connect and sends ConnectAckMsg from poll() once the future resolves; a failed registration marks the connection dead instead. Because the peer treats the acknowledgement as permission to send, a peer that may fetch is still a peer we have already registered — the wire protocol is unchanged. Registrations carry a per-peer generation, bumped on every add and remove, so a worker that finishes building handles for a peer that was reaped or re-registered meanwhile releases those handles instead of reviving a stale peer. A pending registration counts as session pending work so the engine keeps ticking until the acknowledgement goes out. Signed-off-by: Liran Schour <lirans@il.ibm.com>
|
I ran this PR on the same cell and case chain I used for #55962, to check whether registration alone covers #55179. Setup: GLM-5.2-FP8, two prefill pods (DP8/TP1, 8xH200 each) and a decode pair,
The busy rows stay in the baseline band with this PR, the same as the registration-only #55962 arm, and drop to 0.6-0.9 s only once the control plane is serviced during the model wait. Idle rows are the same on every tree. On the model in the description: I measured registration directly with trace stamps on both peers in an earlier round, and So I would keep this change (the reap-race generation guard is worth having on its own), but not as the fix for #55179: that needs the per-step gating addressed. |
…rule The NixlTransport threading note says the manager lock must cover every agent entry point. Peer registration moved onto a worker thread, as NixlConnector's handshake executor already does and as vllm-project#57139 proposes for this transport, deliberately runs outside that lock. Say so, so the note stays accurate when that lands. Signed-off-by: Liran Schour <lirans@il.ibm.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ruff's pydocstyle D413 requires a blank line after the last docstring section. The new add_remote_peer_async docstring ended its Returns block directly against the closing quotes, which made the pinned ruff-check pre-commit hook rewrite the file and fail CI. Matches the existing style of every other Returns: block in this file. Signed-off-by: Liran Schour <lirans@il.ibm.com>
|
/ci run |
|
✅ Triggered Buildkite CI #93394 for commit |
|
✅ @liranschour, CI is now available for this PR.
|
CI selector (shadow): 1 test steps (1 jobs) instead of 62 (78 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (1)
Would skip (today's rules run them) (61)
Would add (today's rules do not run them) (0)none AMD mirrors: would skip (53)
AMD mirrors: would add (0)none 6 changed files · base |
|
/ci run |
|
❌ This PR is 14 commits behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #94025 for commit |
This branch has not been deployed
4 participants
Addresses #55179.
Problem
P2PSession._on_connectcalledNixlTransport.add_remote_peerinline before replying withConnectAckMsg. Registration isadd_remote_agentplus aprep_xfer_dlistover every block of the peer's region — 19,599 blocks per rank for the 64 GiB CPU tier in the reported deployment — so a cold handshake ran that work inside a scheduling iteration on both peers, and the lookup queued behind the acknowledgement waited for it.Scope: registration only, not the per-step gating
The issue proposes two independent changes. This PR implements only the first, because the reported numbers point at registration as the dominant term.
Within one scheduler step,
TieringManager.on_schedule_endrunsget_finished_jobs()→_poll_once()(drain + dispatch) beforeserve_external_requests(parent), so receive-and-answer collapse into a single tick. The busy source's critical path is therefore about four ticks:ConnectMsg→ ack,LookupMsg→LookupRespMsg,FetchMsg→write_blocks, and NIXL completion →TransferDoneMsg. The destination was idle in the measurement, so its ticks are effectively free.That gives
13.6 s ≈ 2·T_reg + 4·T_step. At the issue's own stated ~1 s/step, per-step gating accounts for ~4 s, leaving ~9.6 s for two registrations. For gating alone to explain 13.6 s you would need ~3.4 s/step, which contradicts the same description. Per-step gating is real as a mechanism but is the minority term here, so it is deliberately left out of this PR and should be justified by its own measurement.Change
DataTransport.add_remote_peer_async— new method whose default implementation runsadd_remote_peerinline and returns an already-resolved future. Transports with cheap registration are unaffected and need no override.NixlTransportoverrides it with a single-workerThreadPoolExecutor, mirroring the NIXL connector's_handshake_initiation_executor.add_remote_peeris split into_build_peer_handles(the expensive agent wireup plus per-block descriptors, run unlocked — holding a lock across it would stall the scheduler thread for exactly as long as the work being moved off it) and_publish_peer(the atomic install).P2PSessionstarts the registration in_on_connectand sendsConnectAckMsgfrompoll()once the future resolves; a failed registration marks the connection dead so the manager reaps the session.The acknowledgement is the safety barrier
ConnectAckMsgis what clears the peer's_send_readygate, so the peer cannot send aFetchMsguntil it arrives. Gating the ack on registration completion preserves the invariant that any peer permitted to fetch is a peer already registered here, sowrite_blockscan never observe an unregistered peer as a result of this change. The wire protocol is unchanged.Generation guard
Registrations carry a per-peer generation, bumped on every add and every remove. A worker that finishes building handles under a stale generation releases them instead of publishing. This closes the reap race:
_reap_dead_sessionscallssession.close()and then_data.remove_remote_peer(pid), and without the guard a registration landing after that would leave a registered peer with no session plus leaked NIXL handles. A pending registration also counts as session pending work, so the engine keeps ticking until the acknowledgement goes out.No
NIXL_THREAD_SYNC_STRICTThe issue suggests opening the agent with
NIXL_THREAD_SYNC_STRICTor_RW. That is not reachable throughnixl_agent_config: the Python wrapper hard-codesagent_config.syncMode, selectingSTRICTonly whenenable_listenis true, which would additionally start a listener thread and bind a port that this tier does not use.It is also unnecessary.
NixlConnectoralready runsadd_remote_agent,get_xfer_descsandprep_xfer_dliston its single-worker handshake executor concurrently with transfer submission on the main thread, with the defaultsyncMode=NONE. This PR adopts the threading model vLLM already ships rather than introducing a new one, which also removes the need to measure STRICT's cost on the transfer path.Why this does not duplicate an existing PR
Draft #55962 covers this ground as part of a larger change that also services the control plane during model steps, and does so by adding a
poll_pending_workhook toKVConnectorBase_V1plus changes toEngineCoreand both executors.This PR is materially different in approach and scope:
KVConnectorBase_V1. The audit behind this is that onlyServerRole.serve_external_requestsand its four helpers need theParentManagerhandle — the whole client role, the entire message-dispatch table, and all NIXL submit/poll/cancel work are already parent-free, andon_lookupalready defers resolution to the existingserve_external_requestswindow. No fresh scheduler-thread window is required.vllm/v1/engine/core.py,vllm/v1/executor/{abstract,uniproc_executor,multiproc_executor}.py,kv_connector/v1/base.py,multi_connector.py, oroffloading_connector.py.Landing this first also makes #55962's remaining half measurable on its own terms: with registration off the critical path, whatever delay is left on a busy source is the per-step gating component.
Tests
Run from
venv/(this checkout's working environment):New tests, one behavior each:
ConnectAckMsgis withheld while registration is in flight, and sent once it resolves.ConnectMsgmid-registration is a protocol error.remove_remote_peeris discarded and its handles released.The session test fake gained a
defer_registrationmode so the acknowledgement gating is exercised deterministically rather than relying on thread timing.Model evaluation
No eval was run, because this change cannot affect model output: it reorders when a control-plane acknowledgement is sent and moves NIXL registration to another thread. Which blocks are transferred, and their contents, are untouched, and the wire protocol is unchanged.
Performance validation is still outstanding. The change targets a timing property that unit tests cannot measure, and no multi-pod deploy has been run yet. The measurement that validates it is TTFT on a busy source with a cold peer versus a warm peer: if cold-peer TTFT drops toward warm-peer, registration was the dominant cost. This PR stays in draft until that result is in.
AI assistance
AI assistance (Claude Code) was used for the analysis, implementation and tests in this PR. The submitting human is reviewing every changed line and is running the hardware validation described above before marking it ready for review.
🤖 Generated with Claude Code