zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[TransferEngine] Fix compilation bugs of nvmeof transport (#174)

This commit is contained in:
Feng Ren 2025-03-31 10:00:40 +08:00 committed by GitHub
parent 7f2b2b536d
commit 206198bc20
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 246 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mooncake-transfer-engine/include/transport/nvmeof_transport/cufile_context.h
View File

@ -66,8 +66,7 @@ class CuFileContext {
/// Create a GDS segment from file name. Return NULL on error.
explicit CuFileContext(const char *filename) {
int fd = open(filename, O_RDWR | O_DIRECT, 0664);
LOG(INFO) << "open " << filename << " get " << fd;
int fd = open(filename, O_RDWR | O_DIRECT);
memset(&desc, 0, sizeof(desc));
desc.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;
desc.handle.fd = fd;

20
mooncake-transfer-engine/include/transport/nvmeof_transport/nvmeof_transport.h
View File

@ -37,6 +37,10 @@ class NVMeoFTransport : public Transport {
BatchID allocateBatchID(size_t batch_size) override;
Status submitTransferTask(
const std::vector<TransferRequest *> &request_list,
const std::vector<TransferTask *> &task_list) override;
Status submitTransfer(BatchID batch_id,
const std::vector<TransferRequest> &entries) override;
@ -44,16 +48,11 @@ class NVMeoFTransport : public Transport {
TransferStatus &status) override;
Status freeBatchID(BatchID batch_id) override;
private:
void startTransfer(Slice *slice);
private:
struct NVMeoFBatchDesc {
size_t desc_idx_;
std::vector<TransferStatus> transfer_status;
std::vector<std::pair<uint64_t, uint64_t>>
task_to_slices; // task id -> (slice_begin, slice_num)
// unsigned nr_completed;
};
struct pair_hash {
template <class T1, class T2>
std::size_t operator()(const std::pair<T1, T2> &pair) const {
@ -85,16 +84,13 @@ class NVMeoFTransport : public Transport {
return 0;
}
void addSliceToTask(void *source_addr, uint64_t slice_len,
uint64_t target_start, TransferRequest::OpCode op,
TransferTask &task, const char *file_path);
void addSliceToCUFileBatch(void *source_addr, uint64_t file_offset,
uint64_t slice_len, uint64_t desc_id,
TransferRequest::OpCode op, CUfileHandle_t fh);
const char *getName() const override { return "nvmeof"; }
std::unordered_map<BatchID, int> batch_to_cufile_desc_;
std::unordered_map<std::pair<SegmentHandle, uint64_t>,
std::shared_ptr<CuFileContext>, pair_hash>
segment_to_context_;

9
mooncake-transfer-engine/include/transport/transport.h
View File

@ -104,10 +104,8 @@ class Transport {
uint64_t dest_addr;
} tcp;
struct {
const char *file_path;
uint64_t start;
uint64_t length;
uint64_t buffer_id;
uint64_t offset;
int cufile_desc;
} nvmeof;
struct {
void *remote_filename;
@ -130,14 +128,13 @@ class Transport {
};
struct TransferTask {
volatile uint64_t slice_count = 0;
volatile uint64_t success_slice_count = 0;
volatile uint64_t failed_slice_count = 0;
volatile uint64_t transferred_bytes = 0;
volatile bool is_finished = false;
uint64_t total_bytes = 0;
BatchID batch_id = 0;
};
struct BatchDesc {

17
mooncake-transfer-engine/scripts/register.py
View File

@ -1,3 +1,4 @@
# /usr/bin/python
# Copyright 2024 KVCache.AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
@ -27,21 +28,13 @@ if __name__ == "__main__":
os.environ.pop("HTTPS_PROXY", None)
os.environ.pop("http_proxy", None)
os.environ.pop("https_proxy", None)
etcd_server = sys.argv[1]
segment_name = "nvmeof/" + sys.argv[2]
etcd_host = sys.argv[1]
segment_name = "mooncake/nvmeof/" + sys.argv[2]
files = sys.argv[3:]
local_server_name = socket.gethostname()
etcd_host = sys.argv[1]
files = sys.argv[2:]
server_name = socket.gethostname()
print(server_name)
segment_name = "mooncake/nvmeof/" + server_name
print(segment_name)
etcd = etcd3.client(host=etcd_host, port=2379)
value = {}
value['server_name'] = server_name
value['protocol'] = "nvmeof"
@ -51,7 +44,9 @@ if __name__ == "__main__":
buffer = {}
buffer['length'] = os.path.getsize(file)
buffer['file_path'] = file
buffer['local_path_map'] = {}
local_path_map = {}
local_path_map[server_name] = file
buffer['local_path_map'] = local_path_map
value['buffers'].append(buffer)
print(json.dumps(value))

8
mooncake-transfer-engine/src/multi_transport.cpp
View File

@ -54,7 +54,8 @@ Status MultiTransport::freeBatchID(BatchID batch_id) {
if (!batch_desc.task_list[task_id].is_finished) {
LOG(ERROR) << "BatchID cannot be freed until all tasks are done";
return Status::BatchBusy(
"BatchID cannot be freed until all tasks are done"); }
"BatchID cannot be freed until all tasks are done");
}
}
delete &batch_desc;
#ifdef CONFIG_USE_BATCH_DESC_SET
@ -88,6 +89,7 @@ Status MultiTransport::submitTransfer(
std::to_string(request.target_id));
}
auto &task = batch_desc.task_list[task_id];
task.batch_id = batch_id;
++task_id;
submit_tasks[transport].request_list.push_back(
(TransferRequest *)&request);
@ -95,7 +97,7 @@ Status MultiTransport::submitTransfer(
}
for (auto &entry : submit_tasks) {
auto status = entry.first->submitTransferTask(entry.second.request_list,
entry.second.task_list);
entry.second.task_list);
if (!status.ok()) {
LOG(ERROR) << "MultiTransport: Failed to submit transfer task to "
<< entry.first->getName();
@ -106,7 +108,7 @@ Status MultiTransport::submitTransfer(
}
Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
TransferStatus &status) {
TransferStatus &status) {
auto &batch_desc = *((BatchDesc *)(batch_id));
const size_t task_count = batch_desc.task_list.size();
if (task_id >= task_count) {

6
mooncake-transfer-engine/tests/CMakeLists.txt
View File

@ -10,6 +10,12 @@ add_executable(rdma_transport_test2 rdma_transport_test2.cpp)
target_link_libraries(rdma_transport_test2 PUBLIC transfer_engine gtest gtest_main )
# add_test(NAME rdma_transport_test2 COMMAND rdma_transport_test2)
if (USE_NVMEOF)
add_executable(nvmeof_transport_test nvmeof_transport_test.cpp)
target_link_libraries(nvmeof_transport_test PUBLIC transfer_engine gtest gtest_main )
# add_test(NAME nvmeof_transport_test COMMAND nvmeof_transport_test)
endif()
add_executable(tcp_transport_test tcp_transport_test.cpp)
target_link_libraries(tcp_transport_test PUBLIC transfer_engine gtest gtest_main )
add_test(NAME tcp_transport_test COMMAND tcp_transport_test)

217
mooncake-transfer-engine/tests/nvmeof_transport_test.cpp Normal file
View File

@ -0,0 +1,217 @@
// 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 <gflags/gflags.h>
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <sys/time.h>
#include <cstdlib>
#include <fstream>
#include <iomanip>
#include <memory>
#include "transfer_engine.h"
#include "transport/transport.h"
using namespace mooncake;
namespace mooncake {
static std::string getHostname() {
char hostname[256];
if (gethostname(hostname, 256)) {
PLOG(ERROR) << "Failed to get hostname";
return "";
}
return hostname;
}
DEFINE_string(local_server_name, getHostname(),
"Local server name for segment discovery");
DEFINE_string(metadata_server, "127.0.0.1:2379", "etcd server host address");
DEFINE_string(mode, "initiator",
"Running mode: initiator or target. Initiator node read/write "
"data blocks from target node");
DEFINE_string(operation, "read", "Operation type: read or write");
DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|tcp");
DEFINE_string(device_name, "erdma_1",
"Device name to use, valid if protocol=rdma");
DEFINE_string(nic_priority_matrix, "",
"Path to RDMA NIC priority matrix file (Advanced)");
// python /workspace/Mooncake/mooncake-transfer-engine/scripts/register.py localhost test_nvmeof /workspace/sample
DEFINE_string(segment_id, "nvmeof/test_nvmeof", "Segment ID to access data");
static void *allocateMemoryPool(size_t size, int socket_id,
bool from_vram = false) {
return numa_alloc_onnode(size, socket_id);
}
static void freeMemoryPool(void *addr, size_t size) { numa_free(addr, size); }
class NVMeofTransportTest : public ::testing::Test {
public:
std::shared_ptr<mooncake::TransferMetadata> metadata_client;
void *addr = nullptr;
std::pair<std::string, uint16_t> hostname_port;
std::unique_ptr<mooncake::TransferEngine> engine;
const size_t ram_buffer_size = 1ull << 30;
Transport *xport;
std::string nic_priority_matrix;
void **args;
mooncake::Transport::SegmentID segment_id;
std::shared_ptr<TransferMetadata::SegmentDesc> segment_desc;
uint64_t remote_base;
protected:
void SetUp() override {
static int offset = 0;
google::InitGoogleLogging("NVMeofTransportTest");
FLAGS_logtostderr = 1;
// disable topology auto discovery for testing.
engine = std::make_unique<TransferEngine>(false);
hostname_port = parseHostNameWithPort(FLAGS_local_server_name);
engine->init(FLAGS_metadata_server, FLAGS_local_server_name.c_str(),
hostname_port.first.c_str(),
hostname_port.second + offset++);
xport = nullptr;
args = (void **)malloc(2 * sizeof(void *));
args[0] = nullptr;
xport = engine->installTransport("nvmeof", args);
ASSERT_NE(xport, nullptr);
addr = allocateMemoryPool(ram_buffer_size, 0, false);
int rc = engine->registerLocalMemory(addr, ram_buffer_size, "cpu:0");
ASSERT_EQ(rc, 0);
segment_id = engine->openSegment(FLAGS_segment_id.c_str());
bindToSocket(0);
segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id);
remote_base = 0;
}
void TearDown() override {
google::ShutdownGoogleLogging();
engine->unregisterLocalMemory(addr);
freeMemoryPool(addr, ram_buffer_size);
}
};
TEST_F(NVMeofTransportTest, MultiWrite) {
const size_t kDataLength = 4096000;
int times = 10;
while (times--) {
for (size_t offset = 0; offset < kDataLength; ++offset)
*((char *)(addr) + offset) = 'a' + lrand48() % 26;
auto batch_id = xport->allocateBatchID(1);
Status s;
TransferRequest entry;
entry.opcode = TransferRequest::WRITE;
entry.length = kDataLength;
entry.source = (uint8_t *)(addr);
entry.target_id = segment_id;
entry.target_offset = remote_base;
s = xport->submitTransfer(batch_id, {entry});
LOG_ASSERT(s.ok());
bool completed = false;
TransferStatus status;
while (!completed) {
Status s = xport->getTransferStatus(batch_id, 0, status);
ASSERT_EQ(s, Status::OK());
if (status.s == TransferStatusEnum::COMPLETED)
completed = true;
else if (status.s == TransferStatusEnum::FAILED) {
LOG(INFO) << "FAILED";
completed = true;
}
}
s = xport->freeBatchID(batch_id);
ASSERT_EQ(s, Status::OK());
}
}
TEST_F(NVMeofTransportTest, MultipleRead) {
const size_t kDataLength = 4096000;
int times = 10;
while (times--) {
for (size_t offset = 0; offset < kDataLength; ++offset)
*((char *)(addr) + offset) = 'a' + lrand48() % 26;
auto batch_id = xport->allocateBatchID(1);
Status s;
TransferRequest entry;
entry.opcode = TransferRequest::WRITE;
entry.length = kDataLength;
entry.source = (uint8_t *)(addr);
entry.target_id = segment_id;
entry.target_offset = remote_base;
s = xport->submitTransfer(batch_id, {entry});
LOG_ASSERT(s.ok());
bool completed = false;
TransferStatus status;
while (!completed) {
Status s = xport->getTransferStatus(batch_id, 0, status);
ASSERT_EQ(s, Status::OK());
if (status.s == TransferStatusEnum::COMPLETED)
completed = true;
else if (status.s == TransferStatusEnum::FAILED) {
LOG(INFO) << "FAILED";
completed = true;
}
}
s = engine->freeBatchID(batch_id);
ASSERT_EQ(s, Status::OK());
}
times = 10;
while (times--) {
auto batch_id = xport->allocateBatchID(1);
int ret = 0;
TransferRequest entry;
entry.opcode = TransferRequest::READ;
entry.length = kDataLength;
entry.source = (uint8_t *)(addr) + kDataLength;
entry.target_id = segment_id;
entry.target_offset = remote_base;
Status s;
s = xport->submitTransfer(batch_id, {entry});
ASSERT_EQ(s, Status::OK());
bool completed = false;
TransferStatus status;
while (!completed) {
Status s = xport->getTransferStatus(batch_id, 0, status);
ASSERT_EQ(s, Status::OK());
if (status.s == TransferStatusEnum::COMPLETED)
completed = true;
else if (status.s == TransferStatusEnum::FAILED) {
completed = true;
}
}
s = xport->freeBatchID(batch_id);
ASSERT_EQ(s, Status::OK());
ret = memcmp((uint8_t *)(addr), (uint8_t *)(addr) + kDataLength,
kDataLength);
ASSERT_EQ(ret, 0);
}
engine->unregisterLocalMemory(addr);
freeMemoryPool(addr, ram_buffer_size);
}
} // namespace mooncake
int main(int argc, char **argv) {
gflags::ParseCommandLineFlags(&argc, &argv, false);
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}