Skip to content

libfabric: PT-owns-endpoint with MPSC lock-free ring #1949

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
oa-aws wants to merge 1 commit into
base: main
Choose a base branch
from
+285 −28
Open

libfabric: PT-owns-endpoint with MPSC lock-free ring #1949

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 23 additions & 8 deletions src/plugins/libfabric/libfabric_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,16 +259,28 @@ nixlLibfabricCudaCtx::cudaUpdateCtxPtr(void *address, int expected_dev, bool &wa

was_updated = false;

if (expected_dev == -1) return -1;
if (myDevId_ != -1 && expected_dev != myDevId_) return -1;
if (expected_dev == -1) {
return -1;
}
if (myDevId_ != -1 && expected_dev != myDevId_) {
return -1;
}

ret = cudaQueryAddr(address, is_dev, dev, ctx, pci_bus_id);
if (ret) return ret;
if (!is_dev) return 0;
if (dev != expected_dev) return -1;
if (ret) {
return ret;
}
if (!is_dev) {
return 0;
}
if (dev != expected_dev) {
return -1;
}

if (pthrCudaCtx_) {
if (pthrCudaCtx_ != ctx) return -1;
if (pthrCudaCtx_ != ctx) {
return -1;
}
return 0;
}

Expand All @@ -282,7 +294,9 @@ nixlLibfabricCudaCtx::cudaUpdateCtxPtr(void *address, int expected_dev, bool &wa
int
nixlLibfabricCudaCtx::cudaSetCtx() {
CUresult result;
if (NULL == pthrCudaCtx_) return 0;
if (NULL == pthrCudaCtx_) {
return 0;
}

result = cuCtxSetCurrent(pthrCudaCtx_);
return (CUDA_SUCCESS == result);
Expand Down Expand Up @@ -1183,7 +1197,8 @@ nixlLibfabricEngine::postXferDescriptors(nixlLibfabricReq::OpType op_type,
desc_submitted_count,
desc_idx,
desc_count,
xfer_base_offset);
xfer_base_offset,
local[desc_idx].devId);

if (status != NIXL_SUCCESS) {
NIXL_ERROR << "prepareAndSubmitTransfer failed for descriptor " << desc_idx
Expand Down
139 changes: 128 additions & 11 deletions src/utils/libfabric/libfabric_rail.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@
#include "serdes/serdes.h"
#include "libfabric_common.h"

#ifdef HAVE_CUDA
#include <cuda_runtime.h>

@fengjica fengjica Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this included here because some new cuda runtimes functions are added?

@oa-aws oa-aws Jul 16, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just took it form the original PR. I think it was in preparation for setting cuda context per deferred request in the progress thread (whereas now this is done in nixlLibfabricEngine::postXferDescriptors before calling rail_manager_.prepareAndSubmitTransfer, which assumes request execution in the current thread). This change is also required (i.e. when PT=1 do not set CUDA context in the engine, but rather let the PT do that).

#endif

#include <cstring>
#include <stdexcept>
#include <stack>
Expand Down Expand Up @@ -168,7 +172,9 @@ ControlRequestPool::~ControlRequestPool() {

void
ControlRequestPool::cleanup() {
if (buffer_chunks_.empty()) return; // Already cleaned up
if (buffer_chunks_.empty()) {
return; // Already cleaned up
}

for (auto &chunk : buffer_chunks_) {
if (chunk.mr) {
Expand Down Expand Up @@ -712,8 +718,8 @@ nixlLibfabricRail::setXferIdCallback(std::function<void(uint32_t)> callback) {

// Per-rail completion processing - handles one rail's CQ with configurable blocking behavior
nixl_status_t
nixlLibfabricRail::progressCompletionQueue() const {
// Completion processing
nixlLibfabricRail::progressCompletionQueue() {
// Completion processing FIRST — process arrivals before posting new items
struct fi_cq_data_entry completions[NIXL_LIBFABRIC_CQ_BATCH_SIZE];

int ret;
Expand Down Expand Up @@ -747,7 +753,8 @@ nixlLibfabricRail::progressCompletionQueue() const {
// CQ lock released here - completion is now local data

if (ret == -FI_EAGAIN) {
return NIXL_IN_PROG; // No completions available
// SQ full - do inline CQ progress to free slots, then retry
// No completions - fall through to drainPostQueue
}

if (ret > 0) {
Expand All @@ -762,10 +769,14 @@ nixlLibfabricRail::progressCompletionQueue() const {
}

NIXL_DEBUG << "Processed " << ret << " completions on rail " << rail_id;
return NIXL_SUCCESS;
}

return NIXL_ERR_BACKEND; // Unexpected case
// PT-owns-endpoint: drain pending posts AFTER processing completions
if (progress_thread_enabled_) {
drainPostQueue();
}

return (ret > 0) ? NIXL_SUCCESS : NIXL_IN_PROG;
}

// Route completion to appropriate handler (rail-specific)
Expand Down Expand Up @@ -1018,9 +1029,7 @@ nixlLibfabricRail::postRecv(nixlLibfabricReq *req) const {
}

nixl_status_t
nixlLibfabricRail::postSend(uint64_t immediate_data,
fi_addr_t dest_addr,
nixlLibfabricReq *req) const {
nixlLibfabricRail::postSend(uint64_t immediate_data, fi_addr_t dest_addr, nixlLibfabricReq *req) {
if (req->buffer_size == 0 || req->buffer_size > NIXL_LIBFABRIC_SEND_RECV_BUFFER_SIZE) {
NIXL_ERROR << "Invalid message size=" << req->buffer_size
<< " (max: " << NIXL_LIBFABRIC_SEND_RECV_BUFFER_SIZE << ")";
Expand Down Expand Up @@ -1063,6 +1072,7 @@ nixlLibfabricRail::postSend(uint64_t immediate_data,
}

if (ret == -FI_EAGAIN) {
// SQ full - do inline CQ progress to free slots, then retry
Comment thread
oa-aws marked this conversation as resolved.
// Resource temporarily unavailable - retry indefinitely for all providers
attempt++;

Expand Down Expand Up @@ -1099,6 +1109,111 @@ nixlLibfabricRail::postSend(uint64_t immediate_data,
return NIXL_ERR_BACKEND;
}

/****************************************
* PT-owns-endpoint: post queue methods
****************************************/

void
nixlLibfabricRail::enqueuePost(const nixlLibfabricPostRequest &req) {
while (!post_ring_.push(req)) {
std::this_thread::yield();

@fengjica fengjica Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it better to yield and retry, or abort to let let users know that something is full (and let user retry)?

@oa-aws oa-aws Jul 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I changed the implementation due to code rabbit comments, but the question remains the same: what should we do in case of a full ring buffer? I think the best option for now is to retry a few times until some request timeout is reached and let user know we failed, so the ring buffer size can be configured to larger value for the workload in question.

}
}

nixl_status_t
nixlLibfabricRail::drainPostQueue() {
nixlLibfabricPostRequest pr;
int posted = 0;
[[maybe_unused]] int current_cuda_device = -1;

@fengjica fengjica Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not used anywhere?

@oa-aws oa-aws Jul 16, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It should be used, see comments above.


while (post_ring_.pop(pr)) {

@fengjica fengjica Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This while loop will not break until the post_ring is empty. With one PT, it's possible that it's stuck here, and do not get a chance to progress other rails, right?

@oa-aws oa-aws Jul 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct. I was worried about that too. This is how it was done in the original PR, and we need to think about a solution. See previous comment. Currently adding some compile time constant as a batch limit.

int ret = -FI_EAGAIN;

while (ret == -FI_EAGAIN) {
if (pr.type == nixlLibfabricPostRequest::WRITE) {
Comment on lines +1127 to +1133

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
rg -n -C3 'device_id|current_cuda_device|cudaSetDevice' \
  src/utils/libfabric/libfabric_rail.cpp \
  src/utils/libfabric/libfabric_rail_manager.cpp \
  src/plugins/libfabric/libfabric_backend.cpp

Repository: ai-dynamo/nixl

Length of output: 14213

🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '1124,1215p' src/utils/libfabric/libfabric_rail.cpp
printf '\n====\n'
sed -n '1215,1295p' src/utils/libfabric/libfabric_rail.cpp
printf '\n====\n'
sed -n '460,490p' src/utils/libfabric/libfabric_rail_manager.cpp
printf '\n====\n'
sed -n '1128,1178p' src/plugins/libfabric/libfabric_backend.cpp

Repository: ai-dynamo/nixl

Length of output: 10441

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C3 'struct nixlLibfabricPostRequest|device_id|current_cuda_device|cudaSetDevice|enqueuePost|drainPostQueue' src/utils/libfabric src/plugins/libfabric

Repository: ai-dynamo/nixl

Length of output: 36704

Switch to the queued request’s CUDA device before posting it.

pr.device_id is copied into the queue, but drainPostQueue() never applies it. The direct post path already calls cudaSetDevice() for VRAM transfers, so the progress thread can post on the wrong GPU unless it switches first.

🤖 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/utils/libfabric/libfabric_rail.cpp` around lines 1127 - 1133, Update
drainPostQueue() to switch to pr.device_id before posting each queued VRAM
request, reusing the existing CUDA device-tracking logic and only calling
cudaSetDevice when the queued device differs from the current device. Preserve
the existing behavior for non-VRAM requests and error handling.

struct iovec iov{};
iov.iov_base = pr.local_addr;
iov.iov_len = pr.length;

struct fi_rma_iov rma_iov{};
rma_iov.addr = pr.remote_addr;
rma_iov.len = pr.length;
rma_iov.key = pr.remote_key;

void *desc = pr.local_desc;
struct fi_msg_rma msg = {};
msg.msg_iov = &iov;
msg.desc = &desc;
msg.iov_count = 1;
msg.addr = pr.dest_addr;
msg.rma_iov = &rma_iov;
msg.rma_iov_count = 1;
msg.context = &pr.req->ctx;
msg.data = pr.immediate_data;

ret = fi_writemsg(endpoint, &msg, pr.fi_flags | FI_REMOTE_CQ_DATA);
} else {
struct iovec iov{};
iov.iov_base = pr.local_addr;
iov.iov_len = pr.length;

struct fi_rma_iov rma_iov{};
rma_iov.addr = pr.remote_addr;
rma_iov.len = pr.length;
rma_iov.key = pr.remote_key;

void *desc = pr.local_desc;
struct fi_msg_rma msg = {};
msg.msg_iov = &iov;
msg.desc = &desc;
msg.iov_count = 1;
msg.addr = pr.dest_addr;
msg.rma_iov = &rma_iov;
msg.rma_iov_count = 1;
msg.context = &pr.req->ctx;

ret = fi_readmsg(endpoint, &msg, pr.fi_flags);
}

if (ret == -FI_EAGAIN) {
struct fi_cq_data_entry cq_buf[16];
int cq_ret = fi_cq_read(cq, cq_buf, 16);
if (cq_ret > 0) {
for (int c = 0; c < cq_ret; c++) {
processCompletionQueueEntry(&cq_buf[c]);

@fengjica fengjica Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we defer processing them so that post can be drained faster? Just a thought, not necessary for this PR.

@oa-aws oa-aws Jul 16, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This eventually increments an atomic integer. Unless there is other heavy computation, I don't think we can gain much. The question is, is there a use case where a completion callback would be heavy (i.e. someone developing on top of the libfabric plugin)? If so, then this should have a design, because I don't think it is trivial (e.g. does someone rely on counters to catch up to tell request is finished? If so, would deferring cause unwanted delays? Maybe user should declare whether the callback should be deferred? etc.).

@fengjica fengjica Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right.

Upon checking the callbacks today in the code, I also notice a bug. Because I was thinking a notif needs to be sent out after all completions are done, for READ. I only found one place in the code in checkXfer() for that. It means, if the user (initiator) does not poll status through checkXfer() into our code, we do not send the notif to the target. https://github.com/ai-dynamo/nixl/blob/main/src/plugins/libfabric/libfabric_backend.cpp#L1452-L1454. I was expecting when PT=1, this notif should be sent to target earlier. If you also think this is the expected behavior, I'll open a bug.

No need to address it in this PR, too.

}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}
}

if (ret != 0) {
NIXL_ERROR << "drainPostQueue: post failed ret=" << ret << " (" << fi_strerror(-ret)
<< ") on rail " << rail_id;
if (pr.req && pr.req->completion_callback) {
pr.req->completion_callback();
}
continue;
Comment on lines +1189 to +1195

@coderabbitai coderabbitai Bot Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift

Do not report failed posts as successful completions.

The backend callback only increments completed_requests, so invoking it here makes an operation that was never posted appear successful. The request is also never returned to its pool. Record an asynchronous failure, release pr.req, then notify completion through a failure-aware contract.

🤖 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/utils/libfabric/libfabric_rail.cpp` around lines 1189 - 1195, Update the
ret != 0 failure branch in drainPostQueue so failed posts are not reported
through the success-only completion_callback. Record the asynchronous failure,
release pr.req back to its request pool, and invoke the request’s failure-aware
completion notification contract while preserving the existing error log and
continue behavior.

@oa-aws oa-aws Jul 16, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is incorrect, otherwise the call to check whether a multi-part/segment request finished will never declare success, as the counters will never match up if a completion is missing, even if failed. In addition to that, maybe a fix should be added to tell the total request status, in case some sub-segments posting failed.

}
posted++;
if (posted % 32 == 0) {

@fengjica fengjica Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess this 32 is something that can be tuned later to find an optimal ratio of draining vs reading cq, is that the case?

@oa-aws oa-aws Jul 16, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think so. This is coming from the original PR, and I was wondering when to benchmark fine-tuning of it. We should also (as you mentioned above) enforce some batch-size limit so that post requests from other rails are executed in a timely manner. This budgeting is not trivial though, so I suggest for this PR we put some constants that work reasonably, and for next PRs we improve on that based on perf test results.

struct fi_cq_data_entry cq_buf[16];
int cq_ret = fi_cq_read(cq, cq_buf, 16);
if (cq_ret > 0) {
for (int c = 0; c < cq_ret; c++) {
processCompletionQueueEntry(&cq_buf[c]);
}
} else if (cq_ret < 0 && cq_ret != -FI_EAGAIN) {
struct fi_cq_err_entry err_entry = {};
fi_cq_readerr(cq, &err_entry, 0);
NIXL_ERROR << "CQ error in drain interleave on rail " << rail_id << ": "
<< fi_strerror(err_entry.err);
}
}
}

return NIXL_SUCCESS;
}

nixl_status_t
nixlLibfabricRail::postWrite(const void *local_buffer,
size_t length,
Expand All @@ -1108,7 +1223,7 @@ nixlLibfabricRail::postWrite(const void *local_buffer,
uint64_t remote_addr,
uint64_t remote_key,
nixlLibfabricReq *req,
uint64_t fi_flags) const {
uint64_t fi_flags) {
// Validation
if (!req) {
NIXL_ERROR << "Invalid request for write on rail " << rail_id;
Expand Down Expand Up @@ -1160,6 +1275,7 @@ nixlLibfabricRail::postWrite(const void *local_buffer,
}

if (ret == -FI_EAGAIN) {
// SQ full - do inline CQ progress to free slots, then retry
// Resource temporarily unavailable - retry indefinitely for all providers
attempt++;

Expand Down Expand Up @@ -1203,7 +1319,7 @@ nixlLibfabricRail::postRead(void *local_buffer,
fi_addr_t dest_addr,
uint64_t remote_addr,
uint64_t remote_key,
nixlLibfabricReq *req) const {
nixlLibfabricReq *req) {
// Validation
if (!req) {
NIXL_ERROR << "Invalid request for read on rail " << rail_id;
Expand Down Expand Up @@ -1242,6 +1358,7 @@ nixlLibfabricRail::postRead(void *local_buffer,
}

if (ret == -FI_EAGAIN) {
// SQ full - do inline CQ progress to free slots, then retry
// Resource temporarily unavailable - retry indefinitely for all providers
attempt++;

Expand Down
Loading
Loading