mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

implement cpp methods for p2p connection (#1539) 

* implement cpp methods for p2p connection

* Apply suggestions from code review

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* lint

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Donghun Lee 이동훈 2026-02-12 13:27:50 +09:00 committed by GitHub
parent 45305495a3
commit dc432620cf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 196 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
mooncake-transfer-engine/tent/include/tent/transfer_engine.h
View File

@ -81,11 +81,40 @@ struct tent_notifi_record {
char msg[4096];
};
typedef struct tent_notifi_record tent_notifi_record_t;
struct tent_notifi_info {
int num_records;
struct tent_notifi_record* records;
};
typedef struct tent_notifi_info tent_notifi_info;
#define PERM_LOCAL_READ_WRITE (0)
#define PERM_GLOBAL_READ_ONLY (1)
#define PERM_GLOBAL_READ_WRITE (2)
#define TRANSPORT_RDMA (0)
#define TRANSPORT_MNNVL (1)
#define TRANSPORT_SHM (2)
#define TRANSPORT_NVLINK (3)
#define TRANSPORT_GDS (4)
#define TRANSPORT_IOURING (5)
#define TRANSPORT_TCP (6)
#define TRANSPORT_ASCEND_DIRECT (7)
#define TRANSPORT_UNSPEC (8)
struct tent_memory_options {
char location[64];
int permission; /* PERM_LOCAL_READ_WRITE, etc. */
int transport_type; /* TRANSPORT_RDMA, etc. */
char shm_path[256];
size_t shm_offset;
int internal; /* 0 = false, nonzero = true */
};
typedef struct tent_memory_options tent_memory_options_t;
void tent_load_config_from_file(const char* path);
void tent_set_config(const char* key, const char* value);
@ -142,6 +171,30 @@ int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id,
int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id,
tent_status_t* status);
int tent_available(tent_engine_t engine);
int tent_register_memory_with_perm(tent_engine_t engine, void* addr,
size_t size, int permission);
int tent_register_memory_batch(tent_engine_t engine, void** addrs,
size_t* sizes, size_t count, int permission);
int tent_unregister_memory_batch(tent_engine_t engine, void** addrs,
size_t* sizes, size_t count);
int tent_allocate_memory_ex(tent_engine_t engine, void** addr, size_t size,
tent_memory_options_t* opts);
int tent_register_memory_ex(tent_engine_t engine, void* addr, size_t size,
tent_memory_options_t* opts);
int tent_register_memory_batch_ex(tent_engine_t engine, void** addrs,
size_t* sizes, size_t count,
tent_memory_options_t* opts);
int tent_task_status_list(tent_engine_t engine, tent_batch_id_t batch_id,
tent_status_t* statuses, size_t* count);
#ifdef __cplusplus
}
#endif // __cplusplus

143
mooncake-transfer-engine/tent/src/transfer_engine_c.cpp
View File

@ -330,3 +330,146 @@ int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id,
xfer_status->transferred_bytes = internal_status.transferred_bytes;
return 0;
}
// =========================================================================
// Helper: convert C tent_memory_options_t to C++ MemoryOptions
// =========================================================================
static mooncake::tent::MemoryOptions convert_options(
const tent_memory_options_t* opts) {
mooncake::tent::MemoryOptions options;
if (opts->location[0] != '\0') {
options.location = opts->location;
}
options.perm = (mooncake::tent::Permission)opts->permission;
options.type = (mooncake::tent::TransportType)opts->transport_type;
if (opts->shm_path[0] != '\0') {
options.shm_path = opts->shm_path;
}
options.shm_offset = opts->shm_offset;
options.internal = opts->internal != 0;
return options;
}
// =========================================================================
// New C API functions for Python parity
// =========================================================================
int tent_available(tent_engine_t engine) {
if (!engine) return 0;
return CAST(engine)->available() ? 1 : 0;
}
int tent_register_memory_with_perm(tent_engine_t engine, void* addr,
size_t size, int permission) {
CHECK_POINTER(engine);
CHECK_POINTER(addr);
auto perm = static_cast<mooncake::tent::Permission>(permission);
auto status = CAST(engine)->registerLocalMemory({addr}, {size}, perm);
if (!status.ok()) {
LOG(ERROR) << "tent_register_memory_with_perm: " << status.ToString();
return -1;
}
return 0;
}
int tent_register_memory_batch(tent_engine_t engine, void** addrs,
size_t* sizes, size_t count, int permission) {
CHECK_POINTER(engine);
CHECK_POINTER(addrs);
CHECK_POINTER(sizes);
std::vector<void*> addr_list(addrs, addrs + count);
std::vector<size_t> size_list(sizes, sizes + count);
auto perm = static_cast<mooncake::tent::Permission>(permission);
auto status = CAST(engine)->registerLocalMemory(addr_list, size_list, perm);
if (!status.ok()) {
LOG(ERROR) << "tent_register_memory_batch: " << status.ToString();
return -1;
}
return 0;
}
int tent_unregister_memory_batch(tent_engine_t engine, void** addrs,
size_t* sizes, size_t count) {
CHECK_POINTER(engine);
CHECK_POINTER(addrs);
std::vector<void*> addr_list(addrs, addrs + count);
std::vector<size_t> size_list;
if (sizes) {
size_list.assign(sizes, sizes + count);
}
auto status = CAST(engine)->unregisterLocalMemory(addr_list, size_list);
if (!status.ok()) {
LOG(ERROR) << "tent_unregister_memory_batch: " << status.ToString();
return -1;
}
return 0;
}
int tent_allocate_memory_ex(tent_engine_t engine, void** addr, size_t size,
tent_memory_options_t* opts) {
CHECK_POINTER(engine);
CHECK_POINTER(addr);
CHECK_POINTER(opts);
auto options = convert_options(opts);
auto status = CAST(engine)->allocateLocalMemory(addr, size, options);
if (!status.ok()) {
LOG(ERROR) << "tent_allocate_memory_ex: " << status.ToString();
return -1;
}
return 0;
}
int tent_register_memory_ex(tent_engine_t engine, void* addr, size_t size,
tent_memory_options_t* opts) {
CHECK_POINTER(engine);
CHECK_POINTER(addr);
CHECK_POINTER(opts);
auto options = convert_options(opts);
auto status = CAST(engine)->registerLocalMemory({addr}, {size}, options);
if (!status.ok()) {
LOG(ERROR) << "tent_register_memory_ex: " << status.ToString();
return -1;
}
return 0;
}
int tent_register_memory_batch_ex(tent_engine_t engine, void** addrs,
size_t* sizes, size_t count,
tent_memory_options_t* opts) {
CHECK_POINTER(engine);
CHECK_POINTER(addrs);
CHECK_POINTER(sizes);
CHECK_POINTER(opts);
std::vector<void*> addr_list(addrs, addrs + count);
std::vector<size_t> size_list(sizes, sizes + count);
auto options = convert_options(opts);
auto status =
CAST(engine)->registerLocalMemory(addr_list, size_list, options);
if (!status.ok()) {
LOG(ERROR) << "tent_register_memory_batch_ex: " << status.ToString();
return -1;
}
return 0;
}
int tent_task_status_list(tent_engine_t engine, tent_batch_id_t batch_id,
tent_status_t* statuses, size_t* count) {
CHECK_POINTER(engine);
CHECK_POINTER(batch_id);
CHECK_POINTER(statuses);
CHECK_POINTER(count);
std::vector<mooncake::tent::TransferStatus> status_list;
auto status = CAST(engine)->getTransferStatus(batch_id, status_list);
if (!status.ok()) {
LOG(ERROR) << "tent_task_status_list: " << status.ToString();
return -1;
}
size_t to_copy = std::min(status_list.size(), *count);
for (size_t i = 0; i < to_copy; ++i) {
statuses[i].status = (int)status_list[i].s;
statuses[i].transferred_bytes = status_list[i].transferred_bytes;
}
*count = status_list.size();
return 0;
}