PLUGINS/UCX: wait for in-flight transfers before release() returns - #2044
vedularaghu wants to merge 2 commits into
Conversation
|
👋 Hi vedularaghu! 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:
📝 WalkthroughWalkthrough
ChangesUCX request lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Release as nixlUcxBackendReqH::release()
participant Drain as drainRequest()
participant Worker as UCX worker
participant Request as UCX request
Release->>Drain: Cancel and drain request
Drain->>Worker: Progress worker
Worker->>Request: Poll completion
Request-->>Drain: Return completion state
Drain-->>Release: Return drained or timed out
Release->>Request: Release or intentionally leak
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 2
🤖 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_backend.cpp`:
- Around line 141-151: Replace the multi-line block comment in ucx_backend.cpp
with equivalent regular // comments at the same location, keeping the existing
rationale about draining the worker and the ucp_memh_put lifetime issue intact.
Update only the comment style around the UCX zcopy completion handling; do not
change the surrounding logic or symbols such as ucp_request_cancel,
ucp_request_free, or ucp_memh_put.
- Around line 152-155: Update the request-wait logic around
ucp_request_check_status in the UCX backend so UCS_ERR_UNSUPPORTED cannot cause
the loop to exit while the request remains in flight. Use a supported UCX
completion mechanism, or retain the request and associated memory handle until
asynchronous completion before calling ucp_request_free or deregistering memory.
🪄 Autofix
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: 87731186-cfa3-459c-8060-ef56494d3026
📒 Files selected for processing (1)
src/plugins/ucx/ucx_backend.cpp
3f304f1 to
067f176
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
067f176 to
ec52236
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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_backend.cpp`:
- Around line 141-155: Update the drain-timeout path in release() so a false
result from drainRequest(req) retains ownership of the transfer’s memory handle
until UCX reaches terminal completion, rather than only leaking req. Ensure
requests_.clear(), conn_.reset(), and any caller-triggered deregistration cannot
release the backing memory while ucp_memh_put() may still run; alternatively,
block release() until terminal completion.
- Line 193: Update drainRequest around worker_->progress() to inspect its
returned progress count and briefly sleep when it is zero, matching the backoff
behavior used in mem_list.cpp. Preserve the existing loop and timeout behavior
while avoiding busy-polling during periods with no UCX progress.
🪄 Autofix
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: 83b7eca2-e90a-4981-9c77-e7b957d5e2db
📒 Files selected for processing (1)
src/plugins/ucx/ucx_backend.cpp
## 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.
ec52236 to
15f357e
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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_backend.cpp`:
- Around line 175-191: Validate the durations returned for
NIXL_UCX_REQUEST_DRAIN_TIMEOUT and NIXL_UCX_WARNING_TIMEOUT as strictly positive
before entering the drain loop in release(). Reject zero or negative values,
preventing immediate warning retries and leaked in-flight requests.
🪄 Autofix
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: 89532ee6-4b4a-438c-8040-d601e41d94aa
📒 Files selected for processing (1)
src/plugins/ucx/ucx_backend.cpp
15f357e to
06938ef
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
| // TODO: Need process this properly. | ||
| // it may not be enough to cancel UCX request | ||
| worker_->reqCancel(req); | ||
| if (!drainRequest(req)) { |
There was a problem hiding this comment.
Currently there is no proper "abort" functionality in NIXL, we are discussing it.
@mkhazraee
But I'm afraid that this approach does not really solve the problem, just hides it a bit by doing 10s extra polling, but then it still fails the same way..
There was a problem hiding this comment.
Agreed - this isn't an abort, and I don't want to claim it is.
To split the two cases apart:
Transfer can still make progress. This is the ordinary releaseXferReq()-mid-transfer path, and today release() returns with the RMA still outstanding. The new gtest measures it: with the drain removed, only 8 128 of 67 108 864 bytes of a READ had landed when releaseXferReq() returned (0 bytes with no progress thread). The caller is free to deregisterMem() at that point, and ucp_mem_unmap() frees a memh the request still holds a raw pointer to. The drain does close that window, and it costs ~300 ms for a 64 MiB transfer.
Transfer can never make progress. You're right that the timeout doesn't fix anything here. After 10 s release() returns anyway, the operation is still outstanding, and the caller's memory still isn't safe to deregister - all the deadline buys is an error line instead of silence. Only a real abort helps, and that's yours to design.
So I'd frame this as: it fixes the case where waiting is sufficient, and it makes the case where it isn't sufficient visible instead of silent. If you'd rather not carry the 10 s knob at all and wait for proper abort support, we're happy to hold or close it - the part we'd like to avoid keeping is release() returning, in the normal case, while UCX still holds a pointer to a memh the caller is about to unmap.
Separately: the production evidence we gathered for #2047 supports your diagnosis on the other thread. On a wedged pod (NIXL 1.3.2, UCX debug logging) there were 156 set_ep_failed status Endpoint timeout on lane[N] events over the wedge, 27 of them on the exact rank whose handle stalled, while NIXL reported NIXL_IN_PROG throughout. A FAILED endpoint with a permanently outstanding request is exactly what you described, and it points at your checkConnection()-on-NIXL_IN_PROG POC as the primary fix rather than anything in these two PRs. Details in #2047.
brminich
left a comment
brminich
left a comment
There was a problem hiding this comment.
it does not seem to solve a real issue.
Pls submit a unit test, which would show this fix is working
| worker_->reqCancel(req); | ||
| if (!drainRequest(req)) { | ||
| // Still in flight past the deadline, so the request object is not | ||
| // ours to reclaim: freeing it now lets a later completion write into |
There was a problem hiding this comment.
no, ucp_request_free is not doing anything with uncompleted request. Request is returned to memory pool only when it is completed and ucp_request_free() is called
There was a problem hiding this comment.
You're right, and thanks for catching it - I had the request lifetime wrong. ucp_request_release_common() only marks an uncompleted request UCP_REQUEST_FLAG_RELEASED; ucp_request_put() happens from ucp_request_complete() when it finally completes, so the request is never handed back early and freeing it in flight is safe.
That also means the "leak the request on timeout" branch I had here was a straight regression - a real req_mp leak for no benefit - so I've removed it. reqRelease() is unconditional again, as before this PR.
Re-deriving the crash from scratch, the hazard is on the memh, not the request. We post RMA with UCP_OP_ATTR_FIELD_MEMH and a memh from ucp_mem_map(). ucp_datatype_contig_iter_init() stores it raw, without taking a reference, because a user memh isn't refcounted:
dt_iter->type.contig.memh = param->memh;and on completion ucp_datatype_iter_cleanup() -> ucp_datatype_iter_mem_dereg_single() -> ucp_memh_put() dereferences memh->context and memh->parent. If the caller has deregistered in between, ucp_mem_unmap() -> ucp_memh_cleanup() has already ucs_free()d it, and that dereference is the SIGSEGV in ucp_memh_put() we saw.
So what release() has to guarantee is "UCX is done with the memh before the caller may unmap it", and ucp_request_cancel() can't provide it for RMA since it only acts on UCP_REQUEST_FLAG_RECV_TAG. The PR description has been rewritten around this; please re-read it rather than the original.
## What? Wait for `pendingReqs` to drain in `nixlUcxCompositeBackendReqH::release()` before resetting the shared state. ## Why? Setting `sharedState_->status` to a failed value stops new chunks from starting, but it does not complete chunk requests that have already been posted, and resetting `sharedState_` only drops this handle's reference to it. The chunk handles are cancelled and released later, on the threadpool worker that owns them. So `release()` can return while a chunk request is still in flight. The caller is then free to deregister its memory, which is the same hazard as the non-composite path: chunks are posted with `UCP_OP_ATTR_FIELD_MEMH`, and UCX holds that `ucp_mem_h` in the request as a plain pointer without taking a reference. When the operation completes, `ucp_datatype_iter_cleanup()` calls `ucp_memh_put()` on it, which dereferences `memh->context` and `memh->parent`. If `ucp_mem_unmap()` has already `ucs_free()`d that memh, the completion faults inside `ucp_memh_put()`. Only the threadpool worker that owns a chunk may progress it, so the wait polls the pending counter instead of driving progress itself. The wait is bounded by `NIXL_UCX_REQUEST_DRAIN_TIMEOUT` (default 10s) so that `release()` cannot block forever on a wedged chunk; if the deadline expires the handle is released anyway and an error 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 the composite half of the fix in the preceding commit and only closes the window in combination with it: a cancelled chunk is drained by `nixlUcxBackendReqH::release()`, and this wait is what keeps the composite `release()` from returning before that has happened. ## How was it tested? The `ucx_threadpool` and `ucx_threadpool_no_pt` parameterisations of `TestTransferRelease.InFlightXferIsDrainedBeforeReleaseReturns` cover the composite path. With this commit reverted and the preceding one kept, they fail with half the destination written (progress thread) or none of it (no progress thread), while the non-composite parameterisations still pass.
06938ef to
68fd77e
Compare
|
@brminich @iyastreb - pushed a rewrite. Summary of what changed and what I conceded: 1. You were right about 2. The bug is on the memory handle, and it is real. RMA is posted with 3. Unit test, as asked.
4. @iyastreb - agreed on the limitation, no argument from me. The drain works when the transfer can still make progress. When it can't, the deadline expires, Also worth flagging on your other point: we went back to the UCX debug logs from the wedged production pod, and they support your #2045 is the composite/threadpool half of this change and now stacks on this branch, since neither closes the window alone. |
68fd77e to
4d758dc
Compare
vedularaghu
changed the title
|
@brminich @iyastreb - folded #2045 into this PR and closed it, so there is one change and one argument to evaluate rather than two. The composite/threadpool path is the same bug and the two halves interlock (neither closes the window alone), which was not obvious with them split across PRs. Kept as two commits so each code path stays separately reviewable, and because each has its own negative control. While re-verifying the merged branch I got two things worth adding, both reproducible on demand with the fix reverted: Reverting the simple-path drain kills all four parameterisations, including the threadpool ones that still have the composite wait - the wait is meaningless if the chunk release it waits for does not drain. 5 of 14 runs also died on UCX's own fatal assertion rather than merely failing the data check: UCX refusing to destroy a TCP endpoint that still has operations queued on it. Reverting only the composite wait segfaulted 1 run in 3, with the read landing in a buffer the caller had already freed: The chunk drain does run - just on the threadpool worker, after the composite Being precise about what that does and does not show: neither reproduces the exact With both commits: 25 consecutive runs of the four cases, no failures and no crashes. The concession on |
|
Currently real abort if not supported by NIXL. |
|
Hi All We have a similar issue in the POSIX / io_uring backend where we can have outstanding requests in the ring. In my opinion the correct behavior is to panic (= terminate the process) when the handle is dropped and there are outstanding requests because of this issue you mentioned in the PR description:
I was circulating this document internally before. We need to agree on a NIXL wide approach and then all the baekcnd |
4 participants
What?
Make
release()wait for UCX to be done with a transfer's memory before it returns, on both the simple-handle and the composite/threadpool path.nixlUcxBackendReqH::release()until every in-flight request has reached a terminal state.nixlUcxCompositeBackendReqH::release()wait forpendingReqsto drain before resetting the shared state.Two commits because they touch different classes and each has its own negative control, but they are one invariant and neither half closes the window alone - see below. This supersedes #2045, which was folded in here so the argument only has to be evaluated once.
Correction to the original description
@brminich is right, and the first version of this PR was wrong on the mechanism. It claimed
ucp_request_free()"does not synchronously complete internal requests" and implied the use-after-free came from freeing an in-flight request. Inucp_request_release_common():An uncompleted request is only marked
UCP_REQUEST_FLAG_RELEASED; it goes back to the pool fromucp_request_complete()when it finally completes. Soucp_request_free()on an in-flight request is safe and is the supported way to hand it back.The first version also leaked the request when the drain deadline expired, on the theory that freeing it was the hazard. That was strictly worse than the code it replaced: a genuine
req_mpleak for no benefit. It is gone -reqRelease()is unconditional, exactly as before this PR.Why? (re-derived)
ucp_request_cancel()is a no-op for RMA: it only acts on requests carryingUCP_REQUEST_FLAG_RECV_TAG. Sorelease()can return with the operation still outstanding. That much of the original TODO holds.The hazard that creates is on the memory handle, not on the request. NIXL posts RMA with
UCP_OP_ATTR_FIELD_MEMHand a memh fromucp_mem_map()(nixlUcxEp::read/writeinucx_utils.cpp). UCX stores that pointer in the request without taking a reference, because a user memh is not reference counted:When the operation completes,
ucp_datatype_iter_cleanup()callsucp_datatype_iter_mem_dereg_single()->ucp_memh_put(), which dereferencesmemh->contextandmemh->parentbefore doing anything else. If the caller has deregistered in the meantime,ucp_mem_unmap()->ucp_memh_cleanup()has alreadyucs_free()d that memh, and the completion faults insideucp_memh_put(). That matches the SIGSEGV we saw in production.More generally, the operation is still reading and writing the caller's buffers and still holds the endpoint. So the invariant
release()has to establish is not "the request object is reclaimed" but "UCX is done with this transfer's memory before the caller may unmap or free it".Distinct from the rkey-unpack / endpoint-teardown race fixed in #1987 - that one races unpack against teardown, this one races RMA completion against the caller reclaiming memory.
Why both halves are needed
On the threadpool path the two changes interlock:
release()sets the shared status to a failure, which makes the owning worker runnixlUcxChunkBackendReqH::complete()->nixlUcxBackendReqH::release()on the chunk. Without commit 1 that release does not drain, so the chunk's requests can still be in flight when it returns.release()does not wait for that to happen at all - it returns before the worker has even looked at the chunk.The controls below show this directly: reverting commit 1 alone breaks all four parameterisations, including the threadpool ones that still have commit 2.
Unit test
@brminich asked for a test showing the fix works.
TestTransferRelease.InFlightXferIsDrainedBeforeReleaseReturnsintest/gtest/test_transfer.cpp:postXferReq()returnedNIXL_IN_PROG),releaseXferReq()returns,A read completes only once its data has landed locally, so this observes the drain directly rather than relying on a crash. The data path is pinned to
UCX_TLS=tcp, because over shm/self UCX copies inline and nothing is ever in flight to drain. 64 descriptors puts the batch above the fixture'ssplit_batch_size, so theucx_threadpool*parameterisations take the composite path. No RDMA hardware needed; the four cases take 0.2-0.6 s each.Negative control 1 - revert commit 1 (simple-path drain), keep commit 2
releaseXferReq()returneducxucx_no_ptucx_threadpoolucx_threadpool_no_ptAll four fail, including the threadpool cases that still have the composite wait - the wait is meaningless if the chunk release it waits for does not drain.
5 of 14 runs also died inside UCX with its own fatal assertion, rather than just failing the data check:
That is UCX refusing to destroy a TCP endpoint that still has operations queued on it.
Negative control 2 - revert commit 2 (composite wait), keep commit 1
releaseXferReq()returneducx,ucx_no_ptucx_threadpoolucx_threadpool_no_pt1 of 3 runs segfaulted instead, with the read landing in a buffer the caller had already freed:
The chunk drain does run, but on the threadpool worker, after the composite
release()has already returned and the caller has freed the destination.To be precise about what these do and do not show: neither control reproduces the exact
ucp_memh_put()SIGSEGV we saw in production - these are over TCP, not RDMA, and the fault lands at a different site. They are the same root cause (release()returning while the RMA is live) surfacing wherever the transport happens to touch the reclaimed resource first.With both commits, 25 consecutive runs of the four cases: no failures, no crashes.
Test plan
debugoptimized, UCX 1.20, built clean from this branchtest/gtest/unit: 84/84test/gtest*TestTransfer*: 57/57, including the four new casesclang-format-diff-19clean on every line added againstmainWhat this does not fix
@iyastreb is right that this is not an abort, and I do not want to oversell it.
The drain closes the window whenever the transfer can still make progress - which is the ordinary
releaseXferReq()-mid-transfer case the test reproduces. It does not help when the transfer can never complete: after the deadlinerelease()returns anyway, the operation is still outstanding, and the caller's memory still is not safe to reclaim. All the timeout buys there is an error message instead of silence.That residual case needs the real abort primitive you are discussing, and if you would rather this waited for that, we are happy to hold it. What we would like to avoid is
release()continuing to return, in the ordinary case, while UCX still holds the caller's buffers and endpoint.Separately, on the
NIXL_IN_PROG/FAILED-endpoint point from the other thread: production UCX debug logs from a wedged pod back your diagnosis - 156set_ep_failed status Endpoint timeout on lane[N]events over the wedge, 27 on the exact rank whose handle stalled, all while NIXL reportedNIXL_IN_PROG. Written up in #2047, which we have reframed as a backstop rather than a fix for that.