SIGN IN SIGN UP

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