Fixes a race in the UCX backend between synchronous rkey unpack and endpoint teardown. - #1987
Conversation
|
👋 Hi chaunceyjiang! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesUCX failure handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/plugins/ucx/ucx_utils.cpp (1)
157-170: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winTransition state before releasing a closed endpoint.
Lines 160 and 170 clear
ephbut leavestateasCONNECTED. A subsequently unblockedunpackRkey()therefore passes its state check and calls UCX with a null endpoint. Set a terminal state (DISCONNECTEDon successful close; an appropriate failure state on close failure) before clearingeph.Proposed fix
if (request == nullptr) { + setState(nixl::ucx::ep_state_t::DISCONNECTED); eph = nullptr; return NIXL_SUCCESS; } if (UCS_PTR_IS_ERR(request)) { + setState(nixl::ucx::ep_state_t::FAILED); eph = nullptr; return nixl::ucx::ucsToNixlStatus(UCS_PTR_STATUS(request)); } ucp_request_free(request); + setState(nixl::ucx::ep_state_t::DISCONNECTED); eph = nullptr; return NIXL_SUCCESS;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/plugins/ucx/ucx_utils.cpp` around lines 157 - 170, Update the CONNECTED branch of the endpoint close logic to transition state before clearing eph: set state to DISCONNECTED on successful or null-request close, and set the appropriate failure state when UCS_PTR_IS_ERR(request) is true, then clear eph and return. Ensure subsequent unpackRkey() state checks cannot use a null endpoint.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/plugins/ucx/ucx_utils.h`:
- Line 64: Rename the private member state to state_ in the relevant class, and
update every reference to this member accordingly while preserving its atomic
type and initialization.
- Around line 46-59: Add Doxygen-style block comments for the public
nixlUcxEpOps type and unpackRkey() function, documenting the callback and
context lifetime requirements and the status behavior returned by the
operations. Cover both the endpoint-operation callbacks and the associated
unpackRkey API without changing their implementation.
In `@test/gtest/unit/ucx/meson.build`:
- Line 1: The copyright headers in test/gtest/unit/ucx/meson.build lines 1-1 and
test/gtest/unit/ucx/ucx_ep_rkey_test.cpp lines 1-4 are missing the required “All
rights reserved.” wording. Update both SPDX header blocks so the NVIDIA
copyright line ends with “NVIDIA CORPORATION & AFFILIATES. All rights reserved.”
In `@test/gtest/unit/ucx/ucx_ep_rkey_test.cpp`:
- Around line 21-147: Run clang-format on the affected test code, including
CoordinatedEpOps::ops, the callback lambdas, the initializer, and asynchronous
test expressions. Preserve the test logic while applying the repository’s
standard formatting throughout the commented range.
---
Outside diff comments:
In `@src/plugins/ucx/ucx_utils.cpp`:
- Around line 157-170: Update the CONNECTED branch of the endpoint close logic
to transition state before clearing eph: set state to DISCONNECTED on successful
or null-request close, and set the appropriate failure state when
UCS_PTR_IS_ERR(request) is true, then clear eph and return. Ensure subsequent
unpackRkey() state checks cannot use a null endpoint.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 3c6cfdd2-3e98-4e0b-af90-a9a928d538b2
📒 Files selected for processing (8)
src/plugins/ucx/rkey.cppsrc/plugins/ucx/rkey.hsrc/plugins/ucx/ucx_backend.cppsrc/plugins/ucx/ucx_utils.cppsrc/plugins/ucx/ucx_utils.htest/gtest/unit/meson.buildtest/gtest/unit/ucx/meson.buildtest/gtest/unit/ucx/ucx_ep_rkey_test.cpp
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/plugins/ucx/ucx_utils.cpp (1)
89-117: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftSerialize UCX endpoint lifetime across close, error callbacks, and rkey unpack.
checkTxState()is only an atomic read, anducp_ep_rkey_unpack()still dereferences the raw endpoint after a separate preflight. Without a shared synchronization boundary around endpoint close/clearing and metadata construction, the endpoint can disappear or fail between the check and the UCX call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/plugins/ucx/ucx_utils.cpp` around lines 89 - 117, Serialize UCX endpoint lifetime operations across the error callback and endpoint users: introduce or reuse a shared synchronization boundary covering close/endpoint clearing, metadata construction around ucp_ep_rkey_unpack(), and the state/error-callback paths in nixlUcxEp::checkTxState, nixlUcxEp::setState, and the callback shown here. In src/plugins/ucx/ucx_utils.cpp lines 89-117, update the callback/state handling to participate in that synchronization; in src/plugins/ucx/ucx_backend.cpp lines 969-974, hold the same boundary through endpoint validation and ucp_ep_rkey_unpack() rather than relying on a separate preflight check.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/plugins/ucx/ucx_utils.cpp`:
- Around line 135-138: Replace the deprecated ucp_ep_close_nb() call in the
force-close path with ucp_ep_close_nbx(), using UCP_EP_CLOSE_FLAG_FORCE and the
existing request handling. Confirm the project’s minimum UCX version supports
the newer API, and keep the connected path’s existing newer-API usage
consistent.
- Around line 134-152: Make closeImpl’s FAILED handling idempotent: after
force-closing and releasing a valid endpoint in the FAILED case, transition the
state to DISCONNECTED, or return immediately when FAILED already has no endpoint
handle. Preserve the existing remote-disconnect result for the initial failed
closure while preventing subsequent calls from invoking ucp_ep_close_nb with
null.
In `@test/gtest/unit/ucx/ucx_ep_rkey_test.cpp`:
- Line 26: Add an assertion immediately after capturing the native endpoint in
the test’s relevant flow to verify that getEp() returns a non-null valid handle.
Ensure the assertion specifically checks the captured nativeEndpoint, so the
test cannot pass when both endpoint retrievals return nullptr.
---
Outside diff comments:
In `@src/plugins/ucx/ucx_utils.cpp`:
- Around line 89-117: Serialize UCX endpoint lifetime operations across the
error callback and endpoint users: introduce or reuse a shared synchronization
boundary covering close/endpoint clearing, metadata construction around
ucp_ep_rkey_unpack(), and the state/error-callback paths in
nixlUcxEp::checkTxState, nixlUcxEp::setState, and the callback shown here. In
src/plugins/ucx/ucx_utils.cpp lines 89-117, update the callback/state handling
to participate in that synchronization; in src/plugins/ucx/ucx_backend.cpp lines
969-974, hold the same boundary through endpoint validation and
ucp_ep_rkey_unpack() rather than relying on a separate preflight check.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: ee2a4e92-e78b-401e-90b2-6cc8036db892
📒 Files selected for processing (5)
src/plugins/ucx/ucx_backend.cppsrc/plugins/ucx/ucx_utils.cppsrc/plugins/ucx/ucx_utils.htest/gtest/unit/ucx/meson.buildtest/gtest/unit/ucx/ucx_ep_rkey_test.cpp
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/plugins/ucx/ucx_utils.cpp (1)
89-106: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftSerialize endpoint state transitions with endpoint-handle access.
err_cb()andcloseImpl()can both observeCONNECTED, whilesetState()performs a separate load/assert/store. One thread can close or nullephwhile the callback reads it, and the callback can then overwrite the close transition or trigger thenew_state != old_stateassertion. Protect the state andephlifetime with a shared synchronization protocol; atomic state alone is insufficient.Also applies to: 113-117
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/plugins/ucx/ucx_utils.cpp` around lines 89 - 106, Synchronize the endpoint state and eph lifetime across err_cb() and closeImpl() using the same shared mutex or equivalent protocol, rather than relying on the atomic state alone. Hold that synchronization while validating/accessing eph and while performing the state transition, so closeImpl() cannot null or close the handle concurrently and err_cb() cannot overwrite the close transition or violate setState()’s transition assertion.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/plugins/ucx/ucx_utils.cpp`:
- Around line 89-106: Synchronize the endpoint state and eph lifetime across
err_cb() and closeImpl() using the same shared mutex or equivalent protocol,
rather than relying on the atomic state alone. Hold that synchronization while
validating/accessing eph and while performing the state transition, so
closeImpl() cannot null or close the handle concurrently and err_cb() cannot
overwrite the close transition or violate setState()’s transition assertion.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: ebe75fb7-19f6-4f1e-a6be-b1f5befb63c1
📒 Files selected for processing (2)
src/plugins/ucx/ucx_utils.cppsrc/plugins/ucx/ucx_utils.h
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/plugins/ucx/ucx_utils.cpp (1)
121-148: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftSerialize
ephuse with endpoint teardown.
state_is atomic, butephis not protected by a lifetime lock. A rkey path can pass its state check and call UCX while this path force-closes and nulls the same native handle. Hold a shared lock across the state check plusucp_ep_rkey_unpack(), and an exclusive lock throughoutcloseImpl(), to prevent the reported close-during-unpack crash.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/plugins/ucx/ucx_utils.cpp` around lines 121 - 148, Protect endpoint handle lifetime with the existing synchronization mechanism: acquire an exclusive lock for the entire closeImpl() operation, and hold a shared lock across the rkey path’s state validation and ucp_ep_rkey_unpack() call. Ensure eph cannot be force-closed or nulled while the rkey operation is using it, while preserving the current state and error handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/plugins/ucx/ucx_utils.cpp`:
- Around line 121-148: Protect endpoint handle lifetime with the existing
synchronization mechanism: acquire an exclusive lock for the entire closeImpl()
operation, and hold a shared lock across the rkey path’s state validation and
ucp_ep_rkey_unpack() call. Ensure eph cannot be force-closed or nulled while the
rkey operation is using it, while preserving the current state and error
handling.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 0d9a3fa1-a0de-4c33-9e27-bb0a989755b0
📒 Files selected for processing (1)
src/plugins/ucx/ucx_utils.cpp
|
/build |
1 similar comment
|
/build |
|
Hi @iyastreb, it looks like this CI failure is unrelated to my code. |
|
🤖 CI Triage Agent — TL;DR: The vLLM and SGLang sanity stages failed because the model prefetch step hit HuggingFace's anonymous rate limit (HTTP 429) — an infra/environment issue, not the PR's code. Fix by authenticating the download with an Full analysisSummary: Stages "Run vLLM sanity" (#715) and "Run SGLang sanity" (#674) failed; the wheel builds themselves succeeded (the aarch64 build_helper failures #488/#489 are the pre-sanity pipeline-stop steps, superseded by the successful rebuilds #536/#567/#598). Root cause: In Implicated commit: unknown (environmental — HuggingFace anonymous rate limit; not caused by the PR under test) File: Suggested fix:
Related: none found. |
|
🤖 CI Triage Agent — TL;DR: The "Allocate DL EP Environment" stage failed because the Slurm Full analysisSummary: Stage 252 ( Root cause: Implicated commit: unknown — not code-related; commit fb6f84a only rides along with the CI infra failure. File: CI pipeline Slurm allocation step ( Suggested fix: Re-run the build — this is a transient resource-contention timeout on the Related: none found. |
|
/ok to test a8030d1 |
@iyastreb This test is now included in CI. |
|
/build |
|
/ok to test 926fc09 |
|
🤖 CI Triage Agent — TL;DR: The "Run CPP tests" stage failed at Full analysisSummary: The Root cause: Implicated commit: unknown — the File: Suggested fix: Either (a) guard the invocation like the adjacent DOCA tests — Related: PR #1987 (the triggering PR, but not the cause); none of the searched issues track this |
|
🤖 CI Triage Agent — TL;DR: The build and compile succeeded; the job failed in the "Allocate DL EP Environment" stage because a Slurm Full analysisSummary: Root cause: Implicated commit: none — not caused by [REDACTED:Hex High Entropy String] / PR #1987. File: Jenkins pipeline allocation step invoking Suggested fix: Retry the build once GB200 ( Related: none found.
|
|
🤖 CI Triage Agent — TL;DR: The Full analysisSummary: Post-build GPU environment allocation failed; Root cause: In stage 200 the pipeline ran Implicated commit: none — not a code regression (unrelated to PR #1987's UCX rkey/endpoint race fix at commit 6a8856a). File: n/a — CI infrastructure (SLURM allocation in the pipeline Suggested fix: Retry the build once GB200 ( Related: PR #1987 (the triggering change, unrelated to the failure). |
|
Hi @iyastreb, it looks like this CI failure is unrelated to my code. |
## What? Drain the worker in `nixlUcxBackendH::release()` until every in-flight request has completed, instead of relying on `ucp_request_cancel()` alone. ## Why? `release()` currently cancels outstanding requests and frees them immediately, with a TODO noting "it may not be enough to cancel UCX request". It isn't: `ucp_request_cancel()` is a no-op for RMA/zcopy operations, and `ucp_request_free()` does not synchronously complete internal requests. So a request can still be in flight after `release()` returns. When the caller subsequently deregisters its memory, `ucp_mem_unmap()` frees the memory handle while a zcopy completion is still pending. That completion later runs on the progress thread and dereferences the freed handle, giving a use-after-free that surfaces as a SIGSEGV inside `ucp_memh_put()`. Draining until the request reaches a terminal state guarantees the zcopy completion callback runs while the memh is still valid. This is distinct from the rkey-unpack / endpoint-teardown race fixed in ai-dynamo#1987 — that one races unpack against teardown, this one races zcopy completion against `ucp_mem_unmap()`. ## How was it tested? Carried in production at Fireworks AI on a NIXL 0.10.0-based build for disaggregated-prefill KV cache transfer, where the use-after-free was originally observed. Forward-ported here to `main`, which still contains the original TODO and the same cancel-and-free path.
## What? Drain the worker in `nixlUcxBackendH::release()` until every in-flight request has completed, instead of relying on `ucp_request_cancel()` alone. ## Why? `release()` currently cancels outstanding requests and frees them immediately, with a TODO noting "it may not be enough to cancel UCX request". It isn't: `ucp_request_cancel()` is a no-op for RMA/zcopy operations, and `ucp_request_free()` does not synchronously complete internal requests. So a request can still be in flight after `release()` returns. When the caller subsequently deregisters its memory, `ucp_mem_unmap()` frees the memory handle while a zcopy completion is still pending. That completion later runs on the progress thread and dereferences the freed handle, giving a use-after-free that surfaces as a SIGSEGV inside `ucp_memh_put()`. Draining until the request reaches a terminal state guarantees the zcopy completion callback runs while the memh is still valid. The drain is bounded by `NIXL_UCX_REQUEST_DRAIN_TIMEOUT` (default 10s) so that `release()` cannot block forever on a transfer that is wedged on an endpoint which is still alive but no longer making progress. If the deadline expires the request is deliberately leaked rather than freed, since freeing a request that is still in flight is the very use-after-free this change is avoiding, and the error log tells the caller its memory is not safe to deregister. Progress is reported every `NIXL_UCX_WARNING_TIMEOUT` (default 5s), matching the existing wait loop in `mem_list.cpp`. This is distinct from the rkey-unpack / endpoint-teardown race fixed in ai-dynamo#1987 — that one races unpack against teardown, this one races zcopy completion against `ucp_mem_unmap()`. ## How was it tested? Carried in production at Fireworks AI on a NIXL 0.10.0-based build for disaggregated-prefill KV cache transfer, where the use-after-free was originally observed. Forward-ported here to `main`, which still contains the original TODO and the same cancel-and-free path. Co-authored-by: Cursor <cursoragent@cursor.com>
## What? Drain the worker in `nixlUcxBackendH::release()` until every in-flight request has completed, instead of relying on `ucp_request_cancel()` alone. ## Why? `release()` currently cancels outstanding requests and frees them immediately, with a TODO noting "it may not be enough to cancel UCX request". It isn't: `ucp_request_cancel()` is a no-op for RMA/zcopy operations, and `ucp_request_free()` does not synchronously complete internal requests. So a request can still be in flight after `release()` returns. When the caller subsequently deregisters its memory, `ucp_mem_unmap()` frees the memory handle while a zcopy completion is still pending. That completion later runs on the progress thread and dereferences the freed handle, giving a use-after-free that surfaces as a SIGSEGV inside `ucp_memh_put()`. Draining until the request reaches a terminal state guarantees the zcopy completion callback runs while the memh is still valid. The drain is bounded by `NIXL_UCX_REQUEST_DRAIN_TIMEOUT` (default 10s) so that `release()` cannot block forever on a transfer that is wedged on an endpoint which is still alive but no longer making progress. If the deadline expires the request is deliberately leaked rather than freed, since freeing a request that is still in flight is the very use-after-free this change is avoiding, and the error log tells the caller its memory is not safe to deregister. Progress is reported every `NIXL_UCX_WARNING_TIMEOUT` (default 5s), matching the existing wait loop in `mem_list.cpp`. This is distinct from the rkey-unpack / endpoint-teardown race fixed in ai-dynamo#1987 — that one races unpack against teardown, this one races zcopy completion against `ucp_mem_unmap()`. ## How was it tested? Carried in production at Fireworks AI on a NIXL 0.10.0-based build for disaggregated-prefill KV cache transfer, where the use-after-free was originally observed. Forward-ported here to `main`, which still contains the original TODO and the same cancel-and-free path. Co-authored-by: Cursor <cursoragent@cursor.com>
## What? Drain the worker in `nixlUcxBackendH::release()` until every in-flight request has completed, instead of relying on `ucp_request_cancel()` alone. ## Why? `release()` currently cancels outstanding requests and frees them immediately, with a TODO noting "it may not be enough to cancel UCX request". It isn't: `ucp_request_cancel()` is a no-op for RMA/zcopy operations, and `ucp_request_free()` does not synchronously complete internal requests. So a request can still be in flight after `release()` returns. When the caller subsequently deregisters its memory, `ucp_mem_unmap()` frees the memory handle while a zcopy completion is still pending. That completion later runs on the progress thread and dereferences the freed handle, giving a use-after-free that surfaces as a SIGSEGV inside `ucp_memh_put()`. Draining until the request reaches a terminal state guarantees the zcopy completion callback runs while the memh is still valid. The drain is bounded by `NIXL_UCX_REQUEST_DRAIN_TIMEOUT` (default 10s) so that `release()` cannot block forever on a transfer that is wedged on an endpoint which is still alive but no longer making progress. If the deadline expires the request is deliberately leaked rather than freed, since freeing a request that is still in flight is the very use-after-free this change is avoiding, and the error log tells the caller its memory is not safe to deregister. Progress is reported every `NIXL_UCX_WARNING_TIMEOUT` (default 5s), matching the existing wait loop in `mem_list.cpp`. This is distinct from the rkey-unpack / endpoint-teardown race fixed in ai-dynamo#1987 — that one races unpack against teardown, this one races zcopy completion against `ucp_mem_unmap()`. ## How was it tested? Carried in production at Fireworks AI on a NIXL 0.10.0-based build for disaggregated-prefill KV cache transfer, where the use-after-free was originally observed. Forward-ported here to `main`, which still contains the original TODO and the same cancel-and-free path.
…les (#2027) Follow-up to #1811 / #1987. `postXferReq` guards the remote section by name only, so after a disconnect invalidates (frees) gen-N metadata and the peer is re-registered (gen-N+1), posting a gen-N xfer handle dereferences the freed `nixlUcxPublicMetadata` pinned in `targetDescs` — segfault in `sendXferRangeBatch`/`ucp_put_nbx`. Reject stale-generation handles with `NIXL_ERR_NOT_FOUND` (same guard in `estimateXferCost`, which dereferences the same metadata). Verified with a fault-injection reproducer: without the fix the crash reproduces reliably; with it, none. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Prevented remote transfer requests from using outdated metadata after the associated remote resource is re-registered. * Added validation to reject invalid or stale handles during cost estimation and transfer submission. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: jjthomas <jjthomas@users.noreply.github.com>
## What? Drain the worker in `nixlUcxBackendReqH::release()` until every in-flight request has reached a terminal state, instead of returning as soon as `ucp_request_cancel()` and `ucp_request_free()` have been called. ## Why? `release()` currently cancels outstanding requests and frees them, with a TODO noting "it may not be enough to cancel UCX request". It isn't, but not for the reason the request-lifetime reading suggests. `ucp_request_cancel()` is a no-op for RMA: it only acts on requests carrying `UCP_REQUEST_FLAG_RECV_TAG`. So `release()` can return with the operation still outstanding. The hazard that creates is on the *memory handle*, not on the request. NIXL posts RMA with `UCP_OP_ATTR_FIELD_MEMH` and a memh from `ucp_mem_map()`. UCX stores that pointer in the request (`ucp_datatype_iter`, `type.contig.memh`) without taking a reference - a user memh is not reference counted. When the operation completes, `ucp_datatype_iter_cleanup()` calls `ucp_datatype_iter_mem_dereg_single()` -> `ucp_memh_put()`, which dereferences `memh->context` and `memh->parent`. If the caller has deregistered in the meantime, `ucp_mem_unmap()` -> `ucp_memh_cleanup()` has already `ucs_free()`d that memh, and the completion faults inside `ucp_memh_put()`. Releasing the request object itself is safe either way: `ucp_request_release_common()` returns an *uncompleted* request to the pool only by marking it `UCP_REQUEST_FLAG_RELEASED`; UCX calls `ucp_request_put()` when the request later completes. Draining protects the memory handle, not the request, so `reqRelease()` is now called unconditionally. The drain is bounded by `NIXL_UCX_REQUEST_DRAIN_TIMEOUT` (default 10s) so that `release()` cannot block forever on a wedged transfer. If the deadline expires, the request is still released and an error tells the caller its memory is not safe to deregister - there is nothing better `release()` can do without a real abort primitive. Progress is reported every `NIXL_UCX_WARNING_TIMEOUT` (default 5s), matching the existing wait loop in `mem_list.cpp`. This is distinct from the rkey-unpack / endpoint-teardown race fixed in ai-dynamo#1987 - that one races unpack against teardown, this one races RMA completion against `ucp_mem_unmap()`. ## How was it tested? `TestTransferRelease.InFlightXferIsDrainedBeforeReleaseReturns` in `test/gtest/test_transfer.cpp` posts a 64 MiB READ over a TCP-pinned UCX data path, releases the handle while the read is still in flight, and checks the destination is fully written by the time `releaseXferReq()` returns. A read completes only once its data has landed locally, so this observes the drain directly and needs no RDMA hardware. Also carried in production at Fireworks AI on a NIXL 0.10.0-based build for disaggregated-prefill KV cache transfer, where the `ucp_memh_put()` SIGSEGV was originally observed.
## What? Drain the worker in `nixlUcxBackendReqH::release()` until every in-flight request has reached a terminal state, instead of returning as soon as `ucp_request_cancel()` and `ucp_request_free()` have been called. The composite/threadpool half of the same invariant is in the following commit; neither closes the window on its own. ## Why? `release()` currently cancels outstanding requests and frees them, with a TODO noting "it may not be enough to cancel UCX request". It isn't, but not for the reason the request-lifetime reading suggests. `ucp_request_cancel()` is a no-op for RMA: it only acts on requests carrying `UCP_REQUEST_FLAG_RECV_TAG`. So `release()` can return with the operation still outstanding. The hazard that creates is on the *memory handle*, not on the request. NIXL posts RMA with `UCP_OP_ATTR_FIELD_MEMH` and a memh from `ucp_mem_map()`. UCX stores that pointer in the request (`ucp_datatype_iter`, `type.contig.memh`) without taking a reference - a user memh is not reference counted. When the operation completes, `ucp_datatype_iter_cleanup()` calls `ucp_datatype_iter_mem_dereg_single()` -> `ucp_memh_put()`, which dereferences `memh->context` and `memh->parent`. If the caller has deregistered in the meantime, `ucp_mem_unmap()` -> `ucp_memh_cleanup()` has already `ucs_free()`d that memh, and the completion faults inside `ucp_memh_put()`. Releasing the request object itself is safe either way: `ucp_request_release_common()` returns an *uncompleted* request to the pool only by marking it `UCP_REQUEST_FLAG_RELEASED`; UCX calls `ucp_request_put()` when the request later completes. Draining protects the memory handle, not the request, so `reqRelease()` is now called unconditionally. The drain is bounded by `NIXL_UCX_REQUEST_DRAIN_TIMEOUT` (default 10s) so that `release()` cannot block forever on a wedged transfer. If the deadline expires, the request is still released and an error tells the caller its memory is not safe to deregister - there is nothing better `release()` can do without a real abort primitive. Progress is reported every `NIXL_UCX_WARNING_TIMEOUT` (default 5s), matching the existing wait loop in `mem_list.cpp`. This is distinct from the rkey-unpack / endpoint-teardown race fixed in ai-dynamo#1987 - that one races unpack against teardown, this one races RMA completion against `ucp_mem_unmap()`. ## How was it tested? `TestTransferRelease.InFlightXferIsDrainedBeforeReleaseReturns` in `test/gtest/test_transfer.cpp` posts a 64 MiB READ over a TCP-pinned UCX data path, releases the handle while the read is still in flight, and checks the destination is fully written by the time `releaseXferReq()` returns. A read completes only once its data has landed locally, so this observes the drain directly and needs no RDMA hardware. With the drain removed and nothing else changed, the `ucx` parameterisation sees 8 128 of 67 108 864 bytes landed when `releaseXferReq()` returns, and `ucx_no_pt` sees 0. With the drain in place both are complete. Also carried in production at Fireworks AI on a NIXL 0.10.0-based build for disaggregated-prefill KV cache transfer, where the `ucp_memh_put()` SIGSEGV was originally observed.
…les (#2027) (#2104) Cherry pick of #2027 --------- Follow-up to #1811 / #1987. `postXferReq` guards the remote section by name only, so after a disconnect invalidates (frees) gen-N metadata and the peer is re-registered (gen-N+1), posting a gen-N xfer handle dereferences the freed `nixlUcxPublicMetadata` pinned in `targetDescs` — segfault in `sendXferRangeBatch`/`ucp_put_nbx`. Reject stale-generation handles with `NIXL_ERR_NOT_FOUND` (same guard in `estimateXferCost`, which dereferences the same metadata). Verified with a fault-injection reproducer: without the fix the crash reproduces reliably; with it, none. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Prevented remote transfer requests from using outdated metadata after the associated remote resource is re-registered. * Added validation to reject invalid or stale handles during cost estimation and transfer submission. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- ## What? _Describe what this PR is doing._ ## Why? _Justification for the PR. If there is an existing issue/bug, please reference it. For bug fixes, the 'Why?' and 'What?' can be merged into a single item._ ## How? _It is optional, but for complex PRs, please provide information about the design, architecture, approach, etc._ Signed-off-by: jjthomas <jjthomas@users.noreply.github.com> Co-authored-by: James Thomas <jamesjoethomas@gmail.com>
|
@iyastreb We need to think about how to add test coverage for this kind of issues |
6 participants
Fix #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:
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:
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
Summary by CodeRabbit
Bug Fixes
Tests