Nixlbench: register remote IOVs only for storage backends - #1772
Merged
aranadive merged 1 commit intoJun 15, 2026
Merged
Conversation
iyastreb
requested review from
a team,
aranadive,
brminich and
ovidiusm
as code owners
June 15, 2026 12:49
|
👋 Hi iyastreb! Thank you for contributing to ai-dynamo/nixl. Your PR reviewers will review your contribution then trigger the CI to test your changes. 🚀 |
Contributor
📝 WalkthroughWalkthroughIn ChangesConditional Remote Memory Registration
Estimated code review effort🎯 2 (Simple) | ⏱️ ~5 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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)
benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp (1)
1347-1358:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFix partial registration lifecycle to avoid leaked/inconsistent agent state on errors.
When storage remote registration fails after local registration succeeds, the function returns without rolling back local registration. Similarly, deregistration does local first, so a remote failure can leave remote state behind. This creates inconsistent state during
--reregister_memfailure paths.Suggested fix
static nixl_status_t registerIterationMem(nixlAgent *agent, const std::vector<xferBenchIOV> &local_iov, const std::vector<xferBenchIOV> &remote_iov, nixlBackendH *backend_engine) { nixl_opt_args_t reg_args; reg_args.backends.push_back(backend_engine); nixl_reg_dlist_t local_reg = iovListToNixlRegDlist(local_iov, GET_SEG_TYPE(true)); nixl_status_t rc = agent->registerMem(local_reg, ®_args); if (rc != NIXL_SUCCESS) { return rc; } if (xferBenchConfig::isStorageBackend()) { nixl_reg_dlist_t remote_reg = iovListToNixlRegDlist(remote_iov, getRemoteSegType()); rc = agent->registerMem(remote_reg, ®_args); if (rc != NIXL_SUCCESS) { + // Best-effort rollback to keep registration state consistent. + (void)agent->deregisterMem(local_reg, ®_args); return rc; } } return NIXL_SUCCESS; } static nixl_status_t deregisterIterationMem(nixlAgent *agent, const std::vector<xferBenchIOV> &local_iov, const std::vector<xferBenchIOV> &remote_iov, nixlBackendH *backend_engine) { nixl_opt_args_t reg_args; reg_args.backends.push_back(backend_engine); - nixl_reg_dlist_t local_reg = iovListToNixlRegDlist(local_iov, GET_SEG_TYPE(true)); - nixl_status_t rc = agent->deregisterMem(local_reg, ®_args); - if (rc != NIXL_SUCCESS) { - return rc; - } + nixl_status_t rc = NIXL_SUCCESS; if (xferBenchConfig::isStorageBackend()) { nixl_reg_dlist_t remote_reg = iovListToNixlRegDlist(remote_iov, getRemoteSegType()); rc = agent->deregisterMem(remote_reg, ®_args); if (rc != NIXL_SUCCESS) { return rc; } } + + nixl_reg_dlist_t local_reg = iovListToNixlRegDlist(local_iov, GET_SEG_TYPE(true)); + rc = agent->deregisterMem(local_reg, ®_args); + if (rc != NIXL_SUCCESS) { + return rc; + } return NIXL_SUCCESS; }Also applies to: 1373-1384
🤖 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 `@benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp` around lines 1347 - 1358, The code has a partial registration lifecycle issue where if local memory registration succeeds but remote registration fails (in the storage backend path), the function returns without rolling back the local registration, leaving inconsistent agent state. At benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp lines 1347-1358 (anchor), when the remote registerMem call for remote_reg fails after local registration succeeds, add a deregistration call to rollback the local_reg before returning the error code. The same fix pattern must be applied at lines 1373-1384 (siblings) in the deregistration path, where if the remote deregistration fails after local deregistration succeeds, you must handle the error appropriately to prevent leaving remote state behind. Ensure error handling maintains consistent agent state across all failure paths.
🤖 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 `@benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp`:
- Around line 1347-1358: The code has a partial registration lifecycle issue
where if local memory registration succeeds but remote registration fails (in
the storage backend path), the function returns without rolling back the local
registration, leaving inconsistent agent state. At
benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp lines 1347-1358 (anchor),
when the remote registerMem call for remote_reg fails after local registration
succeeds, add a deregistration call to rollback the local_reg before returning
the error code. The same fix pattern must be applied at lines 1373-1384
(siblings) in the deregistration path, where if the remote deregistration fails
after local deregistration succeeds, you must handle the error appropriately to
prevent leaving remote state behind. Ensure error handling maintains consistent
agent state across all failure paths.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: ef45e777-32c5-4b43-930b-9f22aae63065
📒 Files selected for processing (1)
benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp
Contributor
Author
|
/build |
aranadive
approved these changes
Jun 15, 2026
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
2 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?
Bugfix https://nvbugspro.nvidia.com/bug/6310482
Fix
--reregister_memfailing for memory/network backends (e.g. UCX) with"VRAM is detected as host by UCX".
Why?
registerIterationMem()/deregisterIterationMem()register the remote IOVslocally on every iteration, unconditionally. For memory backends the remote side
is a peer process's memory (exchanged via metadata), not a local allocation, so
registering those peer addresses locally is invalid — UCX can't resolve them and
reports VRAM as host. The normal alloc/dealloc path only registers remote memory
for storage backends; the per-iteration path missed that gate. Regression from
#1474.
How?
Gate the remote register/deregister on
isStorageBackend(), matchingallocateMemory()/deallocateMemory(). Memory backends now (de)register onlylocal memory per iteration, as before.
Summary by CodeRabbit