Repository navigation
PLUGINS/UCX: Add scatter-gather list (SGL) put path - #1835
Conversation
|
👋 Hi michal-shalev! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
61383da to
0aa8bbb
Compare
|
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:
📝 WalkthroughWalkthroughThe PR adds optional UCX SGL API detection, transfer descriptors, endpoint posting helpers, runtime enablement, and an SGL-based write-transfer path with request and flush tracking. ChangesUCX SGL support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant nixlUcxEngine
participant sglXfer
participant nixlUcxEp
participant UCX
nixlUcxEngine->>sglXfer: build local and remote segment state
nixlUcxEngine->>nixlUcxEp: postSgl(localView, remoteView, count, req)
nixlUcxEp->>UCX: ucp_put_nbx(...)
UCX-->>nixlUcxEp: request pointer or completion status
nixlUcxEp-->>nixlUcxEngine: NIXL_IN_PROG or completion status
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: 3
🤖 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 1323-1325: The SGL fast path in ucx_backend.cpp currently skips
the per-endpoint batching logic in the send path. Update the NIXL_WRITE branch
in the transfer routine (the block that calls sendXferSgl) so it only takes this
route when the full remote range remote[start_idx:end_idx] resolves to a single
connection, or move the batching/connection-homogeneity check into sendXferSgl
itself. Use the existing transfer flow and metadata helpers in this backend to
verify the range before returning the SGL path.
- Around line 1249-1293: The sendXferSgl() path is using a single
getConnection(remote_agent) for the whole range, which can post remote
addrs/rkeys through the wrong UCX endpoint when descriptors span multiple
connections. Update ucx_backend.cpp so the SGL send is scoped to one connection
by grouping descriptors by rmd->conn (like the non-SGL batching logic) and
posting each group on its own conn->getEp(worker_id), or detect mixed
connections and fall back to the non-SGL path. Keep the existing
local_sgl/remote_sgl setup but build it per connection instead of once for the
whole range.
In `@src/plugins/ucx/ucx_utils.h`:
- Around line 112-117: The public endpoint API `postSgl()` in `ucx_utils.h`
needs a Doxygen block documenting the SGL lifetime contract. Add a `/** ... */`
comment above `postSgl()` that clearly states the caller-owned
`ucp_dt_local_sgl_t` and `ucp_dt_remote_sgl_t` arrays must remain valid until
the associated request completes, so users know the stored UCX pointers must
outlive completion. Use the function name `postSgl()` and the `nixlUcxReq`
request type in the comment to make the ownership and completion requirement
explicit.
🪄 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: 8909565d-7350-4d42-8b9f-c50ca03ce3eb
📒 Files selected for processing (5)
meson.buildsrc/plugins/ucx/ucx_backend.cppsrc/plugins/ucx/ucx_backend.hsrc/plugins/ucx/ucx_utils.cppsrc/plugins/ucx/ucx_utils.h
2218f6c to
9610c29
Compare
michal-shalev
changed the title
9610c29 to
a647483
Compare
michal-shalev
changed the title
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_sgl.h`:
- Around line 33-39: Add a Doxygen-style documentation block immediately before
the public sglXfer class declaration, describing its ownership semantics and
stating that the arrays referenced by the UCX request must outlive that request.
Keep the constructor and class behavior unchanged.
🪄 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: c41315bf-f0fc-449f-810c-1c4ddde17bce
📒 Files selected for processing (4)
src/plugins/ucx/meson.buildsrc/plugins/ucx/ucx_backend.cppsrc/plugins/ucx/ucx_sgl.cppsrc/plugins/ucx/ucx_sgl.h
|
/build |
|
I have just submitted PR for perf degradation caused by previous SGL changes |
|
/build |
|
/build |
|
🤖 CI Triage Agent — TL;DR: The build compiled fine; the job failed because the Slurm Full analysisSummary: Stage 226 "Allocate DL EP Environment" (ucx-v1.22.x branch) failed when Slurm could not grant a compute node within the immediate-allocation deadline. Root cause: After a successful NIXL/nixlbench build and Docker image push (~17:04), the pipeline ran Implicated commit: unknown — not caused by commit c92dd11 or PR #1835; this is a Slurm scheduling/capacity failure external to the code. File: N/A (CI infra — Slurm allocation step invoking Suggested fix: This is an infrastructure/capacity retry situation, not a code fix. Concretely: (1) re-run the build — the queue was simply full of GB200 jobs; (2) if this recurs, raise the Related: none |
|
/build |
Good catch on the NIXL issue, and thanks for fixing it in #2017. PR #2017 addresses a pre-existing issue in the default NIXL path that became more noticeable after openucx/ucx#11258 increased the size of PR #1835 introduces a separate, optional SGL path behind an environment variable and does not modify the default execution path. If you believe there is a similar issue in this PR, please point to the specific code path or demonstrate the regression here. |
|
/build |
|
/ok to test 514c03e |
## What? Add a UCX scatter-gather (SGL) put path. ## Why? Posting one UCX request per descriptor is wasteful for multi-element writes. An SGL put hands UCX the whole descriptor list in one call, reducing per-element overhead. ## How? Off by default, gated at two layers: - **Build-time**: `meson.build` probes for `ucp_dt_local_sgl_t` and defines `HAVE_UCX_SGL_API` only when UCX exposes the SGL API. All SGL code is `#ifdef`-guarded. - **Run-time**: `NIXL_UCX_SGL_ENABLE` (default `false`) toggles it; if set on a build without SGL support, it warns and falls back to the per-descriptor path. For `NIXL_WRITE`, a new `sglXfer` gathers the descriptor range into parallel arrays (addrs, lengths, memhs, rkeys) kept on the request handle. `sendXferRange` short-circuits to `sendXferSgl`, issuing one `postSgl` (`ucp_put_nbx` + `ucp_dt_make_sgl()`) and a single `flushEp` per batch. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added optional UCX scatter-gather list (SGL) offload for write transfers. * When supported and enabled, write operations are submitted using SGL for more efficient handling across multiple segments. * **Compatibility** * UCX SGL support is auto-detected at build time; builds without it won’t include the feature. * Runtime enablement is controlled by `NIXL_UCX_SGL_ENABLE`; if SGL isn’t compiled in, it is forced off with a warning. * When SGL can’t be used, the system falls back to the existing transfer behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
What?
Add a UCX scatter-gather (SGL) put path.
Why?
Posting one UCX request per descriptor is wasteful for multi-element writes. An SGL put hands UCX the whole descriptor list in one call, reducing per-element overhead.
How?
Off by default, gated at two layers:
meson.buildprobes forucp_dt_local_sgl_tand definesHAVE_UCX_SGL_APIonly when UCX exposes the SGL API. All SGL code is#ifdef-guarded.NIXL_UCX_SGL_ENABLE(defaultfalse) toggles it; if set on a build without SGL support, it warns and falls back to the per-descriptor path.For
NIXL_WRITE, a newsglXfergathers the descriptor range into parallel arrays (addrs, lengths, memhs, rkeys) kept on the request handle.sendXferRangeshort-circuits tosendXferSgl, issuing onepostSgl(ucp_put_nbx+ucp_dt_make_sgl()) and a singleflushEpper batch.Summary by CodeRabbit
NIXL_UCX_SGL_ENABLE; if SGL isn’t compiled in, it is forced off with a warning.