1072 lines
41 KiB
C++
1072 lines
41 KiB
C++
// Copyright 2024 KVCache.AI
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "transfer_engine_py.h"
|
|
|
|
#include <cassert>
|
|
#include <numeric>
|
|
#include <fstream>
|
|
|
|
#include <pybind11/stl.h>
|
|
#include "transport/rpc_communicator/rpc_interface.h"
|
|
|
|
#ifdef USE_MNNVL
|
|
#include "transport/nvlink_transport/nvlink_transport.h"
|
|
#endif
|
|
|
|
#ifdef USE_INTRA_NVLINK
|
|
#include "transport/intranode_nvlink_transport/intranode_nvlink_transport.h"
|
|
#endif
|
|
|
|
#ifdef USE_CUDA
|
|
#include <cuda_runtime.h>
|
|
#endif
|
|
|
|
static void *(*allocateMemory)(size_t) = nullptr;
|
|
static void (*freeMemory)(void *) = nullptr;
|
|
static std::string g_protocol;
|
|
|
|
// Handle allocateMemory function pointer based on protocol
|
|
void initMemoryAllocator(const char *protocol) {
|
|
if (allocateMemory != nullptr) {
|
|
LOG(WARNING) << "Memory allocator already initialized with: "
|
|
<< g_protocol;
|
|
return;
|
|
}
|
|
g_protocol = protocol;
|
|
if (strcmp(protocol, "nvlink") == 0) {
|
|
#ifdef USE_MNNVL
|
|
allocateMemory = [](size_t s) -> void * {
|
|
return mooncake::NvlinkTransport::allocatePinnedLocalMemory(s);
|
|
};
|
|
freeMemory = [](void *p) {
|
|
mooncake::NvlinkTransport::freePinnedLocalMemory(p);
|
|
};
|
|
LOG(INFO) << "Selected MNNVL (NVLink) memory allocator";
|
|
#else
|
|
LOG(ERROR) << "Protocol 'nvlink' requires -DUSE_MNNVL=ON";
|
|
#endif
|
|
} else if (strcmp(protocol, "nvlink_intra") == 0) {
|
|
#ifdef USE_INTRA_NVLINK
|
|
allocateMemory = [](size_t s) -> void * {
|
|
return mooncake::IntraNodeNvlinkTransport::
|
|
allocatePinnedLocalMemory(s);
|
|
};
|
|
freeMemory = [](void *p) {
|
|
mooncake::IntraNodeNvlinkTransport::freePinnedLocalMemory(p);
|
|
};
|
|
LOG(INFO) << "Selected Intra-NVLink memory allocator";
|
|
#else
|
|
LOG(ERROR) << "Protocol 'nvlink_intra' requires -DUSE_INTRA_NVLINK=ON";
|
|
#endif
|
|
} else {
|
|
// default fallback
|
|
allocateMemory = malloc;
|
|
freeMemory = free;
|
|
LOG(WARNING) << "Using default malloc/free for protocol: " << protocol;
|
|
}
|
|
}
|
|
|
|
TransferEnginePy::TransferEnginePy() {
|
|
const int64_t kNanosPerSecond = 1000 * 1000 * 1000;
|
|
if (getenv("MC_TRANSFER_TIMEOUT")) {
|
|
int timeout_sec = std::max(5, atoi(getenv("MC_TRANSFER_TIMEOUT")));
|
|
transfer_timeout_nsec_ = timeout_sec * kNanosPerSecond;
|
|
} else {
|
|
transfer_timeout_nsec_ = 30 * kNanosPerSecond;
|
|
}
|
|
}
|
|
|
|
TransferEnginePy::~TransferEnginePy() {
|
|
for (auto &handle : handle_map_) engine_->closeSegment(handle.second);
|
|
handle_map_.clear();
|
|
engine_.reset();
|
|
for (auto &buffer : buffer_list_) freeMemory(buffer);
|
|
buffer_list_.clear();
|
|
for (auto &buffer : large_buffer_list_) freeMemory(buffer);
|
|
large_buffer_list_.clear();
|
|
}
|
|
|
|
std::vector<std::string> buildDeviceFilter(const std::string &device_names) {
|
|
std::stringstream ss(device_names);
|
|
std::string item;
|
|
std::vector<std::string> tokens;
|
|
while (getline(ss, item, ',')) {
|
|
tokens.push_back(item);
|
|
}
|
|
return tokens;
|
|
}
|
|
|
|
std::pair<std::string, std::string> parseConnectionString(
|
|
const std::string &conn_string) {
|
|
std::pair<std::string, std::string> result;
|
|
std::string proto = "etcd";
|
|
std::string domain;
|
|
std::size_t pos = conn_string.find("://");
|
|
|
|
if (pos != std::string::npos) {
|
|
proto = conn_string.substr(0, pos);
|
|
domain = conn_string.substr(pos + 3);
|
|
} else if (conn_string == P2PHANDSHAKE) {
|
|
proto = "";
|
|
domain = P2PHANDSHAKE;
|
|
} else {
|
|
domain = conn_string;
|
|
}
|
|
|
|
result.first = proto;
|
|
result.second = domain;
|
|
return result;
|
|
}
|
|
|
|
std::string buildConnString(const std::string &metadata_type,
|
|
const std::string &metadata_server) {
|
|
if (metadata_server == P2PHANDSHAKE) {
|
|
return P2PHANDSHAKE;
|
|
}
|
|
|
|
std::string conn_string = metadata_server;
|
|
if (conn_string.find("://") == std::string::npos)
|
|
conn_string = metadata_type + "://" + metadata_server;
|
|
return conn_string;
|
|
}
|
|
|
|
int TransferEnginePy::initialize(const char *local_hostname,
|
|
const char *metadata_server,
|
|
const char *protocol,
|
|
const char *device_name) {
|
|
initMemoryAllocator(protocol);
|
|
|
|
auto conn_string = parseConnectionString(metadata_server);
|
|
return initializeExt(local_hostname, conn_string.second.c_str(), protocol,
|
|
device_name, conn_string.first.c_str());
|
|
}
|
|
|
|
int TransferEnginePy::initializeExt(const char *local_hostname,
|
|
const char *metadata_server,
|
|
const char *protocol,
|
|
const char *device_name,
|
|
const char *metadata_type) {
|
|
std::string proto = protocol ? std::string(protocol) : "";
|
|
std::string conn_string = buildConnString(metadata_type, metadata_server);
|
|
|
|
auto device_name_safe = device_name ? std::string(device_name) : "";
|
|
auto device_filter = buildDeviceFilter(device_name_safe);
|
|
|
|
#ifdef USE_EFA
|
|
// When using EFA protocol, we still need topology discovery but won't
|
|
// auto-install RDMA
|
|
bool use_efa = (proto == "efa");
|
|
// Disable auto_discover to prevent RDMA transport installation, we'll
|
|
// install EFA manually
|
|
engine_ = std::make_unique<TransferEngine>(false, device_filter);
|
|
// Manually discover topology for EFA to populate device list
|
|
if (use_efa) {
|
|
engine_->getLocalTopology()->discover(device_filter);
|
|
LOG(INFO) << "Topology discovery complete for EFA. Found "
|
|
<< engine_->getLocalTopology()->getHcaList().size()
|
|
<< " devices.";
|
|
}
|
|
#else
|
|
engine_ = std::make_unique<TransferEngine>(true, device_filter);
|
|
#endif
|
|
|
|
if (getenv("MC_LEGACY_RPC_PORT_BINDING")) {
|
|
auto hostname_port = parseHostNameWithPort(local_hostname);
|
|
int ret =
|
|
engine_->init(conn_string, local_hostname,
|
|
hostname_port.first.c_str(), hostname_port.second);
|
|
if (ret) return -1;
|
|
} else {
|
|
// the last two params are unused
|
|
int ret = engine_->init(conn_string, local_hostname, "", 0);
|
|
if (ret) return -1;
|
|
}
|
|
|
|
#ifdef USE_EFA
|
|
// Install EFA transport when protocol is "efa"
|
|
if (use_efa) {
|
|
LOG(INFO)
|
|
<< "Installing EFA transport as requested by protocol parameter";
|
|
auto transport = engine_->installTransport("efa", nullptr);
|
|
if (!transport) {
|
|
LOG(ERROR) << "Failed to install EFA transport";
|
|
return -1;
|
|
}
|
|
LOG(INFO) << "EFA transport installed successfully";
|
|
}
|
|
#endif
|
|
|
|
free_list_.resize(kSlabSizeKBTabLen);
|
|
return 0;
|
|
}
|
|
|
|
int TransferEnginePy::getRpcPort() { return engine_->getRpcPort(); }
|
|
|
|
char *TransferEnginePy::allocateRawBuffer(size_t capacity) {
|
|
auto buffer = allocateMemory(capacity);
|
|
if (!buffer) return nullptr;
|
|
int ret = engine_->registerLocalMemory(buffer, capacity, kWildcardLocation);
|
|
if (ret) {
|
|
freeMemory(buffer);
|
|
return nullptr;
|
|
}
|
|
return (char *)buffer;
|
|
}
|
|
|
|
int TransferEnginePy::findClassId(size_t size) {
|
|
if (size > 1024ull * kSlabSizeKB[kMaxClassId]) return -1;
|
|
for (int i = kMaxClassId - 1; i >= 0; --i)
|
|
if (size > 1024ull * kSlabSizeKB[i]) return i + 1;
|
|
return 0;
|
|
}
|
|
|
|
int TransferEnginePy::doBuddyAllocate(int class_id) {
|
|
if (class_id == kMaxClassId) {
|
|
auto buffer = allocateRawBuffer(kDefaultBufferCapacity);
|
|
buffer_list_.push_back(buffer);
|
|
for (size_t offset = 0; offset < kDefaultBufferCapacity;
|
|
offset += 1024ull * kSlabSizeKB[kMaxClassId])
|
|
free_list_[kMaxClassId].push(buffer + offset);
|
|
return 0;
|
|
}
|
|
if (free_list_[class_id + 1].empty()) {
|
|
int ret = doBuddyAllocate(class_id + 1);
|
|
if (ret) return ret;
|
|
}
|
|
assert(!free_list_[class_id + 1].empty());
|
|
char *buffer = free_list_[class_id + 1].top();
|
|
free_list_[class_id + 1].pop();
|
|
free_list_[class_id].push(buffer);
|
|
free_list_[class_id].push(buffer + kSlabSizeKB[class_id] * 1024);
|
|
return 0;
|
|
}
|
|
|
|
uintptr_t TransferEnginePy::allocateManagedBuffer(size_t length) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
int class_id = findClassId(length);
|
|
if (class_id < 0) {
|
|
char *buffer = allocateRawBuffer(length);
|
|
if (buffer) large_buffer_list_.insert(buffer);
|
|
return (uintptr_t)buffer;
|
|
}
|
|
if (free_list_[class_id].empty())
|
|
if (doBuddyAllocate(class_id)) return 0;
|
|
assert(!free_list_[class_id].empty());
|
|
char *buffer = free_list_[class_id].top();
|
|
free_list_[class_id].pop();
|
|
return (uintptr_t)buffer;
|
|
}
|
|
|
|
int TransferEnginePy::freeManagedBuffer(uintptr_t buffer_addr, size_t length) {
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
auto buffer = (char *)buffer_addr;
|
|
int class_id = findClassId(length);
|
|
if (class_id < 0) {
|
|
large_buffer_list_.erase(buffer);
|
|
engine_->unregisterLocalMemory(buffer);
|
|
freeMemory(buffer);
|
|
return 0;
|
|
}
|
|
free_list_[class_id].push(buffer);
|
|
return 0;
|
|
}
|
|
|
|
int TransferEnginePy::transferSyncWrite(const char *target_hostname,
|
|
uintptr_t buffer,
|
|
uintptr_t peer_buffer_address,
|
|
size_t length) {
|
|
return transferSync(target_hostname, buffer, peer_buffer_address, length,
|
|
TransferOpcode::WRITE);
|
|
}
|
|
|
|
int TransferEnginePy::transferSyncRead(const char *target_hostname,
|
|
uintptr_t buffer,
|
|
uintptr_t peer_buffer_address,
|
|
size_t length) {
|
|
return transferSync(target_hostname, buffer, peer_buffer_address, length,
|
|
TransferOpcode::READ);
|
|
}
|
|
|
|
int TransferEnginePy::batchTransferSyncWrite(
|
|
const char *target_hostname, std::vector<uintptr_t> buffers,
|
|
std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths) {
|
|
return batchTransferSync(target_hostname, buffers, peer_buffer_addresses,
|
|
lengths, TransferOpcode::WRITE);
|
|
}
|
|
|
|
int TransferEnginePy::batchTransferSyncRead(
|
|
const char *target_hostname, std::vector<uintptr_t> buffers,
|
|
std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths) {
|
|
return batchTransferSync(target_hostname, buffers, peer_buffer_addresses,
|
|
lengths, TransferOpcode::READ);
|
|
}
|
|
|
|
batch_id_t TransferEnginePy::batchTransferAsyncWrite(
|
|
const char *target_hostname, const std::vector<uintptr_t> &buffers,
|
|
const std::vector<uintptr_t> &peer_buffer_addresses,
|
|
const std::vector<size_t> &lengths) {
|
|
return batchTransferAsync(target_hostname, buffers, peer_buffer_addresses,
|
|
lengths, TransferOpcode::WRITE);
|
|
}
|
|
|
|
batch_id_t TransferEnginePy::batchTransferAsyncRead(
|
|
const char *target_hostname, const std::vector<uintptr_t> &buffers,
|
|
const std::vector<uintptr_t> &peer_buffer_addresses,
|
|
const std::vector<size_t> &lengths) {
|
|
return batchTransferAsync(target_hostname, buffers, peer_buffer_addresses,
|
|
lengths, TransferOpcode::READ);
|
|
}
|
|
|
|
int TransferEnginePy::transferSync(const char *target_hostname,
|
|
uintptr_t buffer,
|
|
uintptr_t peer_buffer_address, size_t length,
|
|
TransferOpcode opcode,
|
|
TransferNotify *notify) {
|
|
pybind11::gil_scoped_release release;
|
|
Transport::SegmentHandle handle;
|
|
{
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
if (handle_map_.count(target_hostname)) {
|
|
handle = handle_map_[target_hostname];
|
|
} else {
|
|
LOG(INFO)
|
|
<< "transferSync, cache not found, openSegment with target "
|
|
<< target_hostname;
|
|
handle = engine_->openSegment(target_hostname);
|
|
if (handle == (Transport::SegmentHandle)-1) return -1;
|
|
handle_map_[target_hostname] = handle;
|
|
}
|
|
}
|
|
|
|
// TODO this is just a workaround
|
|
// When transfer engine submits one task, it will be dispatch to a worker
|
|
// associated with one local RNIC. If the local RNIC fails to connect to any
|
|
// remote RNIC, it will eventually fail. This allows selecting multiple
|
|
// local RNIC in one transferSync call. Will be fixed in the next revision.
|
|
const int max_retry =
|
|
engine_->numContexts() + 1; // Iter all possible local contexts
|
|
auto start_ts = getCurrentTimeInNano();
|
|
for (int retry = 0; retry < max_retry; ++retry) {
|
|
auto batch_id = engine_->allocateBatchID(1);
|
|
TransferRequest entry;
|
|
if (opcode == TransferOpcode::WRITE) {
|
|
entry.opcode = TransferRequest::WRITE;
|
|
} else {
|
|
entry.opcode = TransferRequest::READ;
|
|
}
|
|
entry.length = length;
|
|
entry.source = (void *)buffer;
|
|
entry.target_id = handle;
|
|
entry.target_offset = peer_buffer_address;
|
|
entry.advise_retry_cnt = retry;
|
|
|
|
Status s =
|
|
notify
|
|
? engine_->submitTransferWithNotify(
|
|
batch_id, {entry},
|
|
TransferMetadata::NotifyDesc{notify->name, notify->msg})
|
|
: engine_->submitTransfer(batch_id, {entry});
|
|
if (!s.ok()) {
|
|
Status segment_status = engine_->CheckSegmentStatus(handle);
|
|
if (!segment_status.ok()) {
|
|
LOG(WARNING)
|
|
<< "submitTransfer failed with target " << target_hostname
|
|
<< ", CheckSegmentStatus not ok, ready to closeSegment";
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
engine_->closeSegment(handle);
|
|
engine_->getMetadata()->removeSegmentDesc(target_hostname);
|
|
handle_map_.erase(target_hostname);
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
TransferStatus status;
|
|
bool completed = false;
|
|
while (!completed) {
|
|
Status s = engine_->getTransferStatus(batch_id, 0, status);
|
|
LOG_ASSERT(s.ok());
|
|
if (status.s == TransferStatusEnum::COMPLETED) {
|
|
engine_->freeBatchID(batch_id);
|
|
return 0;
|
|
} else if (status.s == TransferStatusEnum::FAILED) {
|
|
engine_->freeBatchID(batch_id);
|
|
completed = true;
|
|
} else if (status.s == TransferStatusEnum::TIMEOUT) {
|
|
LOG(INFO) << "Sync data transfer timeout";
|
|
completed = true;
|
|
}
|
|
auto current_ts = getCurrentTimeInNano();
|
|
const int64_t timeout =
|
|
transfer_timeout_nsec_ + length; // 1GiB per second
|
|
if (current_ts - start_ts > timeout) {
|
|
LOG(INFO) << "Sync data transfer timeout after "
|
|
<< current_ts - start_ts << "ns, local buffer "
|
|
<< (void *)buffer << " remote buffer "
|
|
<< (void *)peer_buffer_address << " length "
|
|
<< length;
|
|
return -1;
|
|
}
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
int TransferEnginePy::batchTransferSync(
|
|
const char *target_hostname, std::vector<uintptr_t> buffers,
|
|
std::vector<uintptr_t> peer_buffer_addresses, std::vector<size_t> lengths,
|
|
TransferOpcode opcode, TransferNotify *notify) {
|
|
pybind11::gil_scoped_release release;
|
|
Transport::SegmentHandle handle;
|
|
{
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
if (handle_map_.count(target_hostname)) {
|
|
handle = handle_map_[target_hostname];
|
|
} else {
|
|
handle = engine_->openSegment(target_hostname);
|
|
if (handle == (Transport::SegmentHandle)-1) return -1;
|
|
handle_map_[target_hostname] = handle;
|
|
}
|
|
}
|
|
|
|
if (buffers.size() != peer_buffer_addresses.size() ||
|
|
buffers.size() != lengths.size()) {
|
|
LOG(ERROR)
|
|
<< "buffers, peer_buffer_addresses and lengths have different size";
|
|
return -1;
|
|
}
|
|
|
|
const int max_retry = engine_->numContexts() + 1;
|
|
auto start_ts = getCurrentTimeInNano();
|
|
auto total_length = std::accumulate(lengths.begin(), lengths.end(), 0ull);
|
|
auto batch_size = buffers.size();
|
|
std::vector<TransferRequest> entries;
|
|
for (size_t i = 0; i < batch_size; ++i) {
|
|
TransferRequest entry;
|
|
if (opcode == TransferOpcode::WRITE) {
|
|
entry.opcode = TransferRequest::WRITE;
|
|
} else {
|
|
entry.opcode = TransferRequest::READ;
|
|
}
|
|
entry.length = lengths[i];
|
|
entry.source = (void *)buffers[i];
|
|
entry.target_id = handle;
|
|
entry.target_offset = peer_buffer_addresses[i];
|
|
entry.advise_retry_cnt = 0;
|
|
entries.push_back(entry);
|
|
}
|
|
|
|
for (int retry = 0; retry < max_retry; ++retry) {
|
|
auto batch_id = engine_->allocateBatchID(batch_size);
|
|
Status s =
|
|
notify
|
|
? engine_->submitTransferWithNotify(
|
|
batch_id, entries,
|
|
TransferMetadata::NotifyDesc{notify->name, notify->msg})
|
|
: engine_->submitTransfer(batch_id, entries);
|
|
if (!s.ok()) {
|
|
engine_->freeBatchID(batch_id);
|
|
Status segment_status = engine_->CheckSegmentStatus(handle);
|
|
if (!segment_status.ok()) {
|
|
LOG(WARNING)
|
|
<< "submitTransfer failed with target " << target_hostname
|
|
<< ", CheckSegmentStatus not ok, ready to closeSegment";
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
engine_->closeSegment(handle);
|
|
engine_->getMetadata()->removeSegmentDesc(target_hostname);
|
|
handle_map_.erase(target_hostname);
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
TransferStatus status;
|
|
bool completed = false;
|
|
bool already_freed = false;
|
|
while (!completed) {
|
|
Status s = engine_->getBatchTransferStatus(batch_id, status);
|
|
LOG_ASSERT(s.ok());
|
|
if (status.s == TransferStatusEnum::COMPLETED) {
|
|
engine_->freeBatchID(batch_id);
|
|
return 0;
|
|
} else if (status.s == TransferStatusEnum::FAILED) {
|
|
engine_->freeBatchID(batch_id);
|
|
already_freed = true;
|
|
completed = true;
|
|
} else if (status.s == TransferStatusEnum::TIMEOUT) {
|
|
LOG(INFO) << "Sync data transfer timeout";
|
|
completed = true;
|
|
}
|
|
auto current_ts = getCurrentTimeInNano();
|
|
const int64_t timeout =
|
|
transfer_timeout_nsec_ + total_length; // 1GiB per second
|
|
if (current_ts - start_ts > timeout) {
|
|
LOG(INFO) << "Sync batch data transfer timeout after "
|
|
<< current_ts - start_ts << "ns";
|
|
// TODO: as @doujiang24 mentioned, early free(while there are
|
|
// still waiting tasks) the batch_id may fail and cause memory
|
|
// leak(a known issue).
|
|
if (!already_freed) {
|
|
engine_->freeBatchID(batch_id);
|
|
}
|
|
return -1;
|
|
}
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
batch_id_t TransferEnginePy::batchTransferAsync(
|
|
const char *target_hostname, const std::vector<uintptr_t> &buffers,
|
|
const std::vector<uintptr_t> &peer_buffer_addresses,
|
|
const std::vector<size_t> &lengths, TransferOpcode opcode) {
|
|
pybind11::gil_scoped_release release;
|
|
Transport::SegmentHandle handle;
|
|
{
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
if (handle_map_.count(target_hostname)) {
|
|
handle = handle_map_[target_hostname];
|
|
} else {
|
|
handle = engine_->openSegment(target_hostname);
|
|
if (handle == (Transport::SegmentHandle)-1) return -1;
|
|
handle_map_[target_hostname] = handle;
|
|
}
|
|
}
|
|
|
|
if (buffers.size() != peer_buffer_addresses.size() ||
|
|
buffers.size() != lengths.size()) {
|
|
LOG(ERROR)
|
|
<< "buffers, peer_buffer_addresses and lengths have different size";
|
|
return 0;
|
|
}
|
|
|
|
const int max_retry = engine_->numContexts() + 1;
|
|
auto batch_size = buffers.size();
|
|
std::vector<TransferRequest> entries;
|
|
batch_id_t batch_id = 0;
|
|
for (size_t i = 0; i < batch_size; ++i) {
|
|
TransferRequest entry;
|
|
if (opcode == TransferOpcode::WRITE) {
|
|
entry.opcode = TransferRequest::WRITE;
|
|
} else {
|
|
entry.opcode = TransferRequest::READ;
|
|
}
|
|
entry.length = lengths[i];
|
|
entry.source = (void *)buffers[i];
|
|
entry.target_id = handle;
|
|
entry.target_offset = peer_buffer_addresses[i];
|
|
entry.advise_retry_cnt = 0;
|
|
entries.push_back(entry);
|
|
}
|
|
|
|
for (int retry = 0; retry < max_retry; ++retry) {
|
|
batch_id = engine_->allocateBatchID(batch_size);
|
|
auto batch_desc = reinterpret_cast<BatchDesc *>(batch_id);
|
|
|
|
auto start_ts = getCurrentTimeInNano();
|
|
batch_desc->start_timestamp = start_ts;
|
|
|
|
Status s = engine_->submitTransfer(batch_id, entries);
|
|
if (!s.ok()) {
|
|
engine_->freeBatchID(batch_id);
|
|
return 0;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
|
|
return batch_id;
|
|
}
|
|
|
|
int TransferEnginePy::getBatchTransferStatus(
|
|
const std::vector<batch_id_t> &batch_ids) {
|
|
pybind11::gil_scoped_release release;
|
|
TransferStatus status;
|
|
std::unordered_map<batch_id_t, int64_t> timeout_table{};
|
|
for (auto &batch_id : batch_ids) {
|
|
int64_t total_length = 0;
|
|
auto batch_desc = reinterpret_cast<BatchDesc *>(batch_id);
|
|
const size_t task_count = batch_desc->task_list.size();
|
|
|
|
for (size_t task_id = 0; task_id < task_count; task_id++) {
|
|
auto &task = batch_desc->task_list[task_id];
|
|
for (auto &slice : task.slice_list) {
|
|
total_length += slice->length;
|
|
}
|
|
}
|
|
|
|
timeout_table[batch_id] = total_length + transfer_timeout_nsec_;
|
|
}
|
|
|
|
bool failed_or_timeout = false;
|
|
std::unordered_set<batch_id_t> remove_ids{};
|
|
while (!timeout_table.empty() && !failed_or_timeout) {
|
|
for (auto &entry : timeout_table) {
|
|
auto batch_desc = reinterpret_cast<BatchDesc *>(entry.first);
|
|
Status s = engine_->getBatchTransferStatus(entry.first, status);
|
|
LOG_ASSERT(s.ok());
|
|
if (status.s == TransferStatusEnum::COMPLETED) {
|
|
engine_->freeBatchID(entry.first);
|
|
LOG(INFO) << "Batch Transfer completed!";
|
|
remove_ids.insert(entry.first);
|
|
} else if (status.s == TransferStatusEnum::FAILED) {
|
|
failed_or_timeout = true;
|
|
} else if (status.s == TransferStatusEnum::TIMEOUT) {
|
|
LOG(INFO) << "Sync data transfer timeout";
|
|
}
|
|
auto current_ts = getCurrentTimeInNano();
|
|
if (current_ts - batch_desc->start_timestamp > entry.second) {
|
|
LOG(INFO) << "Sync batch data transfer timeout after "
|
|
<< current_ts - batch_desc->start_timestamp << "ns";
|
|
failed_or_timeout = true;
|
|
}
|
|
}
|
|
|
|
for (auto &remove_id : remove_ids) {
|
|
timeout_table.erase(remove_id);
|
|
}
|
|
|
|
remove_ids.clear();
|
|
}
|
|
|
|
if (failed_or_timeout) {
|
|
for (auto &entry : timeout_table) {
|
|
engine_->freeBatchID(entry.first);
|
|
}
|
|
}
|
|
|
|
return failed_or_timeout ? -1 : 0;
|
|
}
|
|
|
|
batch_id_t TransferEnginePy::transferSubmitWrite(const char *target_hostname,
|
|
uintptr_t buffer,
|
|
uintptr_t peer_buffer_address,
|
|
size_t length) {
|
|
pybind11::gil_scoped_release release;
|
|
Transport::SegmentHandle handle;
|
|
{
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
if (handle_map_.count(target_hostname)) {
|
|
handle = handle_map_[target_hostname];
|
|
} else {
|
|
handle = engine_->openSegment(target_hostname);
|
|
if (handle == (Transport::SegmentHandle)-1) return -1;
|
|
handle_map_[target_hostname] = handle;
|
|
}
|
|
}
|
|
|
|
auto batch_id = engine_->allocateBatchID(1);
|
|
TransferRequest entry;
|
|
entry.opcode = TransferRequest::WRITE;
|
|
entry.length = length;
|
|
entry.source = (void *)buffer;
|
|
entry.target_id = handle;
|
|
entry.target_offset = peer_buffer_address;
|
|
|
|
Status s = engine_->submitTransfer(batch_id, {entry});
|
|
if (!s.ok()) return -1;
|
|
|
|
return batch_id;
|
|
}
|
|
|
|
int TransferEnginePy::transferCheckStatus(batch_id_t batch_id) {
|
|
pybind11::gil_scoped_release release;
|
|
TransferStatus status;
|
|
Status s = engine_->getTransferStatus(batch_id, 0, status);
|
|
LOG_ASSERT(s.ok());
|
|
if (status.s == TransferStatusEnum::COMPLETED) {
|
|
engine_->freeBatchID(batch_id);
|
|
return 1;
|
|
} else if (status.s == TransferStatusEnum::FAILED) {
|
|
engine_->freeBatchID(batch_id);
|
|
return -1;
|
|
} else if (status.s == TransferStatusEnum::TIMEOUT) {
|
|
return -2;
|
|
} else {
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
int TransferEnginePy::batchRegisterMemory(
|
|
std::vector<uintptr_t> buffer_addresses, std::vector<size_t> capacities) {
|
|
pybind11::gil_scoped_release release;
|
|
auto batch_size = buffer_addresses.size();
|
|
std::vector<BufferEntry> buffers;
|
|
for (size_t i = 0; i < batch_size; i++) {
|
|
buffers.push_back(
|
|
BufferEntry{(void *)buffer_addresses[i], capacities[i]});
|
|
}
|
|
return engine_->registerLocalMemoryBatch(buffers, kWildcardLocation);
|
|
}
|
|
|
|
int TransferEnginePy::batchUnregisterMemory(
|
|
std::vector<uintptr_t> buffer_addresses) {
|
|
pybind11::gil_scoped_release release;
|
|
auto batch_size = buffer_addresses.size();
|
|
std::vector<void *> buffers;
|
|
for (size_t i = 0; i < batch_size; i++) {
|
|
buffers.push_back(reinterpret_cast<char *>(buffer_addresses[i]));
|
|
}
|
|
return engine_->unregisterLocalMemoryBatch(buffers);
|
|
}
|
|
|
|
int TransferEnginePy::registerMemory(uintptr_t buffer_addr, size_t capacity) {
|
|
char *buffer = reinterpret_cast<char *>(buffer_addr);
|
|
return engine_->registerLocalMemory(buffer, capacity);
|
|
}
|
|
|
|
int TransferEnginePy::unregisterMemory(uintptr_t buffer_addr) {
|
|
char *buffer = reinterpret_cast<char *>(buffer_addr);
|
|
return engine_->unregisterLocalMemory(buffer);
|
|
}
|
|
|
|
#ifdef USE_CUDA
|
|
|
|
/**
|
|
* @brief Context structure for CUDA-stream-synchronized transfers.
|
|
*
|
|
* This structure holds all necessary data to execute a Mooncake transfer
|
|
* from within a CUDA host callback.
|
|
*/
|
|
struct TransferOnCudaContext {
|
|
std::shared_ptr<TransferEngine> engine;
|
|
Transport::BatchID batch_id;
|
|
std::vector<Transport::TransferRequest> requests;
|
|
uint64_t total_bytes;
|
|
};
|
|
|
|
/**
|
|
* @brief CUDA Host Callback function for triggered transfers.
|
|
*
|
|
* This function is called by the CUDA driver when all preceding operations
|
|
* in the associated stream have completed. It submits the transfer requests
|
|
* and waits synchronously for their completion.
|
|
*
|
|
* @param data Pointer to a TransferOnCudaContext object.
|
|
*/
|
|
void CUDART_CB transfer_on_cuda_callback(void *data) {
|
|
auto *ctx = reinterpret_cast<TransferOnCudaContext *>(data);
|
|
|
|
auto status = ctx->engine->submitTransfer(ctx->batch_id, ctx->requests);
|
|
if (!status.ok()) {
|
|
LOG(ERROR) << "[Mooncake Cuda] Submit failed: " << status.ToString()
|
|
<< " | BatchID: " << ctx->batch_id;
|
|
goto error_exit;
|
|
}
|
|
|
|
Transport::TransferStatus t_status;
|
|
while (true) {
|
|
auto ret = ctx->engine->getBatchTransferStatus(ctx->batch_id, t_status);
|
|
if (!ret.ok()) {
|
|
LOG(ERROR) << "[Mooncake Cuda] Failed to get status for BatchID: "
|
|
<< ctx->batch_id;
|
|
goto error_exit;
|
|
}
|
|
|
|
if (t_status.s == Transport::TransferStatusEnum::COMPLETED) {
|
|
break;
|
|
} else if (t_status.s == Transport::TransferStatusEnum::FAILED) {
|
|
LOG(ERROR) << "[Mooncake Cuda] Transfer failed | BatchID: "
|
|
<< ctx->batch_id << " | Bytes: " << ctx->total_bytes;
|
|
goto error_exit;
|
|
} else if (t_status.s == Transport::TransferStatusEnum::TIMEOUT) {
|
|
LOG(ERROR) << "[Mooncake Cuda] Transfer timeout | BatchID: "
|
|
<< ctx->batch_id;
|
|
goto error_exit;
|
|
}
|
|
}
|
|
|
|
ctx->engine->freeBatchID(ctx->batch_id);
|
|
delete ctx;
|
|
return;
|
|
|
|
error_exit:
|
|
// Since this is a CUDA host callback running in a driver thread,
|
|
// we cannot propagate exceptions or error codes back to the main
|
|
// application. A failure here implies the data transfer required for
|
|
// subsequent stream operations has failed, leaving the system in an
|
|
// inconsistent state. We use _exit(1) to terminate the process
|
|
// immediately and avoid undefined behavior.
|
|
_exit(1);
|
|
}
|
|
|
|
/**
|
|
* @brief Submits a batch of transfer requests synchronized with a CUDA stream.
|
|
*
|
|
* This method schedules a host callback on the provided CUDA stream. The
|
|
* Mooncake transfer will only start after all previous kernels/memcpys on
|
|
* the stream have finished.
|
|
*
|
|
* @param target_hostname Remote host to transfer to/from.
|
|
* @param buffers Local buffer addresses.
|
|
* @param peer_buffer_addresses Remote buffer addresses.
|
|
* @param lengths Length of each transfer in bytes.
|
|
* @param opcode READ or WRITE operation.
|
|
* @param stream_ptr Handle to a CUDA stream (cudaStream_t as uintptr_t).
|
|
*/
|
|
void TransferEnginePy::batchTransferOnCuda(
|
|
const char *target_hostname, const std::vector<uintptr_t> &buffers,
|
|
const std::vector<uintptr_t> &peer_buffer_addresses,
|
|
const std::vector<size_t> &lengths, TransferOpcode opcode,
|
|
uintptr_t stream_ptr) {
|
|
pybind11::gil_scoped_release release;
|
|
Transport::SegmentHandle handle;
|
|
{
|
|
std::lock_guard<std::mutex> guard(mutex_);
|
|
if (handle_map_.count(target_hostname)) {
|
|
handle = handle_map_[target_hostname];
|
|
} else {
|
|
handle = engine_->openSegment(target_hostname);
|
|
if (handle == (Transport::SegmentHandle)-1)
|
|
throw std::runtime_error("Failed to open segment");
|
|
handle_map_[target_hostname] = handle;
|
|
}
|
|
}
|
|
|
|
if (buffers.size() != peer_buffer_addresses.size() ||
|
|
buffers.size() != lengths.size()) {
|
|
LOG(ERROR)
|
|
<< "buffers, peer_buffer_addresses and lengths have different size";
|
|
throw std::runtime_error(
|
|
"buffers, peer_buffer_addresses and lengths have different size");
|
|
}
|
|
|
|
size_t batch_size = buffers.size();
|
|
std::vector<TransferRequest> entries;
|
|
uint64_t total_bytes = 0;
|
|
for (size_t i = 0; i < batch_size; ++i) {
|
|
TransferRequest entry;
|
|
entry.opcode = (opcode == TransferOpcode::WRITE)
|
|
? TransferRequest::WRITE
|
|
: TransferRequest::READ;
|
|
entry.length = lengths[i];
|
|
entry.source = (void *)buffers[i];
|
|
entry.target_id = handle;
|
|
entry.target_offset = peer_buffer_addresses[i];
|
|
entries.push_back(entry);
|
|
total_bytes += lengths[i];
|
|
}
|
|
|
|
auto batch_id = engine_->allocateBatchID(batch_size);
|
|
auto *ctx = new TransferOnCudaContext{engine_, batch_id, std::move(entries),
|
|
total_bytes};
|
|
|
|
cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
|
|
cudaError_t err =
|
|
cudaLaunchHostFunc(stream, transfer_on_cuda_callback, ctx);
|
|
if (err != cudaSuccess) {
|
|
delete ctx;
|
|
engine_->freeBatchID(batch_id);
|
|
throw std::runtime_error(std::string("cudaLaunchHostFunc failed: ") +
|
|
cudaGetErrorString(err));
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief Async WRITE transfer triggered by a CUDA stream.
|
|
*/
|
|
void TransferEnginePy::transferWriteOnCuda(const char *target_hostname,
|
|
uintptr_t buffer,
|
|
uintptr_t peer_buffer_address,
|
|
size_t length,
|
|
uintptr_t stream_ptr) {
|
|
batchTransferOnCuda(target_hostname, {buffer}, {peer_buffer_address},
|
|
{length}, TransferOpcode::WRITE, stream_ptr);
|
|
}
|
|
|
|
/**
|
|
* @brief Async READ transfer triggered by a CUDA stream.
|
|
*/
|
|
void TransferEnginePy::transferReadOnCuda(const char *target_hostname,
|
|
uintptr_t buffer,
|
|
uintptr_t peer_buffer_address,
|
|
size_t length, uintptr_t stream_ptr) {
|
|
batchTransferOnCuda(target_hostname, {buffer}, {peer_buffer_address},
|
|
{length}, TransferOpcode::READ, stream_ptr);
|
|
}
|
|
|
|
/**
|
|
* @brief Batch async WRITE transfer triggered by a CUDA stream.
|
|
*/
|
|
void TransferEnginePy::batchTransferWriteOnCuda(
|
|
const char *target_hostname, const std::vector<uintptr_t> &buffers,
|
|
const std::vector<uintptr_t> &peer_buffer_addresses,
|
|
const std::vector<size_t> &lengths, uintptr_t stream_ptr) {
|
|
batchTransferOnCuda(target_hostname, buffers, peer_buffer_addresses,
|
|
lengths, TransferOpcode::WRITE, stream_ptr);
|
|
}
|
|
|
|
/**
|
|
* @brief Batch async READ transfer triggered by a CUDA stream.
|
|
*/
|
|
void TransferEnginePy::batchTransferReadOnCuda(
|
|
const char *target_hostname, const std::vector<uintptr_t> &buffers,
|
|
const std::vector<uintptr_t> &peer_buffer_addresses,
|
|
const std::vector<size_t> &lengths, uintptr_t stream_ptr) {
|
|
batchTransferOnCuda(target_hostname, buffers, peer_buffer_addresses,
|
|
lengths, TransferOpcode::READ, stream_ptr);
|
|
}
|
|
#endif
|
|
|
|
uintptr_t TransferEnginePy::getFirstBufferAddress(
|
|
const std::string &segment_name) {
|
|
Transport::SegmentHandle segment_id =
|
|
engine_->openSegment(segment_name.c_str());
|
|
auto segment_desc = engine_->getMetadata()->getSegmentDescByID(segment_id);
|
|
if (!segment_desc || segment_desc->buffers.empty()) {
|
|
return 0;
|
|
}
|
|
return segment_desc->buffers[0].addr;
|
|
}
|
|
|
|
std::string TransferEnginePy::getLocalTopology(const char *device_name) {
|
|
pybind11::gil_scoped_release release;
|
|
auto device_name_safe = device_name ? std::string(device_name) : "";
|
|
auto device_filter = buildDeviceFilter(device_name_safe);
|
|
std::shared_ptr<TransferEngine> tmp_engine =
|
|
std::make_shared<TransferEngine>(true, device_filter);
|
|
|
|
std::string metadata_conn_string{"P2PHANDSHAKE"}, local_server_name{};
|
|
tmp_engine->init(metadata_conn_string, local_server_name);
|
|
|
|
return tmp_engine->getLocalTopology()->toString();
|
|
}
|
|
|
|
std::vector<TransferEnginePy::TransferNotify> TransferEnginePy::getNotifies() {
|
|
std::vector<TransferMetadata::NotifyDesc> notifies;
|
|
std::vector<TransferNotify> result;
|
|
|
|
int ret = engine_->getNotifies(notifies);
|
|
if (ret != 0) {
|
|
LOG(ERROR) << "Failed to get notifies: " << ret;
|
|
return result;
|
|
}
|
|
|
|
for (const auto ¬ify : notifies) {
|
|
result.emplace_back(
|
|
TransferEnginePy::TransferNotify{notify.name, notify.notify_msg});
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
namespace py = pybind11;
|
|
|
|
// Implementation of coro_rpc_interface binding function
|
|
void bind_coro_rpc_interface(py::module_ &m) {
|
|
// Note: RpcInterface, ReceivedData and ReceivedTensor are already
|
|
// registered by bind_rpc_interface() so we don't register them again here
|
|
// to avoid duplicate type registration errors. The factory functions are
|
|
// also registered by bind_rpc_interface(), so we don't need to register
|
|
// them again.
|
|
|
|
// Add CoroRPCInterface as an alias to RpcInterface
|
|
m.attr("CoroRPCInterface") = m.attr("RpcInterface");
|
|
}
|
|
|
|
PYBIND11_MODULE(engine, m) {
|
|
py::enum_<TransferEnginePy::TransferOpcode> transfer_opcode(
|
|
m, "TransferOpcode", py::arithmetic());
|
|
transfer_opcode.value("Read", TransferEnginePy::TransferOpcode::READ)
|
|
.value("Write", TransferEnginePy::TransferOpcode::WRITE)
|
|
.export_values();
|
|
|
|
py::class_<TransferEnginePy::TransferNotify>(m, "TransferNotify")
|
|
.def(py::init<>())
|
|
.def(py::init<const std::string &, const std::string &>(),
|
|
py::arg("name"), py::arg("msg"))
|
|
.def_readwrite("name", &TransferEnginePy::TransferNotify::name)
|
|
.def_readwrite("msg", &TransferEnginePy::TransferNotify::msg);
|
|
|
|
auto adaptor_cls =
|
|
py::class_<TransferEnginePy>(m, "TransferEngine")
|
|
.def(py::init<>())
|
|
.def("initialize", &TransferEnginePy::initialize)
|
|
.def("initialize_ext", &TransferEnginePy::initializeExt)
|
|
.def("get_rpc_port", &TransferEnginePy::getRpcPort)
|
|
.def("allocate_managed_buffer",
|
|
&TransferEnginePy::allocateManagedBuffer)
|
|
.def("free_managed_buffer", &TransferEnginePy::freeManagedBuffer)
|
|
.def("transfer_sync_write", &TransferEnginePy::transferSyncWrite)
|
|
.def("transfer_sync_read", &TransferEnginePy::transferSyncRead)
|
|
.def("batch_transfer_sync_write",
|
|
&TransferEnginePy::batchTransferSyncWrite)
|
|
.def("batch_transfer_sync_read",
|
|
&TransferEnginePy::batchTransferSyncRead)
|
|
.def("batch_transfer_async_write",
|
|
&TransferEnginePy::batchTransferAsyncWrite)
|
|
.def("batch_transfer_async_read",
|
|
&TransferEnginePy::batchTransferAsyncRead)
|
|
.def("transfer_sync", &TransferEnginePy::transferSync,
|
|
py::arg("target_hostname"), py::arg("buffer"),
|
|
py::arg("peer_buffer_address"), py::arg("length"),
|
|
py::arg("opcode"), py::arg("notify") = nullptr)
|
|
.def("batch_transfer_sync", &TransferEnginePy::batchTransferSync)
|
|
.def("batch_transfer_async", &TransferEnginePy::batchTransferAsync)
|
|
#ifdef USE_CUDA
|
|
.def("transfer_write_on_cuda",
|
|
&TransferEnginePy::transferWriteOnCuda,
|
|
py::arg("target_hostname"), py::arg("buffer"),
|
|
py::arg("peer_buffer_address"), py::arg("length"),
|
|
py::arg("stream_ptr") = 0)
|
|
.def("transfer_read_on_cuda", &TransferEnginePy::transferReadOnCuda,
|
|
py::arg("target_hostname"), py::arg("buffer"),
|
|
py::arg("peer_buffer_address"), py::arg("length"),
|
|
py::arg("stream_ptr") = 0)
|
|
.def("batch_transfer_write_on_cuda",
|
|
&TransferEnginePy::batchTransferWriteOnCuda,
|
|
py::arg("target_hostname"), py::arg("buffers"),
|
|
py::arg("peer_buffer_addresses"), py::arg("lengths"),
|
|
py::arg("stream_ptr") = 0)
|
|
.def("batch_transfer_read_on_cuda",
|
|
&TransferEnginePy::batchTransferReadOnCuda,
|
|
py::arg("target_hostname"), py::arg("buffers"),
|
|
py::arg("peer_buffer_addresses"), py::arg("lengths"),
|
|
py::arg("stream_ptr") = 0)
|
|
#endif
|
|
.def("get_batch_transfer_status",
|
|
&TransferEnginePy::getBatchTransferStatus)
|
|
.def("transfer_submit_write",
|
|
&TransferEnginePy::transferSubmitWrite)
|
|
.def("transfer_check_status",
|
|
&TransferEnginePy::transferCheckStatus)
|
|
.def("write_bytes_to_buffer", &TransferEnginePy::writeBytesToBuffer)
|
|
.def("read_bytes_from_buffer",
|
|
&TransferEnginePy::readBytesFromBuffer)
|
|
.def("register_memory", &TransferEnginePy::registerMemory)
|
|
.def("unregister_memory", &TransferEnginePy::unregisterMemory)
|
|
.def("batch_register_memory",
|
|
&TransferEnginePy::batchRegisterMemory)
|
|
.def("batch_unregister_memory",
|
|
&TransferEnginePy::batchUnregisterMemory)
|
|
.def("get_local_topology", &TransferEnginePy::getLocalTopology,
|
|
py::arg("device_name") = nullptr)
|
|
.def("get_first_buffer_address",
|
|
&TransferEnginePy::getFirstBufferAddress)
|
|
.def("get_notifies", &TransferEnginePy::getNotifies)
|
|
.def("get_engine", &TransferEnginePy::getEngine)
|
|
.def("get_engine_ptr", &TransferEnginePy::getEnginePtr);
|
|
|
|
adaptor_cls.attr("TransferOpcode") = transfer_opcode;
|
|
|
|
py::class_<TransferEngine, std::shared_ptr<TransferEngine>>(
|
|
m, "InnerTransferEngine");
|
|
|
|
// Bind RpcInterface (this also registers ReceivedData, ReceivedTensor, and
|
|
// factory functions)
|
|
mooncake::bind_rpc_interface(m);
|
|
|
|
// Add CoroRPCInterface as an alias to RpcInterface if needed
|
|
bind_coro_rpc_interface(m);
|
|
}
|