Skip to content

[TransferEngine] Delay RDMA redispatch while rails recover #3555

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

Closed
alogfans wants to merge 1 commit into from
Closed

[TransferEngine] Delay RDMA redispatch while rails recover #3555

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
21 changes: 21 additions & 0 deletions mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

#include <queue>
#include <unordered_set>
#include <vector>

#include "config.h"
#include "rdma_context.h"
Expand Down Expand Up @@ -58,6 +59,11 @@ class WorkerPool {

void redispatch(std::vector<Transport::Slice *> &slice_list, int thread_id,
bool handoff_to_local_worker = false);
void delayRedispatch(Transport::Slice *slice, bool handoff_to_local_worker,
const char *reason);
void releaseReadyDelayedSlices(int thread_id);
bool hasDelayedSlices() const;
uint64_t nextDelayedSliceDueNs() const;

void transferWorker(int thread_id);

Expand All @@ -78,6 +84,8 @@ class WorkerPool {
void markRailFailed(const std::string &peer_nic_path,
bool immediate_pause = false);
bool isRailAvailable(const std::string &peer_nic_path);
uint64_t peerRailPauseRemainingNs(const std::string &peer_nic_path,
uint64_t now);

// Retry helper: increment retry count and return whether retry is allowed
static bool shouldRetrySlice(Transport::Slice *slice);
Expand Down Expand Up @@ -165,6 +173,14 @@ class WorkerPool {
std::unordered_map<std::string, RailState> rail_states_;
std::mutex rail_state_lock_;

struct DelayedSlice {
Transport::Slice *slice = nullptr;
uint64_t due_ns = 0;
bool handoff_to_local_worker = false;
};
mutable std::mutex delayed_slices_lock_;
std::vector<DelayedSlice> delayed_slices_;

// Rail monitor configuration
const static int kRailErrorThreshold = 5; // Errors before pause
// Errors further apart than this are not consecutive, so error_count is
Expand All @@ -173,6 +189,11 @@ class WorkerPool {
const static uint64_t kRailErrorWindowNs = 5000000000ull; // 5 seconds
const static uint64_t kContextRecoveryDelayNs =
30000000000ull; // 30 seconds before a recovered local RNIC is reused
const static uint64_t kDelayedRedispatchMinBackoffNs =
100000000ull; // 100 ms
const static uint64_t kDelayedRedispatchMaxBackoffNs =
2000000000ull; // 2 seconds
const static uint64_t kDelayedRedispatchWakeCapNs = 100000000ull; // 100 ms

// Context-level health tracking
std::atomic<int> context_failure_count_{0};
Expand Down
194 changes: 179 additions & 15 deletions mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,11 @@

#include <sys/epoll.h>

#include <algorithm>
#include <cassert>
#include <chrono>
#include <functional>
#include <thread>

#include "config.h"
#include "memory_location.h"
Expand All @@ -33,6 +36,40 @@ namespace mooncake {

const static int kTransferWorkerCount = globalConfig().workers_per_ctx;

static uint64_t delayedRedispatchMaxWaitNs() {
const uint64_t pause_ns =
globalConfig().rdma_rail_pause_seconds * 1000000000ull;
return std::max<uint64_t>(120000000000ull, pause_ns * 4);
}

static uint64_t delayedRedispatchJitterNs(Transport::Slice *slice,
uint64_t window_ns) {
if (!slice || window_ns == 0) return 0;
auto value = reinterpret_cast<uintptr_t>(slice) >> 4;
value ^= static_cast<uintptr_t>(slice->target_id) * 0x9e3779b97f4a7c15ull;
return static_cast<uint64_t>(value % window_ns);
}

static uint64_t delayedRedispatchBackoffNs(Transport::Slice *slice,
uint64_t now) {
const static uint64_t kMinBackoffNs = 100000000ull; // 100 ms
const static uint64_t kMaxBackoffNs = 2000000000ull; // 2 seconds

uint64_t delay_ns = kMinBackoffNs;
if (slice && slice->ts > 0 && now > static_cast<uint64_t>(slice->ts)) {
const uint64_t elapsed_ns = now - static_cast<uint64_t>(slice->ts);
uint64_t scale = std::max<uint64_t>(1, elapsed_ns / kMinBackoffNs);
while (scale > 1 && delay_ns < kMaxBackoffNs) {
delay_ns <<= 1;
scale >>= 1;
}
delay_ns = std::min(delay_ns, kMaxBackoffNs);
}

const uint64_t jitter_window = std::max<uint64_t>(1, delay_ns / 4);
return delay_ns + delayedRedispatchJitterNs(slice, jitter_window);
}

static std::string resolveBufferLocation(
const TransferMetadata::BufferDesc &buffer, uint64_t offset) {
std::string location = buffer.name;
Expand Down Expand Up @@ -251,8 +288,10 @@ int WorkerPool::submitPostSend(
}
}
if (!found) {
slice->markFailed(); // All rails unavailable
slice->peer_nic_path = peer_nic_path;
delayRedispatch(slice, false, "all peer rails unavailable");
all_rails_failed_count++;
submitted_slice_count++;
continue;
}
}
Expand Down Expand Up @@ -310,6 +349,94 @@ void WorkerPool::enqueuePreparedSlices(SliceList (&slice_list_map)[kShardCount],
}
}

void WorkerPool::delayRedispatch(Transport::Slice *slice,
bool handoff_to_local_worker,
const char *reason) {
if (!slice) return;
const uint64_t now = static_cast<uint64_t>(getCurrentTimeInNano());
if (slice->ts <= 0) slice->ts = static_cast<int64_t>(now);
const uint64_t first_wait_ns = static_cast<uint64_t>(slice->ts);
if (now > first_wait_ns &&
now - first_wait_ns > delayedRedispatchMaxWaitNs()) {
Comment on lines +356 to +360
LOG(ERROR) << "Worker: Giving up delayed redispatch after "
<< (now - first_wait_ns) / 1000000ull
<< " ms, reason=" << reason
<< ", target=" << slice->target_id
<< ", peer_nic=" << slice->peer_nic_path
<< ", local_nic=" << context_.deviceName();
slice->markFailed();
processed_slice_count_++;
return;
}

uint64_t delay_ns = delayedRedispatchBackoffNs(slice, now);
const uint64_t rail_pause_remaining_ns =
peerRailPauseRemainingNs(slice->peer_nic_path, now);
if (rail_pause_remaining_ns > 0) {
const uint64_t pause_jitter_ns =
delayedRedispatchJitterNs(slice, kDelayedRedispatchMinBackoffNs);
delay_ns =
std::max(delay_ns, rail_pause_remaining_ns + pause_jitter_ns);
}

{
std::lock_guard<std::mutex> lock(delayed_slices_lock_);
delayed_slices_.push_back(
{slice, now + delay_ns, handoff_to_local_worker});
}
VLOG(1) << "Worker: delayed redispatch for target " << slice->target_id
<< ", reason=" << reason << ", local_nic=" << context_.deviceName()
<< ", peer_nic=" << slice->peer_nic_path
<< ", delay_ms=" << delay_ns / 1000000ull;
}

bool WorkerPool::hasDelayedSlices() const {
std::lock_guard<std::mutex> lock(delayed_slices_lock_);
return !delayed_slices_.empty();
}

uint64_t WorkerPool::nextDelayedSliceDueNs() const {
std::lock_guard<std::mutex> lock(delayed_slices_lock_);
uint64_t next_due_ns = 0;
for (const auto &entry : delayed_slices_) {
if (entry.due_ns == 0) continue;
if (next_due_ns == 0 || entry.due_ns < next_due_ns)
next_due_ns = entry.due_ns;
}
return next_due_ns;
}

void WorkerPool::releaseReadyDelayedSlices(int thread_id) {
const uint64_t now = static_cast<uint64_t>(getCurrentTimeInNano());
std::vector<DelayedSlice> ready;
{
std::lock_guard<std::mutex> lock(delayed_slices_lock_);
if (delayed_slices_.empty()) return;
std::vector<DelayedSlice> pending;
pending.reserve(delayed_slices_.size());
for (auto &entry : delayed_slices_) {
if (entry.due_ns <= now)
ready.push_back(entry);
else
pending.push_back(entry);
}
delayed_slices_.swap(pending);
}
if (ready.empty()) return;

SliceList remote_retry;
SliceList local_retry;
for (auto &entry : ready) {
if (!entry.slice) continue;
if (!context_.active() || entry.handoff_to_local_worker)
local_retry.push_back(entry.slice);
else
remote_retry.push_back(entry.slice);
}
if (!remote_retry.empty()) redispatch(remote_retry, thread_id, false);
if (!local_retry.empty()) redispatch(local_retry, thread_id, true);
}

int WorkerPool::submitPreparedPostSend(
const std::vector<Transport::Slice *> &slice_list) {
// Called by a different local RNIC's worker during local failover. The
Expand Down Expand Up @@ -442,7 +569,10 @@ void WorkerPool::performPostSend(int thread_id) {
// deadline; only a genuine connection/QP failure can arm or extend it.
if (!isRailAvailable(entry.first) ||
context_.isConnectPaused(entry.first)) {
for (auto &slice : entry.second) failed_slice_list.push_back(slice);
for (auto &slice : entry.second) {
slice->peer_nic_path = entry.first;
delayRedispatch(slice, false, "peer rail is paused");
}
entry.second.clear();
continue;
}
Expand All @@ -454,7 +584,10 @@ void WorkerPool::performPostSend(int thread_id) {
auto endpoint = context_.endpoint(entry.first);
#endif
if (!endpoint) {
for (auto &slice : entry.second) failed_slice_list.push_back(slice);
for (auto &slice : entry.second) {
slice->peer_nic_path = entry.first;
delayRedispatch(slice, false, "endpoint unavailable");
}
entry.second.clear();
continue;
}
Expand Down Expand Up @@ -497,6 +630,10 @@ void WorkerPool::performPostSend(int thread_id) {
if (!has_peer_alternative && local_context_inactive &&
tryHandoffToAnotherLocalWorker(slice)) {
processed_slice_count_++;
} else if (!has_peer_alternative &&
!local_context_inactive) {
delayRedispatch(slice, false,
"endpoint setup failed transiently");
} else {
failed_slice_list.push_back(slice);
}
Expand All @@ -513,8 +650,10 @@ void WorkerPool::performPostSend(int thread_id) {
markRailFailed(entry.first, true);
redispatch_counter_++;
context_.deleteEndpointByPtr(endpoint.get());
for (auto &slice : entry.second)
failed_slice_list.push_back(slice);
for (auto &slice : entry.second) {
delayRedispatch(slice, false,
"timed out waiting for ready ACK");
}
entry.second.clear();
}
continue;
Expand Down Expand Up @@ -768,9 +907,9 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
}
// A local RNIC failure cannot be repaired by keeping this
// worker/context and changing the remote rail. If no other
// local worker can take the slice, fail it immediately.
slice->markFailed();
processed_slice_count_++;
// local worker can take the slice yet, wait for rail recovery
// before deciding every local path is gone.
delayRedispatch(slice, true, "no alternate local RNIC");
continue;
}

Expand All @@ -788,8 +927,8 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
<< "dest_addr=" << (void *)slice->rdma.dest_addr
<< ", length=" << slice->length
<< ", retry_cnt=" << slice->rdma.retry_cnt;
slice->markFailed();
processed_slice_count_++;
delayRedispatch(slice, false,
"peer segment unavailable or no target RNIC");
continue;
}
slice->rdma.dest_rkey =
Expand Down Expand Up @@ -822,13 +961,13 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
}
if (!found) {
LOG(ERROR)
<< "Worker: Cannot redispatch slice because all peer "
<< "Worker: Delaying redispatch because all peer "
"rails are paused for target "
<< slice->target_id
<< ", selected peer=" << peer_nic_path
<< ", retry_cnt=" << slice->rdma.retry_cnt;
slice->markFailed();
processed_slice_count_++;
slice->peer_nic_path = peer_nic_path;
delayRedispatch(slice, false, "all peer rails paused");
continue;
}
}
Expand All @@ -850,7 +989,6 @@ void WorkerPool::redispatch(std::vector<Transport::Slice *> &slice_list,
<< ", length=" << slice->length
<< ", retry_cnt=" << slice->rdma.retry_cnt;
}
slice->ts = 0;
if (use_local_queue) {
collective_slice_queue_[thread_id][peer_nic_path].push_back(
slice);
Expand Down Expand Up @@ -929,7 +1067,6 @@ bool WorkerPool::tryHandoffToAnotherLocalWorker(Transport::Slice *slice) {
slice->rdma.source_lkey =
local_segment_desc->buffers[buffer_id].lkey[device_id];
slice->rdma.endpoint = nullptr;
slice->ts = 0;

std::vector<Transport::Slice *> handoff{slice};
alt_ctx->worker_pool_->submitPreparedPostSend(handoff);
Expand Down Expand Up @@ -961,6 +1098,7 @@ void WorkerPool::transferWorker(int thread_id) {
const bool can_post = workerCanPost(thread_id);
const bool can_poll = workerCanPoll(thread_id);
while (workers_running_.load(std::memory_order_relaxed)) {
releaseReadyDelayedSlices(thread_id);
auto processed_slice_count =
processed_slice_count_.load(std::memory_order_relaxed);
auto submitted_slice_count =
Expand Down Expand Up @@ -992,6 +1130,15 @@ void WorkerPool::transferWorker(int thread_id) {
performPollCq(thread_id);
}
#endif
const uint64_t next_delayed_due_ns = nextDelayedSliceDueNs();
if (next_delayed_due_ns > 0) {
const uint64_t now = static_cast<uint64_t>(getCurrentTimeInNano());
if (next_delayed_due_ns > now) {
const uint64_t sleep_ns = std::min(next_delayed_due_ns - now,
kDelayedRedispatchWakeCapNs);
std::this_thread::sleep_for(std::chrono::nanoseconds(sleep_ns));
}
}
Comment on lines +1133 to +1141
last_wait_ts = getCurrentTimeInNano();
}
}
Expand Down Expand Up @@ -1414,6 +1561,23 @@ bool WorkerPool::isRailAvailable(const std::string &peer_nic_path) {
return false;
}

uint64_t WorkerPool::peerRailPauseRemainingNs(const std::string &peer_nic_path,
uint64_t now) {
if (peer_nic_path.empty()) return 0;

std::lock_guard<std::mutex> lock(rail_state_lock_);
auto it = rail_states_.find(peer_nic_path);
if (it == rail_states_.end()) return 0;
auto &state = it->second;
if (state.pause_until_ns == 0) return 0;
if (now >= state.pause_until_ns) {
state.error_count = 0;
state.pause_until_ns = 0;
return 0;
}
return state.pause_until_ns - now;
}

// Unified retry logic: increment retry count and return whether retry is
// allowed
bool WorkerPool::shouldRetrySlice(Transport::Slice *slice) {
Expand Down
Loading