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

[sufubao]

Summary

In prefill/decode disaggregation with --pd_trans_mode nixl, a decode KV-transfer process can abort inside the UCX progress thread after a transient NIXL_ERR_REMOTE_DISCONNECT.

The failure is preceded by two LightLLM threads concurrently adding the same remote NIXL agent. UCX then hits an internal UD endpoint PSN assertion. LightLLM detects the dead transfer process later and kills the router, taking down the entire fixed PD deployment.

Environment

• LightLLM commit: 217a3208 (includes upstream main at 1ca78982)
• Transfer backend: nixl / UCX
• Topology: 2 nodes, 16 GPUs
  • 3 prefill workers, each TP4
  • 1 decode worker, TP4
• Model: Qwen3.5-27B
• Failure observed under concurrent request load

The exact NIXL/UCX package versions still need to be captured from the failed worker image.

Failure timeline

At 08:07:50, decode device 2 fails to notify its prefill peer:

genNotif: backend 'UCX' returned error status NIXL_ERR_REMOTE_DISCONNECT
nixl_cu13._bindings.nixlRemoteDisconnectError: NIXL_ERR_REMOTE_DISCONNECT

The exception path removes the remote agent. Two worker threads then observe that the same peer is absent and both add it:

WARNING prefill_agent_name <prefill-agent>_2 not exist
WARNING prefill_agent_name <prefill-agent>_2 not exist
INFO Added remote agent <prefill-agent>_2 with mem desc <...0x...b3b0>
INFO Added remote agent <prefill-agent>_2 with mem desc <...0x...3330>

About two seconds later, the NIXL UCX progress thread aborts:

ud_ep.c:732 Assertion 'UCT_UD_PSN_COMPARE(ep->tx.acked_psn, <, ep->tx.psn)' failed:
acked_psn=21 must be smaller than current_psn=3

libucs.so        ucs_handle_error / ucs_fatal_error
libuct_ib.so     uct_ud_ep_process_rx
libuct_ib_mlx5.so
libucp.so        ucp_worker_progress
libplugin_UCX.so nixlUcxWorker::progressLoop
libplugin_UCX.so nixlUcxSharedThread::run

The parent health loop reports the secondary failure around 15 seconds later:

kv trans process for device: 2 dead
device_id 2 kv process is unhealth
Critical LightLLM submodule exited unexpectedly ... exitcode=-9

The exitcode=-9 is not the initiating failure. BaseKVMoveManager.check_trans_process_loop explicitly sends SIGKILL after detecting the dead child.

GPU KV usage was approximately 2% at the time, and no CUDA OOM or NCCL error was present.

Relevant code path

lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py:

• remote_agents is shared by several threads without synchronization.
• connect_add_remote_agent performs an unlocked check-then-add.
• remove_remote_agent invalidates NIXL metadata, disconnects the peer, and releases descriptor handles without coordination with sends, notification polling, or transfers.
• Methods such as send_write_ready_task_to_prefill_node perform an unlocked ensure-agent followed by send_notif.

lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py:

• One NixlKVTransporter is shared by seven daemon threads.
• request_page_loop calls send_write_ready_task_to_prefill_node.
• Its exception handler separately calls remove_remote_agent, leaving a race window between the failed send and cleanup.
• Notification and failure-handling threads can try to reconnect/send to the same peer at the same time.

The prefill transfer process contains similar caller-side removal paths.

Root-cause assessment

Confirmed:

• The first fatal native failure is a UCX UD endpoint assertion, not OOM or NCCL.
• The same remote NIXL agent is added twice immediately before the assertion.
• LightLLM invokes a shared NIXL agent concurrently and does not synchronize remote-agent lifecycle operations.
• The prefill worker involved in the failed request remained alive and continued processing, so it did not crash first.

Strong inference:

A transient remote disconnect triggers concurrent remove/re-add operations for the same agent. The duplicate metadata load or endpoint recreation races with UCX progress, leaving the UD endpoint with a reset/current PSN behind its acknowledged PSN.

Proposed fix

1. Add a transporter-owned threading.RLock and serialize all application-thread calls into the shared NIXL agent:

   • metadata and descriptor access
   • notification polling and sending
   • add/remove remote agent
   • transfer creation/post/status/release
   • shutdown

2. Make each compound operation atomic. In particular:

check peer -> add metadata -> prepare remote handles -> publish peer -> send/post

3. Move peer cleanup into the transporter. A failed send and removal of that exact peer generation must happen under the same lock. Callers should only fail the request and must not independently remove the peer.

4. Do not retry the same notification inside the first fix. Fail the affected request and let a later request perform one serialized lazy reconnect.

5. For robust reconnects, track per-peer state and generation:

   • READY / BROKEN / RECONNECTING
   • a single reconnect owner
   • outstanding transfer handles
   • stale-generation errors must not remove a newer peer

6. When the deployed NIXL version supports it, explicitly construct the NIXL agent with strict thread synchronization. Application-level locking is still required because internal synchronization does not make the LightLLM check/add/prepare/publish sequence atomic.

Regression tests

A deterministic fake-NIXL test can use barriers to force two threads through the previous check-then-add race and verify:

• native add_remote_agent is called once;
• send failure and peer removal cannot interleave with another add/send;
• the next request creates exactly one new peer generation;
• an old-generation exception cannot remove a new generation;
• shutdown does not release descriptor handles while an operation owns them.

A two-node fault-injection test should force one notification to return NIXL_ERR_REMOTE_DISCONNECT and verify that there is no duplicate agent addition, no UCX assertion, and the service remains alive for later requests.

[Ashfaqbs]

Opened a fix: #1505 — guards remote_agents with a lock in connect_add_remote_agent/remove_remote_agent to close the concurrent-add race described above.