Fixes a race in the UCX backend between synchronous rkey unpack and endpoint teardown. (#1987)
Fix https://github.com/ai-dynamo/nixl/issues/1986 ## What? Fixes a race in the UCX backend between synchronous rkey unpack and endpoint teardown. After a remote peer disconnects or restarts, the UCX error callback may force-close an endpoint while, or immediately before, loadRemoteMD() calls ucp_ep_rkey_unpack() on the same endpoint. This can cause a native segfault instead of returning a recoverable NIXL error. This PR: - Serializes rkey unpack with both error-driven and normal endpoint close. - Checks endpoint state before calling ucp_ep_rkey_unpack(). - Returns NIXL_ERR_REMOTE_DISCONNECT when unpack is attempted on a failed endpoint. - Preserves the status through the rkey constructor and internalMDHelper(). - Adds focused tests for close-during-unpack and unpack-after-close. ## Why? In vLLM P/D disaggregation, restarting a Prefill Pod can leave the surviving Decode process handling asynchronous UCX endpoint errors while loading metadata for the replacement Prefill. The reported crash occurs in: ucp_ep_rkey_unpack → nixl::ucx::rkey::unpackUcpRkey → nixlUcxEngine::internalMDHelper → nixlAgent::loadRemoteMD One Decode TP worker segfaults while the other workers return NIXL_ERR_REMOTE_DISCONNECT. The crashed worker brings down the entire Decode instance. NIXL previously called ucp_ep_rkey_unpack() without synchronizing against endpoint close or checking whether the endpoint had already failed. ## How? An endpoint-level mutex now protects: - Endpoint state validation followed by rkey unpack. - Force-close from the UCX error callback. - Normal endpoint close during teardown. The endpoint state remains atomic so existing lock-free transfer-path state checks are preserved. Since the rkey constructor cannot return a NIXL status directly, a typed rkey_error carries the failure status to internalMDHelper(). This allows a remote disconnect to remain NIXL_ERR_REMOTE_DISCONNECT instead of being converted to the generic NIXL_ERR_BACKEND. The tests cover both relevant orderings: rkey unpack starts → endpoint close waits endpoint closes first → unpack returns remote disconnect without calling UCX <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved UCX connection error and disconnect handling for more reliable endpoint state transitions. * Prevented metadata creation when a remote endpoint is unavailable or in a failed transmit state. * Ensured failed connections are closed consistently and reported with the appropriate remote-disconnect status. * **Tests** * Added automated coverage for UCX endpoint connection-reset and failure scenarios. * Updated optional UCX test integration for supported builds. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Ilia Yastrebov <iyastrebov@nvidia.com>
C
Chauncey committed
3e3e083c97550a1ddca5653389897cb4ebe5e86e
Parent: 3d764d0
Committed by GitHub <noreply@github.com>
on 7/29/2026, 4:47:05 AM