680 lines
25 KiB
C++
680 lines
25 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>
|
|
|
|
#ifdef USE_MNNVL
|
|
#include "transport/nvlink_transport/nvlink_transport.h"
|
|
static void *allocateMemory(size_t size) {
|
|
return mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
|
|
}
|
|
static void freeMemory(void *ptr) {
|
|
mooncake::NvlinkTransport::freePinnedLocalMemory(ptr);
|
|
}
|
|
#else
|
|
static void *allocateMemory(size_t size) { return malloc(size); }
|
|
static void freeMemory(void *ptr) { free(ptr); }
|
|
#endif
|
|
|
|
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) {
|
|
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) {
|
|
(void)(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);
|
|
engine_ = std::make_unique<TransferEngine>(true, device_filter);
|
|
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;
|
|
}
|
|
|
|
free_list_.resize(kSlabSizeKBTabLen);
|
|
#ifndef USE_ASCEND
|
|
doBuddyAllocate(kMaxClassId);
|
|
#endif
|
|
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) {
|
|
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;
|
|
}
|
|
}
|
|
|
|
// 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 = engine_->submitTransfer(batch_id, {entry});
|
|
if (!s.ok()) 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) {
|
|
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 = engine_->submitTransfer(batch_id, entries);
|
|
if (!s.ok()) {
|
|
engine_->freeBatchID(batch_id);
|
|
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);
|
|
}
|
|
|
|
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);
|
|
return segment_desc->buffers[0].addr;
|
|
}
|
|
|
|
namespace py = pybind11;
|
|
|
|
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();
|
|
|
|
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)
|
|
.def("batch_transfer_sync", &TransferEnginePy::batchTransferSync)
|
|
.def("batch_transfer_async", &TransferEnginePy::batchTransferAsync)
|
|
.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_first_buffer_address",
|
|
&TransferEnginePy::getFirstBufferAddress);
|
|
|
|
adaptor_cls.attr("TransferOpcode") = transfer_opcode;
|
|
}
|