FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix(pd): guard remote_agents dict with a lock to fix NIXL concurrent-add race (GH-1470) by Ashfaqbs · Pull Request #1505 · ModelTC/LightLLM · GitHub

fix(pd): guard remote_agents dict with a lock to fix NIXL concurrent-add race (GH-1470) - #1505

Merged
hiworldwzj merged 3 commits into
ModelTC:mainfrom
Ashfaqbs:fix/nixl-remote-agents-race
Sep 6, 2026
Merged

fix(pd): guard remote_agents dict with a lock to fix NIXL concurrent-add race (GH-1470)#1505
hiworldwzj merged 3 commits into
ModelTC:mainfrom
Ashfaqbs:fix/nixl-remote-agents-race

Conversation

Copy link
Copy Markdown
Contributor

NixlKVTransporter.remote_agents is a plain dict that several worker threads read and mutate concurrently on the same instance (recv_task_loop, dispatch_task_loop, accept_peer_task_loop, request_page_loop, read_page_to_mems_loop, success_loop, fail_loop are all started as separate daemon threads sharing one transporter).

connect_add_remote_agent() did an unlocked check-then-add:

if remote_agent.agent_name in self.remote_agents:
    return
...
self.remote_agents[remote_agent.agent_name] = remote_agent

and remove_remote_agent() did an unlocked check-then-pop. Every send_* method also does an unlocked if peer_name not in self.remote_agents: connect_add_remote_agent(...).

When a peer transiently disappears (e.g. after NIXL_ERR_REMOTE_DISCONNECT removes it), two threads can both observe it missing and both call nixl_agent.add_remote_agent() for the same peer at the same time. That double-add corrupts UCX endpoint state and aborts the NIXL progress thread with a ud_ep.c PSN assertion, which kills the KV-transfer process and takes down the whole PD router (see #1470 for the full failure timeline and stack trace).

Fix

Add a threading.Lock to NixlKVTransporter and guard the full check-then-add sequence in connect_add_remote_agent() and the check-then-pop in remove_remote_agent() with it, so only one thread can add or remove a given peer at a time. The rest of the call sites' pre-checks (if X not in self.remote_agents) stay as cheap fast-path hints — correctness now comes from the lock inside connect_add_remote_agent, which re-checks membership once it holds the lock.

Fixes #1470

Scope

Kept to the one file / one race described in the issue (nixl_kv_transporter.py). Did not touch the wider unlocked reads in write_blocks_paged etc. since those aren't the failure mode reported in #1470 and would be a separate, larger change.

Verification

  • python -m py_compile on the changed file.
  • No NIXL/UCX hardware available in this environment to reproduce the multi-node race directly; the fix follows directly from the file's own logic (unguarded check-then-add/remove across threads that the class itself spawns) and matches the exact code path called out in the issue.

…add race

connect_add_remote_agent() did an unlocked check-then-add on
self.remote_agents, and remove_remote_agent() did an unlocked pop.
Both are called from several worker threads on the same
NixlKVTransporter instance (recv/dispatch/accept_peer/request-page/
ready-page loops all run as separate threads sharing one transporter).

When a peer briefly looks missing (e.g. after a transient
NIXL_ERR_REMOTE_DISCONNECT removes it), two threads can both see it
absent and both call nixl_agent.add_remote_agent() for the same peer
concurrently, which can corrupt UCX endpoint state and abort the NIXL
progress thread (observed as a ud_ep.c PSN assertion), taking down the
router.

Guard the whole check-then-add and check-then-remove sequences with a
lock so only one thread can add or remove a given peer at a time.

Fixes ModelTC#1470

Copy link
Copy Markdown
Collaborator

Thanks for the PR, @Ashfaqbs! I've been quite busy lately, and I'm sorry I haven't been able to get these changes merged sooner. I appreciate your contribution and your patience!

hiworldwzj merged commit d75e63d into ModelTC:main Sep 6, 2026
1 check passed

Ashfaqbs commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

Thank you for merging the PR and for the update @shihaobai, ! No worries at all about the delay.

Please feel free to reach out in the future if you'd like to collaborate on new features or discuss ideas for the project

This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] NIXL/UCX crashes during concurrent remote-agent reconnect in PD mode

3 participants


Back | FazBrowse Home | New Git URL