Repository navigation
PLUGINS/UCX: Set ERR_HANDLER field only when a handler is provided - #2423
Merged
Merged
Conversation
tomerdav
requested review from
a team,
brminich,
gleon99 and
yosefe
as code owners
October 14, 2026 10:51
|
👋 Hi tomerdav! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Title check The title clearly and concisely describes the main UCX change: setting the ERR_HANDLER field only when a handler is provided. Description check The description includes the required What, Why, and How sections. It explains the uninitialized handler issue, the resulting segfault risk, and the implementation approach.
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Contributor
Author
|
/build |
ovidiusm
reviewed
Oct 14, 2026
Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
tomerdav
changed the title
tomerdav
force-pushed
the
fix/ucx-ep-params-zero-init
branch
from
October 14, 2026 12:32
ca5dbd8 to
de69086
Compare
iyastreb
approved these changes
Oct 14, 2026
Contributor
|
/build |
iyastreb
enabled auto-merge (squash)
October 14, 2026 12:49
ovidiusm
pushed a commit
that referenced
this pull request
Oct 19, 2026
## What?
The UCX backend as the CPU proxy's transport, and the switch that turns
the proxy on:
- `nixlUcxEp::atomicAdd`: a post-mode remote `UCP_ATOMIC_OP_ADD`, for
the counters a receiver waits on to learn
that data has landed (`ucx_utils.{h,cpp}`).
- `nixl::ucx::proxyEngine` (`ucx_proxy_engine.{h,cpp}`): the
`nixlUcxEngine` subclass that `create()` picks for
`device_proxy=true`. It owns a `nixl::proxyRuntime` (#2223) with a
`nixl::proxyTransport` (#2225) over the
engine's UCX workers, and hands out `PROXY` memory views (part 5)
instead of direct ones.
- Proxy threads may outnumber channels, up to one thread per (channel,
peer) ring, while each pass still makes
one call per channel.
- The proxy transport and runtime are ready once constructed: no
`init()`/`shutdown()`, and
`proxyRuntime::create()` returns a running runtime or throws.
- Tests: the backend-parameter table, and the existing single-write
device tests run in both modes. The
real-UCX lifetime test is in the tests PR (#2283).
Direct mode is unchanged and remains the default: without
`device_proxy=true` the engine is main's.
## Why?
Parts 1 to 5 give a GPU a way to hand work to host threads, but nothing
posts that work to a network. This PR
makes UCX the transport, so a kernel written against the device API can
run over any UCX transport the host can
drive (for example where UCX's direct GPU API is unavailable), without
changes to the kernel.
## How?
**Transport** (`ucx_proxy_engine.cpp`)
- The engine's shared workers are partitioned one per (channel, peer)
ring, so each proxy thread progresses only
the UCX workers of the rings it owns, and a ring's operations are
ordered on one UCX worker.
- `submit` posts a PUT (`ucp_put_nbx`) or an atomic add. Before an
atomic, the transport fences the worker, so a
counter update is not visible before the PUTs issued ahead of it on that
ring.
- Request status, terminal release, per-ring progress and quiescence map
directly onto UCX requests and
`ucp_worker_progress`. A post-mode atomic completes inside
`ucp_atomic_op_nbx` and never yields a request, so
only PUTs are tracked.
- `resolveDirectPtrs` fills a remote view's direct pointers with
`ucp_rkey_ptr` where the remote key allows
direct access, and leaves them null otherwise. It skips holes (the null
agent, or no metadata), which the
memory-view manager marks unusable.
**Engine** (`nixl::ucx::proxyEngine`)
- Backend parameters: `device_proxy=true` enables the proxy;
`proxy_channel_count`, `proxy_thread_count`,
`proxy_max_peers`, `proxy_ring_depth` and `proxy_drain_timeout_ms` tune
it (parsed by `nixl::parseProxyConfig`).
Tuning without `device_proxy=true`, `num_threads` and the agent progress
thread are rejected. Without a loaded
device implementation (a build without CUDA), creating an engine with
`device_proxy=true` fails.
- `proxy_thread_count` defaults to the channel count and may go up to
one thread per ring: up to the channel
count each thread owns a block of whole channels, and beyond it each
channel's peers are strided over
thread_count / channel_count threads. Only counts that split the rings
evenly are accepted. Each thread gets
its span of channels and the peers it serves in them when the runtime is
built, so the proxy loop is the
previous per-channel loop over a slice of peers.
- The constructor sizes the shared workers to one per ring (channels x
peers): a smaller `num_workers`, such as
the default of 1 that `getPluginParams()` seeds, is raised to that
count, and a larger one is rejected, since only
the proxy threads progress the workers. It runs with
`NIXL_THREAD_SYNC_RW`, since proxy threads call into UCX outside the
agent
lock, and sets `ucx_error_handling_mode=none`: under `peer`, UCX
disables every AMO protocol and the proxy's
atomics would complete as cancelled, so an explicit `peer` is rejected.
It then creates the runtime, which
starts its threads; the runtime is a member, so it stops its threads
before the base engine destroys the workers.
- `prepMemView` registers the descriptors with the runtime and wraps the
result as a `PROXY` view.
`releaseMemView` retires it: the runtime waits until every operation on
the view reaches a terminal status
and the owning UCX workers are quiescent, and only then frees it. A view
that was never published rolls back
without a runtime drain.
- A failed peer's endpoint does not make release fatal, and a view the
runtime cannot retire logs a warning; a
view that cannot be read is fatal.
**Notifications**: on main, the base engine's notification list relies
on the agent lock, which holds while
only agent calls progress its workers. In proxy mode every shared worker
is a ring worker, so a proxy thread
progresses worker 0 and runs the notification callback outside the agent
lock. `proxyEngine` overrides
`appendNotif` and `getNotifs` with its own mutex, so the regular engine
stays lock-free, and its `getNotifs` only
swaps the list: the proxy threads already progress every shared worker
on each pass, worker 0 with its
notifications and incoming connections included, so the application no
longer competes with them on the ring
workers.
**Build**: the UCX plugin always links `libnixl_device_proxy` and
`libdevice_utils` and builds the transport,
with no compile-time guard; neither library needs CUDA to build. CUDA is
still loaded at run time, through the
device operations library and UCX's own modules.
## Known limitations and follow-ups
To be fixed as follow-up work; none affects the tested configurations.
- **Connection fan-out.** Every UCX worker the engine creates connects
to every peer, as in direct mode, but in
proxy mode there are channels x peers workers and each peer is served by
only one of them, so each rank makes
about channels x peers x peers endpoints, and connection setup grows
with it. Keep `proxy_max_peers` at the
real peer count, as EP does. The fix is to connect only the worker that
serves each peer.
- **Software atomics on GPU memory.** If UCX cannot do a NIC atomic on
the counter's GPU memory (seen with a NIC
that is not PCIe-local to the GPU), it falls back to its software atomic
protocol, which performs the add with a
CPU instruction on the target and crashes the remote process. The
transport should detect that fallback at
setup and fail with a clear error.
- **Unbounded quiesce.** `quiesce` waits for `ucp_worker_flush_nbx`,
which is bounded only by the transport's own
retry timeouts. When the peer above crashed (same host), release waited
indefinitely; the likely cause is a
shared-memory lane, which has no retry timeout, under
`ucx_error_handling_mode=none`, which has no keepalive.
The fix is an application-level deadline that reports the peer as
failed.
- **Proxy threads spin and are not bound to the GPU's NUMA node.**
Across ranks the threads must fit the
physical cores, and threads on the GPU's NUMA node are faster; the proxy
should bind its threads itself and
warn on oversubscription.
## Testing
- L40S over RoCE, main `d1c6fb9e`: unit 191 pass; single-write device
tests 18/18 in direct and proxy mode;
UCX suites 90/90; the proxy cases at 4, 32 and 64 threads over 32
channels (local thread-count override, with
the dropped many-elements test added back locally) 12/12 each; #2283's
real-UCX lifetime test 3/3. Every
commit builds with `-Werror`, and so does the release build;
clang-format clean; PR size 387.
- NIXL EP's CI elastic tests with #2328 on top, on one 8x H100 node: all
six single-node plans, direct and
proxy, NVLink and RDMA only, 88/88.
- `ProxyConfigTest`: parameter parsing, thread splits, and that the
threads' shares cover every ring once.
- Proxy unit, GPU and lifetime tests: #2283. NIXL EP, with performance:
#2328.
## Commits
- `bffebeca` PLUGIN/UCX: add a remote atomic add on the endpoint
- `2a3a1b5d` PLUGIN/UCX: add the device proxy transport
- `5b15101d` PLUGIN/UCX: own a device proxy runtime, enabled by backend
parameters
- `09a8eb00` test: check the device proxy backend parameters
- `71db9bd3` test: exercise UCX proxy mode in device PUT integration
tests
- `5b50ae05` ucx: synchronize base engine notification callbacks and
drains
- `511faf52` test: parse the proxy drain timeout
- `9a41e6c9` test: run the device PUT tests with the execution mode
fixed
- `9a396423` PLUGIN/UCX: keep the proxy declarations private and the
opcode switch exhaustive
- `0261134d` utils: roll back a failed proxy memory view with the
regular release
- `54c54d8b` device_api: pick the proxy ring by the element's remote
agent
- `2f49ebcf` PLUGIN/UCX: move the device proxy into its own engine
- `d4f14f64` utils: let proxy threads own individual rings
- `bba890d7` PLUGIN/UCX: run the device proxy with error mode none
- `f956bf6b` utils: construct the proxy transport and runtime ready to
use
Based on main; part 5 (#2326) has merged. The tests PR #2283 sits on top
of this one.
## Changes in this revision
For hbadihi's and rakhmets's reviews; every change is folded into the
commit it belongs to.
- The proxy engine is `nixl::ucx::proxyEngine` (`2f49ebcf`). Its
transport reads the memory handle through
`getMem()`, so the friend declaration is gone and `nixlUcxEp::write()`
takes the handle as const.
- `proxyRuntime::create()` builds the runtime and starts its workers,
and throws on any failure; `startWorkers()`
and the runtime's stopped state are gone (`f956bf6b`).
- `proxyRuntime::releaseMemView()` returns nothing and warns on failure,
and the engine no longer treats a
failed release as fatal (`f956bf6b`, `2f49ebcf`).
- The runtime, not the memory-view manager, maps each remote agent to a
peer and checks `proxy_max_peers`
(`54c54d8b`).
- The transport frees a completed request with `ucp_request_free()`
directly (`2a3a1b5d`).
- Removed: `deviceChannelViews()`, `proxyChannel::deallocate()`, the
transport header's `proxy_config.h`
include, a few comments, and the `ManyElementsOnePeerPut` test commit,
which nixlbench covers (#2328).
- The `num_threads` error now points to `proxy_thread_count`.
- For ColinNV's comments on #2328, on this PR's code: the peer check is
spelled out, the peer list is reserved,
and its copy size follows the element type (`54c54d8b`).
- Our own change: a `num_workers` below channels x peers, such as the
default of 1 from `getPluginParams()`, is
raised to that count instead of rejected, so applications no longer
erase it (#2328 drops its erases); a larger
one is still rejected, since nothing but the proxy threads would
progress the extra workers (`5b15101d`,
`2f49ebcf`).
Previous revisions:
- Rebased on main at `d1c6fb9e`, which includes #2423 (no garbage error
handler under
`ucx_error_handling_mode=none`) and #2412.
- For iyastreb's review: the proxy is a separate `nixlUcxProxyEngine`
that `create()` picks (`d72887b8`); the base
engine is main's code apart from that. The notification lock moved into
the proxy engine, whose `getNotifs`
no longer progresses the workers, which also removes `getNotifs` stalls
under host polling.
- For rakhmets's review: the proxy engine sets up the runtime in its
constructor and tears it down in its
destructor (`d72887b8`); the transport and runtime are ready once
constructed, and `proxyRuntime::create()`
returns a `unique_ptr` (`0efb4f4c`).
- For Colin's review on #2328: `ucx_error_handling_mode` defaults to
`none` with the proxy and `peer` is rejected
(`b240e1e6`), and `proxyTransport` is no longer copyable (`0efb4f4c`).
- Proxy threads up to one per (channel, peer) ring, even splits only
(`10fa8327`): each thread gets a span of
channels and a peer slice, so the loop keeps one call per channel.
- The base engine's memory-view comment now says that a mode fixed at
compile time is checked against the
handle's tag.
- Rebased on main at `d59f99c3` now that #2326 has merged; the PR shows
only its own 12 commits.
- Review fixes folded into their commits: no stray blank lines in the
engine commit (`d56d2688`), and
`nixlUcxEp::atomicAdd` is const (`b0e23d91`).
- `f1e597b5` picks a proxy ring by the element's remote agent instead of
its index in the view, and
`93eee926` puts to more elements of one peer than `proxy_max_peers`.
- Merged #2326's named status `mode` field (`57311ecc`); no code changes
here.
- Rebased on #2326's merge of main (`d52c5c36`, #2230's Device API
comments); no code changes.
- New commit `5b2a988b`, for CodeRabbit's review:
`discardUnpublishedMemView` is removed. When the engine fails to
allocate a
view's device wrapper, it rolls the proxy view back with the regular
`releaseMemView`, so the runtime API is
prepare and release only. The drain this adds runs only on that
allocation-failure path.
- New commit `53dcc156`, for CodeRabbit's review: the proxy helpers are
declared in the engine's existing
private section (with #2326's `dcf2f55a`, the engine's public interface
is unchanged from main), and the
transport's opcode switch has no `default`, so a new opcode triggers a
compiler warning.
- Rebased on #2326's compile-time mode argument (`dc0959c2`); new commit
`4878665d` adds the single-write case
with the mode fixed. No other changes.
- Rebased on #2326's review fixes (`8251aa55`). Following them, the
transport and engine commits no longer
use `HAVE_NIXL_DEVICE_API`: the plugin always builds the transport and
links the proxy library, and the
build-time rejection of proxy parameters without CUDA is gone (engine
creation fails instead). The engine
uses the `nixl::deviceMemView*` names.
- Stacked directly on #2326 instead of on the tests PR #2283. The two
test `meson.build` changes are made
against main, and `proxy_lifetime` and
`UnpublishedRollbackDoesNotDrainLiveTraffic` moved to #2283. No code
changes.
- Rebased on main at `d790ad4f`; no code changes.
- Restacked after the reorder that makes #2283 the post-feature tests
PR: this PR now sits on #2283 (which sits
on #2326), and `RetirementSelectsTheOwningDevice` moved to #2283.
Includes #2326's review fixes; no code
changes here.
- No longer carries #2325's commit.
- Rebased on #2283's test simplification; no code changes.
- Carries #2325's reworked single commit as a plain commit instead of
merging the old two-commit
branch; the EP release reordering is no longer part of it.
- The engine keeps the device allocator it built the runtime with, and
passes it to
`proxyRuntime::create` and to the device memview wrappers.
- `ProxyConfigTest` covers the new `proxy_drain_timeout_ms` parameter.
- Rebased on #2223's review fixes; no code changes.
- Rebased on main at `40b38570` after #2226 merged.
- `proxy_lifetime` now depends on `ucx_dep`, so it builds against the
configured UCX instead of
whatever UCX headers are on the default include path (an incomplete
system UCX broke the build).
- Follows #2226's review renames: the engine releases a view through
`proxyRuntime::releaseMemView`, and `discardUnpublishedMemView` takes
`proxy_view_handle_t`.
- The transport's submit debug log prints only the fields an operation
uses: an atomic
add has no source, so it logs its destination address and value.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added CUDA proxy support for GPU PUT and atomic-add operations, with
routing between proxy and direct UCX execution.
- Added proxy configuration and proxy-backed memory views.
- Added GPU utilities for coordinating work across warps, blocks, and
grids.
- **Bug Fixes**
- Improved remote metadata and memory-view lifetime handling: existing
views retain invalidated metadata, while new transfers using stale
metadata are rejected.
- Improved device endpoint disconnection by synchronizing CUDA and
releasing memory views before disconnecting peers.
- **Tests**
- Added coverage for proxy transfers, configuration, memory-view
lifetimes, completion handling, and remote metadata invalidation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Tomer Davidor <tdavidor@nvidia.com>
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What?
Set
UCP_EP_PARAM_FIELD_ERR_HANDLERinnixlUcxEponly when the error handler is filled in.Why?
Since #2365, under
UCP_ERR_HANDLING_MODE_NONEthe err handler field is still set infield_maskbuterr_handler.cb/argare left uninitialized. UCX then calls thatgarbage pointer when a peer dies, so the surviving process segfaults.
How?
Add the field to
field_masktogether with the handler, matching the "mask only whatis set" convention from #2017.
Summary by CodeRabbit