Repository navigation
Conversation
Co-authored-by: OpenAI Codex Signed-off-by: levius <2114377220@qq.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
nilig
left a comment
nilig
left a comment
There was a problem hiding this comment.
The published-name check leaves a window when the synchronous add_remote_peer runs alongside an async registration for the same peer. If the older registration fails after the newer one has called add_remote_agent() but before its _publish_peer, the name is not published yet, so the rollback removes the agent the newer registration then publishes. The transport ends up with a published dlist for an agent that no longer exists. Reproduced at 18349008c with the real NixlTransport constructor and a fake external agent:
B entered prep while A was blocked: True
published={'peer:1': 'nixl-peer-name'} registered=set()
AssertionError: published agent was removed
The P2P session only calls add_remote_peer_async, so this is an API-level race, not a path production hits today. Wonder if we could also route the synchronous add_remote_peer through _reg_executor: async registrations are already serialized there, so the published-name check and the rollback would both finish before the next build starts. The synchronous call would then raise when superseded, as the async one does. Probe below in case it is useful as a regression test.
test_probe_window.py
"""Older failed registration A rolls back while newer sync registration B
has added the agent but not yet published it."""
import threading
from unittest.mock import MagicMock, patch
import numpy as np
from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport
def test_rollback_while_newer_registration_is_building():
agent = MagicMock()
with (
patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", return_value=agent),
patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig"),
):
transport = NixlTransport("test:1", memoryview(np.zeros((8, 1024), dtype=np.uint8)))
registered: set[str] = set()
def add_remote_agent(metadata):
registered.add("nixl-peer-name")
return "nixl-peer-name"
a_in_prep, release_a = threading.Event(), threading.Event()
b_in_prep, release_b = threading.Event(), threading.Event()
calls = 0
def prep_xfer_dlist(*args, **kwargs):
nonlocal calls
calls += 1
if calls == 1:
a_in_prep.set()
assert release_a.wait(10)
raise RuntimeError("A failed")
b_in_prep.set()
assert release_b.wait(10)
return MagicMock(name="remote_dlist")
agent.add_remote_agent.side_effect = add_remote_agent
agent.remove_remote_agent.side_effect = registered.discard
agent.prep_xfer_dlist.side_effect = prep_xfer_dlist
errors = []
def register_b():
try:
transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
except BaseException as exc:
errors.append(exc)
try:
old = transport.add_remote_peer_async("peer:1", b"meta", 0x1000, 8, 1024)
assert a_in_prep.wait(5)
transport.remove_remote_peer("peer:1")
b = threading.Thread(target=register_b)
b.start()
concurrent = b_in_prep.wait(1.0)
print(f"B entered prep while A was blocked: {concurrent}")
release_a.set()
try:
old.result(timeout=5)
except RuntimeError as exc:
print(f"A: {exc}")
assert b_in_prep.wait(5)
release_b.set()
b.join(timeout=5)
published = dict(transport._peer_nixl_names)
print(f"published={published} registered={registered} b_errors={errors}")
assert errors == []
assert registered == {"nixl-peer-name"}, "published agent was removed"
assert transport.write_blocks("peer:1", [0], [1]) is not None
finally:
release_a.set()
release_b.set()
transport.close()Drop it into tests/v1/kv_offload/tiering/p2p/ and run pytest --confcutdir=tests/v1/kv_offload/tiering/p2p tests/v1/kv_offload/tiering/p2p/test_probe_window.py -s.
liranschour
left a comment
liranschour
left a comment
There was a problem hiding this comment.
Thanks, this closes gap 4 cleanly. I checked the rollback against the registration flow on main and ran test_data_transport.py locally (64 passed).
The production path looks correct. The session only registers through add_remote_peer_async, so every build runs on the single _reg_executor worker. Between the published check and the release, no other build can publish the same NIXL name. A concurrent remove_remote_peer can at worst cause a second remove_remote_agent, which _release_remote_handles already logs and survives. Re-raising the original exception while _release_remote_handles logs cleanup failures is the right split.
Two inline comments: the sync-path race nilig raised, and a small nit.
Not for this PR, just noting it: re-registering a peer that's already published, without removing it first, overwrites _remote_dlists[peer_id] in _publish_peer and never releases the old dlist handle. That's on main already, so it would be a separate fix.
| gen = self._begin_registration(peer_id) | ||
| nixl_name, remote_dlist = self._build_peer_handles( | ||
| agent_metadata, base_addr, num_blocks, block_len | ||
| peer_id, agent_metadata, base_addr, num_blocks, block_len |
There was a problem hiding this comment.
Agree with nilig's race. I think it's worth fixing here. The rollback's ownership check is only sound because builds run one at a time on _reg_executor. This sync path builds on the caller's thread, so it can overlap an async build. An older failed registration can then remove an agent that this call is about to publish. Routing it through the executor removes the gap:
self.add_remote_peer_async(
peer_id, agent_metadata, base_addr, num_blocks, block_len
).result()Nothing calls this from the registration worker, so it can't deadlock. One behaviour change: a superseded sync registration now raises instead of returning silently. Today it ignores _publish_peer's False, which already looks like a bug to me.
Could you add a regression test as well? An older registration fails while a newer one for the same peer has added the agent but not yet published it, and the newer one should stay registered and usable for write_blocks.
There was a problem hiding this comment.
@liranschour Fixed the sync race, added regression coverage and the serialization comment in 6ed41cc. Full P2P suite: 321 passed; pre-commit passed. Could you please re-review?
| remote_dlist = self._agent.prep_xfer_dlist(nixl_name, xfer_dlist) | ||
| except Exception: | ||
| with self._peer_lock: | ||
| published = nixl_name in self._peer_nixl_names.values() |
There was a problem hiding this comment.
Nit: could you add a short comment here noting that this check relies on builds being serialized on _reg_executor? Otherwise adding workers, or a path that skips the executor, would quietly break the ownership check.
Co-authored-by: OpenAI Codex Signed-off-by: levius <2114377220@qq.com>
|
/ci run |
|
✅ @Levius-Fubuki, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #97288 for commit ℹ️ This PR is 38 commits behind upstream |
…ject#63628) Signed-off-by: levius <2114377220@qq.com> Co-authored-by: vllm-agent <claw@inferact.ai> Signed-off-by: Christopher Mackenzie <christopher.mackenzie.software@gmail.com>
…ject#63628) Signed-off-by: levius <2114377220@qq.com> Co-authored-by: vllm-agent <claw@inferact.ai> Signed-off-by: Christopher Mackenzie <christopher.mackenzie.software@gmail.com>
Overview
Roll back a newly registered NIXL agent when descriptor preparation fails in the P2P KV-offload transport. Related to #63602, gap 4 only.
AI assistance (OpenAI Codex) was used for investigation, implementation, tests, and this description. The human submitter confirmed completion of manual review and authorized submission. The reviewed contribution carries their DCO sign-off. Original-submission results below are the recorded server runs; review follow-up results are the recorded local runs. No additional human test execution is asserted.
Claims
get_xfer_descsorprep_xfer_dlistrolls back the unpublished remote registration.Validation
The table and native validation below describe the original submission,
18349008c; follow-up results are recorded separately below.Base:
a7407ab44557802ca32a0ba2c2707362d7d21563.Tests construct the real
NixlTransport; only the external NIXL boundary is mocked. Native validation uses the production transport with two real NIXL agents, host buffers, UCX TCP/shared memory/CUDA-copy, injected descriptor failures, native metadata queries, retries, and exact byte comparisons. Both descriptor stages and an idempotent repeat are checked.This source-level Python validation used Python 3.12.3 and the server's existing torch 2.12.1+cu130. It did not build main's torch 2.14 native extensions. The command excludes the heavyweight root
tests/conftest.py, and uses the P2P suite's real conftest. No LLM/OpenAI serving run, output-quality evaluation, or RDMA/multi-node fault injection was performed; those are not implied by the transport results.Review follow-up — sync registration
Commit
6ed41ccf1broutesadd_remote_peer()throughadd_remote_peer_async(...).result(), so synchronous and asynchronous builds share the single registration worker. Superseded synchronous registrations now raise, and the rollback check documents its serialization requirement.Validation of this follow-up used macOS/Python 3.12 with the real
NixlTransportconstructor and a mocked external NIXL boundary:18349008cThe before-fix failures were
assert set() == {'peer'}(the newer registration's agent was removed) andFailed: DID NOT RAISE RuntimeError(superseded sync registration returned silently).Both baseline and changed runs emit the same missing generated
vllm._versionwarning. No new native-NIXL/RDMA or serving validation is claimed for this follow-up. The already-existing dlist replacement leak remains outside this PR.Details
The registration worker can register native metadata and then fail before
_publish_peerrecords cleanup ownership._build_peer_handlesnow uses the existing independent teardown helper on that exception path, while preserving the exception and avoiding removal of an already-published native name. All registration builds now use the existing single-worker executor, including the synchronous entry point; the generation/publish flow is unchanged. No steady-state inference work, synchronization, public option, or model arithmetic is added.Duplicate checks covered the issue discussion, open PRs referencing #63602, and P2P/NIXL rollback keywords. #59524 fixes incomplete registration in the separate KV connector
nixl/base_worker.py; this change fixes the P2P secondary-tier transport. #61667 covers consumer-slot quarantine for unconfirmed writes. Contributortejachavaliexpressed interest in the broader issue; this PR does not claim assignment or close the full issue.The existing Buildkite
v1-kv-offloadstep selects the data-transport tests viapytest -v -s v1/kv_offload, and its source dependencies includevllm/v1/andtests/v1/kv_offload.Pull Request Checklist
/pr-checklistskill. (Mandatory for agents, optional for humans).Need help on this PR? Tag
@codesmith-botwith what you need.Need help on this PR? Tag
@codesmith-botwith what you need.