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

[MooncakeAdaptor] reconstruct the adaptor arch (#232)

This commit is contained in:
Teng Ma 2025-04-11 17:54:02 +08:00 committed by GitHub
parent 033ff82d84
commit abad4f758b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 589 additions and 501 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -191,3 +191,5 @@ cmake-build-minsizerel
cmake-build
libetcd_wrapper.h
mooncake-wheel/mooncake/mooncake_master

23
mooncake-integration/CMakeLists.txt
View File

@ -44,17 +44,6 @@ target_link_libraries(mooncake_vllm_adaptor PUBLIC
message("${PYTHON_SYS_PATH}")
install(TARGETS mooncake_vllm_adaptor DESTINATION ${PYTHON_SYS_PATH}/)
pybind11_add_module(mooncake_sglang_adaptor ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
sglang/sglang_adaptor.cpp
)
target_link_libraries(mooncake_sglang_adaptor PUBLIC
transfer_engine
glog
gflags
)
message("${PYTHON_SYS_PATH}")
install(TARGETS mooncake_sglang_adaptor DESTINATION ${PYTHON_SYS_PATH}/)
set(PYTHON_PACKAGE_NAME "mooncake")
pybind11_add_module(engine ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
transfer_engine/transfer_engine_py.cpp
@ -64,6 +53,16 @@ target_link_libraries(engine PUBLIC
glog
gflags
)
pybind11_add_module(store ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
store/store_py.cpp
)
target_link_libraries(store PUBLIC
transfer_engine
glog
gflags
mooncake_store
cachelib_memory_allocator
)
message("${PYTHON_SYS_PATH}")
file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/__init__.py
@ -83,5 +82,5 @@ install(
execute_process(COMMAND chmod 766 \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}/__init__.py\")
"
)
install(TARGETS store DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
install(TARGETS engine DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})

327
mooncake-integration/sglang/sglang_adaptor.cpp
View File

@ -1,327 +0,0 @@
// 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 "sglang_adaptor.h"
#include <cassert>
#ifdef USE_CUDA
#include <bits/stdint-uintn.h>
#include <cuda_runtime.h>
#endif
SGLangAdaptor::SGLangAdaptor() {}
SGLangAdaptor::~SGLangAdaptor() {
for (auto &handle : handle_map_) engine_->closeSegment(handle.second);
handle_map_.clear();
engine_.reset();
for (auto &buffer : buffer_list_) free(buffer);
buffer_list_.clear();
for (auto &buffer : large_buffer_list_) free(buffer);
large_buffer_list_.clear();
}
std::string formatDeviceNames(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);
}
std::string formatted;
for (size_t i = 0; i < tokens.size(); ++i) {
formatted += "\"" + tokens[i] + "\"";
if (i < tokens.size() - 1) {
formatted += ",";
}
}
return formatted;
}
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 {
domain = conn_string;
}
result.first = proto;
result.second = domain;
return result;
}
int SGLangAdaptor::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 SGLangAdaptor::initializeExt(const char *local_hostname,
const char *metadata_server,
const char *protocol, const char *device_name,
const char *metadata_type) {
std::string conn_string = metadata_server;
if (conn_string.find("://") == std::string::npos)
conn_string =
std::string(metadata_type) + "://" + std::string(metadata_server);
// TODO: remove `false` in the feature, it's for keep same API in SGLang.
engine_ = std::make_unique<TransferEngine>(false);
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;
}
xport_ = nullptr;
if (strcmp(protocol, "rdma") == 0) {
auto device_names = formatDeviceNames(device_name);
std::string nic_priority_matrix =
"{\"cpu:0\": [[" + device_names + "], []],"
"\"cuda:0\": [[" + device_names + "], []]}";
void **args = (void **)malloc(2 * sizeof(void *));
args[0] = (void *)nic_priority_matrix.c_str();
args[1] = nullptr;
xport_ = engine_->installTransport("rdma", args);
} else if (strcmp(protocol, "tcp") == 0) {
xport_ = engine_->installTransport("tcp", nullptr);
} else {
LOG(ERROR) << "Unsupported protocol";
return -1;
}
if (!xport_) return -1;
free_list_.resize(kSlabSizeKBTabLen);
doBuddyAllocate(kMaxClassId);
return 0;
}
char *SGLangAdaptor::allocateRawBuffer(size_t capacity) {
auto buffer = malloc(capacity);
if (!buffer) return nullptr;
int ret = engine_->registerLocalMemory(buffer, capacity, "cpu:0");
if (ret) {
free(buffer);
return nullptr;
}
return (char *)buffer;
}
int SGLangAdaptor::findClassId(size_t size) {
if (size > 1024ull * kSlabSizeKB[kMaxClassId]) return -1;
for (int i = kMaxClassId - 2; i >= 0; --i)
if (size > 1024ull * kSlabSizeKB[i]) return i + 1;
return 0;
}
int SGLangAdaptor::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 SGLangAdaptor::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 SGLangAdaptor::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);
free(buffer);
return 0;
}
free_list_[class_id].push(buffer);
return 0;
}
int SGLangAdaptor::transferSync(const char *target_hostname, uintptr_t buffer,
uintptr_t peer_buffer_address, size_t length) {
Transport::SegmentHandle handle;
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::READ;
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;
TransferStatus status;
while (true) {
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);
return -1;
}
}
}
int SGLangAdaptor::transferSyncExt(const char *target_hostname,
uintptr_t buffer,
uintptr_t peer_buffer_address, size_t length,
TransferOpcode opcode) {
Transport::SegmentHandle handle;
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;
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;
Status s = engine_->submitTransfer(batch_id, {entry});
if (!s.ok()) return -1;
TransferStatus status;
while (true) {
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);
return -1;
}
}
}
int SGLangAdaptor::expRegisterMemory(uintptr_t buffer_addr, size_t capacity) {
char *buffer = reinterpret_cast<char *>(buffer_addr);
std::string location = "cpu:0";
#ifdef USE_CUDA
// check pointer on GPU
cudaPointerAttributes attributes;
cudaPointerGetAttributes(&attributes, buffer);
if (attributes.type == cudaMemoryTypeDevice) {
location = "cuda:0";
}
#endif
return engine_->registerLocalMemory(buffer, capacity, location);
}
int SGLangAdaptor::expUnregisterMemory(uintptr_t buffer_addr) {
char *buffer = reinterpret_cast<char *>(buffer_addr);
return engine_->unregisterLocalMemory(buffer);
}
uintptr_t SGLangAdaptor::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(mooncake_sglang_adaptor, m) {
py::enum_<SGLangAdaptor::TransferOpcode> transfer_opcode(
m, "TransferOpcode", py::arithmetic());
transfer_opcode.value("READ", SGLangAdaptor::TransferOpcode::READ)
.value("WRITE", SGLangAdaptor::TransferOpcode::WRITE)
.export_values();
auto adaptor_cls =
py::class_<SGLangAdaptor>(m, "TransferEngine")
.def(py::init<>())
.def("initialize", &SGLangAdaptor::initialize)
.def("initializeExt", &SGLangAdaptor::initializeExt)
.def("allocateManagedBuffer", &SGLangAdaptor::allocateManagedBuffer)
.def("freeManagedBuffer", &SGLangAdaptor::freeManagedBuffer)
.def("transferSyncExt", &SGLangAdaptor::transferSyncExt)
.def("transferSync", &SGLangAdaptor::transferSync)
.def("writeBytesToBuffer", &SGLangAdaptor::writeBytesToBuffer)
.def("readBytesFromBuffer", &SGLangAdaptor::readBytesFromBuffer)
.def("expRegisterMemory", &SGLangAdaptor::expRegisterMemory)
.def("expUnregisterMemory", &SGLangAdaptor::expUnregisterMemory)
.def("getFirstBufferAddress",
&SGLangAdaptor::getFirstBufferAddress);
adaptor_cls.attr("TransferOpcode") = transfer_opcode;
}

106
mooncake-integration/sglang/sglang_adaptor.h
View File

@ -1,106 +0,0 @@
// 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 <pybind11/pybind11.h>
#include <sys/time.h>
#include <cstdlib>
#include <fstream>
#include <iomanip>
#include <memory>
#include <stack>
#include <vector>
#include "common/base/status.h"
#include "transfer_engine.h"
#include "transport/rdma_transport/rdma_transport.h"
#include "transport/transport.h"
using namespace mooncake;
const static size_t kDefaultBufferCapacity = 2ull * 1024 * 1024 * 1024;
const static size_t kSlabSizeKBTabLen = 16;
const static size_t kMaxClassId = kSlabSizeKBTabLen - 1;
const static size_t kSlabSizeKB[] = {
8, 16, 32, 64, 128, 256,
512, 1024, 2 * 1024, 4 * 1024, 8 * 1024, 16 * 1024,
32 * 1024, 64 * 1024, 128 * 1024, 256 * 1024};
class SGLangAdaptor {
public:
enum class TransferOpcode {
READ = 0,
WRITE = 1
};
public:
SGLangAdaptor();
~SGLangAdaptor();
int initialize(const char *local_hostname, const char *metadata_server,
const char *protocol, const char *device_name);
int initializeExt(const char *local_hostname, const char *metadata_server,
const char *protocol, const char *device_name,
const char *metadata_type);
uintptr_t allocateManagedBuffer(size_t length);
int freeManagedBuffer(uintptr_t user_tensor, size_t length);
int transferSync(const char *target_hostname, uintptr_t buffer,
uintptr_t peer_buffer_address, size_t length);
int transferSyncExt(const char *target_hostname, uintptr_t buffer,
uintptr_t peer_buffer_address, size_t length, TransferOpcode opcode);
uintptr_t getFirstBufferAddress(const std::string &segment_name);
int writeBytesToBuffer(uintptr_t dest_address, char *src_ptr,
size_t length) {
memcpy((void *)dest_address, (void *)src_ptr, length);
return 0;
}
pybind11::bytes readBytesFromBuffer(uintptr_t source_address,
size_t length) {
return pybind11::bytes(
static_cast<const char *>(reinterpret_cast<void *>(source_address)),
length);
}
// FOR EXPERIMENT ONLY
int expRegisterMemory(uintptr_t buffer_addr, size_t capacity);
// must be called before SGLangAdaptor::~SGLangAdaptor()
int expUnregisterMemory(uintptr_t buffer_addr);
private:
char *allocateRawBuffer(size_t capacity);
int findClassId(size_t size);
int doBuddyAllocate(int class_id);
private:
std::shared_ptr<TransferEngine> engine_;
Transport *xport_;
std::mutex mutex_;
std::vector<std::stack<char *>> free_list_;
std::vector<char *> buffer_list_;
std::unordered_set<char *> large_buffer_list_;
std::unordered_map<std::string, Transport::SegmentHandle> handle_map_;
};

406
mooncake-integration/store/store_py.cpp Normal file
View File

@ -0,0 +1,406 @@
#include "store_py.h"
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
#include <cstdlib> // for atexit
#include <random>
#include "types.h"
using namespace mooncake;
// ResourceTracker implementation using singleton pattern
ResourceTracker &ResourceTracker::getInstance() {
static ResourceTracker instance;
return instance;
}
ResourceTracker::ResourceTracker() {
// Set up signal handlers
struct sigaction sa;
sa.sa_handler = signalHandler;
sigemptyset(&sa.sa_mask);
sa.sa_flags = 0;
// Register for common termination signals
sigaction(SIGINT, &sa, nullptr); // Ctrl+C
sigaction(SIGTERM, &sa, nullptr); // kill command
sigaction(SIGHUP, &sa, nullptr); // Terminal closed
// Register exit handler
std::atexit(exitHandler);
}
ResourceTracker::~ResourceTracker() {
// Cleanup is handled by exitHandler
}
void ResourceTracker::registerInstance(DistributedObjectStore *instance) {
std::lock_guard<std::mutex> lock(mutex_);
instances_.insert(instance);
}
void ResourceTracker::unregisterInstance(DistributedObjectStore *instance) {
std::lock_guard<std::mutex> lock(mutex_);
instances_.erase(instance);
}
void ResourceTracker::cleanupAllResources() {
std::lock_guard<std::mutex> lock(mutex_);
// Perform cleanup outside the lock to avoid potential deadlocks
for (void *instance : instances_) {
DistributedObjectStore *store =
static_cast<DistributedObjectStore *>(instance);
if (store) {
LOG(INFO) << "Cleaning up DistributedObjectStore instance";
store->tearDownAll();
}
}
}
void ResourceTracker::signalHandler(int signal) {
LOG(INFO) << "Received signal " << signal << ", cleaning up resources";
getInstance().cleanupAllResources();
// Re-raise the signal with default handler to allow normal termination
struct sigaction sa;
sa.sa_handler = SIG_DFL;
sigemptyset(&sa.sa_mask);
sa.sa_flags = 0;
sigaction(signal, &sa, nullptr);
raise(signal);
}
void ResourceTracker::exitHandler() {
LOG(INFO) << "Process exiting, cleaning up resources";
getInstance().cleanupAllResources();
}
static bool isPortAvailable(int port) {
int sock = socket(AF_INET, SOCK_STREAM, 0);
if (sock < 0) return false;
int opt = 1;
setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
struct sockaddr_in addr;
memset(&addr, 0, sizeof(addr));
addr.sin_family = AF_INET;
addr.sin_addr.s_addr = INADDR_ANY;
addr.sin_port = htons(port);
bool available = (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) == 0);
close(sock);
return available;
}
// Get a random available port between min_port and max_port
static int getRandomAvailablePort(int min_port = 12300, int max_port = 14300) {
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<> dis(min_port, max_port);
for (int attempts = 0; attempts < 10; attempts++) {
int port = dis(gen);
std::cout << "port is " << port << std::endl;
if (isPortAvailable(port)) {
return port;
}
}
return -1; // Failed to find available port
}
DistributedObjectStore::DistributedObjectStore() {
// Register this instance with the global tracker
easylog::set_min_severity(easylog::Severity::WARN);
ResourceTracker::getInstance().registerInstance(this);
}
DistributedObjectStore::~DistributedObjectStore() {
// Unregister from the tracker before cleanup
ResourceTracker::getInstance().unregisterInstance(this);
if (client_ && segment_ptr_) {
// Try to unmount the segment using saved local_hostname
ErrorCode rc = client_->UnInit();
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount segment in destructor: "
<< toString(rc);
}
// The unique_ptr will automatically free the memory when reset
segment_ptr_.reset();
client_.reset();
}
}
int DistributedObjectStore::setup(const std::string &local_hostname,
const std::string &metadata_server,
size_t global_segment_size,
size_t local_buffer_size,
const std::string &protocol,
const std::string &rdma_devices,
const std::string &master_server_addr) {
this->protocol = protocol;
// Remove port if hostname already contains one
std::string hostname = local_hostname;
size_t colon_pos = hostname.find(":");
if (colon_pos == std::string::npos) {
// Get a random available port
int port = getRandomAvailablePort();
if (port < 0) {
LOG(ERROR) << "Failed to find available port";
return 1;
}
// Combine hostname with port
this->local_hostname = hostname + ":" + std::to_string(port);
} else {
this->local_hostname = local_hostname;
}
client_ = std::make_unique<mooncake::Client>();
void **args = (protocol == "rdma") ? rdma_args(rdma_devices) : nullptr;
ErrorCode rc = client_->Init(this->local_hostname, metadata_server,
protocol, args, master_server_addr);
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to initialize client: " << toString(rc);
return 1;
}
client_buffer_allocator_ =
std::make_unique<SimpleAllocator>(local_buffer_size);
rc = client_->RegisterLocalMemory(client_buffer_allocator_->getBase(),
local_buffer_size, "cpu:0", false, false);
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to register local memory: " << toString(rc);
return 1;
}
void *ptr = allocate_buffer_allocator_memory(global_segment_size);
if (!ptr) {
LOG(ERROR) << "Failed to allocate segment memory";
return 1;
}
segment_ptr_.reset(ptr);
rc = client_->MountSegment(this->local_hostname, segment_ptr_.get(),
global_segment_size);
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to mount segment: " << toString(rc);
return 1;
}
return 0;
}
int DistributedObjectStore::initAll(const std::string &protocol_,
const std::string &device_name,
size_t mount_segment_size) {
if (client_) {
LOG(ERROR) << "Client is already initialized";
return 1;
}
uint64_t buffer_allocator_size = 1024 * 1024 * 1024;
return setup("localhost:12345", "127.0.0.1:2379", mount_segment_size,
buffer_allocator_size, protocol_, device_name);
}
int DistributedObjectStore::allocateSlices(std::vector<Slice> &slices,
const std::string &value) {
uint64_t offset = 0;
while (offset < value.size()) {
auto chunk_size = std::min(value.size() - offset, kMaxSliceSize);
auto ptr = client_buffer_allocator_->allocate(chunk_size);
if (!ptr) {
// Deallocate any previously allocated slices
for (auto &slice : slices) {
client_buffer_allocator_->deallocate(slice.ptr, slice.size);
}
slices.clear();
return 1;
}
memcpy(ptr, value.data() + offset, chunk_size);
slices.emplace_back(Slice{ptr, chunk_size});
offset += chunk_size;
}
return 0;
}
int DistributedObjectStore::allocateSlices(
std::vector<mooncake::Slice> &slices,
const mooncake::Client::ObjectInfo &object_info, uint64_t &length) {
length = 0;
if (object_info.replica_list.empty()) return -1;
auto &replica = object_info.replica_list[0];
for (auto &handle : replica.buffer_descriptors) {
auto chunk_size = handle.size_;
assert(chunk_size <= kMaxSliceSize);
auto ptr = client_buffer_allocator_->allocate(chunk_size);
if (!ptr) return 1;
slices.emplace_back(Slice{ptr, chunk_size});
length += chunk_size;
}
return 0;
}
char *DistributedObjectStore::exportSlices(
const std::vector<mooncake::Slice> &slices, uint64_t length) {
char *buf = new char[length + 1];
buf[length] = '\0';
uint64_t offset = 0;
for (auto slice : slices) {
memcpy(buf + offset, slice.ptr, slice.size);
offset += slice.size;
}
return buf;
}
int DistributedObjectStore::freeSlices(
const std::vector<mooncake::Slice> &slices) {
for (auto slice : slices) {
client_buffer_allocator_->deallocate(slice.ptr, slice.size);
}
return 0;
}
int DistributedObjectStore::tearDownAll() {
if (!client_) {
LOG(ERROR) << "Client is not initialized";
return 1;
}
ErrorCode rc = client_->UnInit();
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount segment: " << toString(rc);
return 1;
}
client_.reset();
client_buffer_allocator_.reset();
segment_ptr_.reset();
local_hostname = "";
device_name = "";
protocol = "";
return 0;
}
int DistributedObjectStore::put(const std::string &key,
const std::string &value) {
if (!client_) {
LOG(ERROR) << "Client is not initialized";
return 1;
}
ReplicateConfig config;
config.replica_num = 1; // TODO
std::vector<Slice> slices;
int ret = allocateSlices(slices, value);
if (ret) return ret;
ErrorCode error_code = client_->Put(std::string(key), slices, config);
freeSlices(slices);
if (error_code != ErrorCode::OK) return toInt(error_code);
return 0;
}
pybind11::bytes DistributedObjectStore::get(const std::string &key) {
if (!client_) {
LOG(ERROR) << "Client is not initialized";
return pybind11::bytes("\0", 0);
}
mooncake::Client::ObjectInfo object_info;
std::vector<Slice> slices;
const auto kNullString = pybind11::bytes("\0", 0);
ErrorCode error_code = client_->Query(key, object_info);
if (error_code != ErrorCode::OK) return kNullString;
uint64_t str_length = 0;
int ret = allocateSlices(slices, object_info, str_length);
if (ret) return kNullString;
error_code = client_->Get(key, object_info, slices);
if (error_code != ErrorCode::OK) {
freeSlices(slices);
return kNullString;
}
if (slices.size() == 1 && slices[0].size == str_length) {
auto result = pybind11::bytes((char *)slices[0].ptr, str_length);
freeSlices(slices);
return result;
}
const char *str = exportSlices(slices, str_length);
freeSlices(slices);
if (!str) return kNullString;
pybind11::bytes result(str, str_length);
delete[] str;
return result;
}
int DistributedObjectStore::remove(const std::string &key) {
if (!client_) {
LOG(ERROR) << "Client is not initialized";
return 1;
}
ErrorCode error_code = client_->Remove(key);
if (error_code != ErrorCode::OK) return toInt(error_code);
return 0;
}
int DistributedObjectStore::isExist(const std::string &key) {
if (!client_) {
LOG(ERROR) << "Client is not initialized";
return -1;
}
ErrorCode err = client_->IsExist(key);
if (err == ErrorCode::OK) return 1; // Yes
if (err == ErrorCode::OBJECT_NOT_FOUND) return 0; // No
return toInt(err); // Error
}
int64_t DistributedObjectStore::getSize(const std::string &key) {
if (!client_) {
LOG(ERROR) << "Client is not initialized";
return -1;
}
mooncake::Client::ObjectInfo object_info;
ErrorCode error_code = client_->Query(key, object_info);
if (error_code != ErrorCode::OK) {
return toInt(error_code);
}
// Calculate total size from all replicas' handles
int64_t total_size = 0;
if (!object_info.replica_list.empty()) {
auto &replica = object_info.replica_list[0];
for (auto &handle : replica.buffer_descriptors) {
total_size += handle.size_;
}
} else {
LOG(ERROR) << "Internal error: object_info.replica_list_size() is 0";
return -1; // Internal error
}
return total_size;
}
namespace py = pybind11;
PYBIND11_MODULE(store, m) {
py::class_<DistributedObjectStore>(m, "MooncakeDistributedStore")
.def(py::init<>())
.def("setup", &DistributedObjectStore::setup)
.def("init_all", &DistributedObjectStore::initAll)
.def("get", &DistributedObjectStore::get)
.def("put", &DistributedObjectStore::put)
.def("remove", &DistributedObjectStore::remove)
.def("is_exist", &DistributedObjectStore::isExist)
.def("close", &DistributedObjectStore::tearDownAll)
.def("get_size", &DistributedObjectStore::getSize);
}

117
mooncake-integration/store/store_py.h Normal file
View File

@ -0,0 +1,117 @@
#pragma once
#include <pybind11/pybind11.h>
#include <csignal>
#include <mutex>
#include <string>
#include <unordered_set>
#include "allocator.h"
#include "client.h"
#include "utils.h"
class DistributedObjectStore;
// Global resource tracker to handle cleanup on abnormal termination
class ResourceTracker {
public:
// Get the singleton instance
static ResourceTracker &getInstance();
// Register a DistributedObjectStore instance for cleanup
void registerInstance(DistributedObjectStore *instance);
// Unregister a DistributedObjectStore instance
void unregisterInstance(DistributedObjectStore *instance);
private:
ResourceTracker();
~ResourceTracker();
// Prevent copying
ResourceTracker(const ResourceTracker &) = delete;
ResourceTracker &operator=(const ResourceTracker &) = delete;
// Cleanup all registered resources
void cleanupAllResources();
// Signal handler function
static void signalHandler(int signal);
// Exit handler function
static void exitHandler();
std::mutex mutex_;
std::unordered_set<DistributedObjectStore *> instances_;
};
class DistributedObjectStore {
public:
DistributedObjectStore();
~DistributedObjectStore();
int setup(const std::string &local_hostname,
const std::string &metadata_server,
size_t global_segment_size = 1024 * 1024 * 16,
size_t local_buffer_size = 1024 * 1024 * 16,
const std::string &protocol = "tcp",
const std::string &rdma_devices = "",
const std::string &master_server_addr = "127.0.0.1:50051");
int initAll(const std::string &protocol, const std::string &device_name,
size_t mount_segment_size = 1024 * 1024 * 16); // Default 16MB
int put(const std::string &key, const std::string &value);
pybind11::bytes get(const std::string &key);
int remove(const std::string &key);
int tearDownAll();
/**
* @brief Check if an object exists
* @param key Key to check
* @return 1 if exists, 0 if not exists, -1 if error
*/
int isExist(const std::string &key);
/**
* @brief Get the size of an object
* @param key Key of the object
* @return Size of the object in bytes, or -1 if error or object doesn't
* exist
*/
int64_t getSize(const std::string &key);
private:
int allocateSlices(std::vector<mooncake::Slice> &slices,
const std::string &value);
int allocateSlices(std::vector<mooncake::Slice> &slices,
const mooncake::Client::ObjectInfo &object_info,
uint64_t &length);
char *exportSlices(const std::vector<mooncake::Slice> &slices,
uint64_t length);
int freeSlices(const std::vector<mooncake::Slice> &slices);
public:
std::unique_ptr<mooncake::Client> client_ = nullptr;
std::unique_ptr<mooncake::SimpleAllocator> client_buffer_allocator_ =
nullptr;
struct SegmentDeleter {
void operator()(void *ptr) {
if (ptr) {
free(ptr);
}
}
};
std::unique_ptr<void, SegmentDeleter> segment_ptr_;
std::string protocol;
std::string device_name;
std::string local_hostname;
};

2
mooncake-store/tests/stress_cluster_benchmark.py
View File

@ -1,7 +1,7 @@
import os
import time
import random
from mooncake_vllm_adaptor import MooncakeDistributedStore
from mooncake.store import MooncakeDistributedStore
# How to test
# 1. Config the following parameters in `setup`, notice all IP addresses

2
mooncake-store/tests/test_distributed_object_store.py
View File

@ -3,7 +3,7 @@ import os
import time
import threading
import random
from mooncake_vllm_adaptor import MooncakeDistributedStore
from mooncake.store import MooncakeDistributedStore
def get_client(store):

4
mooncake-wheel/mooncake/__init__.py
View File

@ -1,6 +1,2 @@
# Import for backward compatibility
from .mooncake_vllm_adaptor import MooncakeDistributedStore
from .mooncake_vllm_adaptor import mooncake_vllm_adaptor
# Import transfer module
from . import transfer

6
mooncake-wheel/mooncake/transfer/__init__.py
View File

@ -1,6 +0,0 @@
# mooncake.transfer module
# Import all symbols from the engine module to make them directly accessible
from ..engine import TransferEngine, TransferOpcode
# Export the main class and enum for direct access
__all__ = ['TransferEngine', 'TransferOpcode']

3
mooncake-wheel/setup.py
View File

@ -28,9 +28,6 @@ setup(
"*.so",
"mooncake_master",
],
"mooncake.transfer": [
"*.so",
],
},
include_package_data=True,
zip_safe=False,

18
mooncake-wheel/tests/test_distributed_object_store.py
View File

@ -3,7 +3,7 @@ import os
import time
import threading
import random
from mooncake import MooncakeDistributedStore
from mooncake.store import MooncakeDistributedStore
def get_client(store):
@ -68,7 +68,7 @@ class TestDistributedObjectStore(unittest.TestCase):
self.assertEqual(self.store.put(key, test_data), 0)
# Verify data through Get operation
self.assertEqual(self.store.getSize(key), len(test_data))
self.assertEqual(self.store.get_size(key), len(test_data))
retrieved_data = self.store.get(key)
self.assertEqual(retrieved_data, test_data)
@ -79,7 +79,7 @@ class TestDistributedObjectStore(unittest.TestCase):
self.assertEqual(self.store.remove(key), 0)
# Get after remove should return empty bytes
self.assertLess(self.store.getSize(key), 0)
self.assertLess(self.store.get_size(key), 0)
empty_data = self.store.get(key)
self.assertEqual(empty_data, b"")
@ -88,18 +88,18 @@ class TestDistributedObjectStore(unittest.TestCase):
key_2 = "test_exist_key"
# Should not exist initially
self.assertLess(self.store.getSize(key_2), 0)
self.assertEqual(self.store.isExist(key_2), 0)
self.assertLess(self.store.get_size(key_2), 0)
self.assertEqual(self.store.is_exist(key_2), 0)
# Should exist after put
self.assertEqual(self.store.put(key_2, test_data_2), 0)
self.assertEqual(self.store.isExist(key_2), 1)
self.assertEqual(self.store.getSize(key_2), len(test_data_2))
self.assertEqual(self.store.is_exist(key_2), 1)
self.assertEqual(self.store.get_size(key_2), len(test_data_2))
# Should not exist after remove
self.assertEqual(self.store.remove(key_2), 0)
self.assertLess(self.store.getSize(key_2), 0)
self.assertEqual(self.store.isExist(key_2), 0)
self.assertLess(self.store.get_size(key_2), 0)
self.assertEqual(self.store.is_exist(key_2), 0)
def test_concurrent_stress_with_barrier(self):
"""Test concurrent Put/Get operations with multiple threads using barrier."""

37
mooncake-wheel/tests/test_import_structure.py
View File

@ -2,34 +2,43 @@
import unittest
class TestImportStructure(unittest.TestCase):
def test_backward_compatibility(self):
"""Test that the old import style still works."""
from mooncake import MooncakeDistributedStore
from mooncake import mooncake_vllm_adaptor
# Just verify we can create instances
store = MooncakeDistributedStore()
adaptor = mooncake_vllm_adaptor()
# Restart this test when it is finished
# def test_backward_compatibility(self):
# """Test that the old import style still works."""
# from mooncake_vllm_adaptor import MooncakeDistributedStore
# import mooncake_vllm_adaptor
self.assertIsNotNone(store)
self.assertIsNotNone(adaptor)
# # Just verify we can create instances
# store = MooncakeDistributedStore()
# adaptor = mooncake_vllm_adaptor.mooncake_vllm_adaptor()
# self.assertIsNotNone(store)
# self.assertIsNotNone(adaptor)
def test_new_import_structure(self):
"""Test that the new import structure works."""
import mooncake.transfer
import mooncake.engine
# Verify the module exists
self.assertIsNotNone(mooncake.transfer)
self.assertIsNotNone(mooncake.engine)
# Verify direct access to TransferEngine
self.assertIsNotNone(mooncake.transfer.TransferEngine)
self.assertIsNotNone(mooncake.engine.TransferEngine)
# Verify direct access to TransferOpcode
self.assertIsNotNone(mooncake.transfer.TransferOpcode)
self.assertIsNotNone(mooncake.engine.TransferOpcode)
from mooncake.store import MooncakeDistributedStore
# Just verify we can create instances
store = MooncakeDistributedStore()
self.assertIsNotNone(store)
def test_direct_import(self):
"""Test direct import of specific components."""
from mooncake.transfer import TransferEngine, TransferOpcode
from mooncake.engine import TransferEngine, TransferOpcode
# Verify direct imports work
self.assertIsNotNone(TransferEngine)

22
mooncake-wheel/tests/transfer_engine_initiator_test.py
View File

@ -1,6 +1,6 @@
import unittest
import os
from mooncake import mooncake_vllm_adaptor
from mooncake.engine import TransferEngine
class TestVLLMAdaptorTransfer(unittest.TestCase):
@ -12,7 +12,7 @@ class TestVLLMAdaptorTransfer(unittest.TestCase):
cls.protocol = os.getenv("PROTOCOL", "tcp") # "rdma" or "tcp"
cls.circle = int(os.getenv("CIRCLE", 1000))
cls.adaptor = mooncake_vllm_adaptor()
cls.adaptor = TransferEngine()
ret = cls.adaptor.initialize(
cls.initiator_server_name,
cls.metadata_server,
@ -33,8 +33,8 @@ class TestVLLMAdaptorTransfer(unittest.TestCase):
adaptor = self.adaptor
circles = self.circle
src_addr = adaptor.getFirstBufferAddress(self.initiator_server_name)
dst_addr = adaptor.getFirstBufferAddress(self.target_server_name)
src_addr = adaptor.get_first_buffer_address(self.initiator_server_name)
dst_addr = adaptor.get_first_buffer_address(self.target_server_name)
for i in range(circles):
str_len = random.randint(16, 256)
@ -42,28 +42,28 @@ class TestVLLMAdaptorTransfer(unittest.TestCase):
data_len = len(src_data)
#Write to local buffer
result = adaptor.writeBytesToBuffer(src_addr, src_data, data_len)
result = adaptor.write_bytes_to_buffer(src_addr, src_data, data_len)
self.assertEqual(result, 0, f"[{i}] writeBytesToBuffer failed")
#Write to the remote end
result = adaptor.transferSyncExt(
self.target_server_name, src_addr, dst_addr, data_len, adaptor.TransferOpcode.WRITE
result = adaptor.transfer_sync_write(
self.target_server_name, src_addr, dst_addr, data_len
)
self.assertEqual(result, 0, f"[{i}] WRITE transferSyncExt failed")
#Clear the local buffer
clear_data = bytes([0] * data_len)
result = adaptor.writeBytesToBuffer(src_addr, clear_data, data_len)
result = adaptor.write_bytes_to_buffer(src_addr, clear_data, data_len)
self.assertEqual(result, 0, f"[{i}] Clear buffer failed")
#Read it back from the remote end
result = adaptor.transferSyncExt(
self.target_server_name, src_addr, dst_addr, data_len, adaptor.TransferOpcode.READ
result = adaptor.transfer_sync_read(
self.target_server_name, src_addr, dst_addr, data_len
)
self.assertEqual(result, 0, f"[{i}] READ transferSyncExt failed")
#Verify data consistency
read_back = adaptor.readBytesFromBuffer(src_addr, data_len)
read_back = adaptor.read_bytes_from_buffer(src_addr, data_len)
self.assertEqual(read_back, src_data, f"[{i}] Data mismatch")
print(f"[✓] {circles} iterations of random write-read passed successfully.")

4
mooncake-wheel/tests/transfer_engine_target.py
View File

@ -1,12 +1,12 @@
import os
from mooncake import mooncake_vllm_adaptor
from mooncake.engine import TransferEngine
target_server_name = os.getenv("TARGET_SERVER_NAME", "127.0.0.1:12345")
initiator_server_name = os.getenv("INITIATOR_SERVER_NAME", "127.0.0.1:12347")
metadata_server = os.getenv("MC_METADATA_SERVER", "127.0.0.1:2379")
protocol = os.getenv("PROTOCOL", "tcp") # Protocol type: "rdma" or "tcp"
target = mooncake_vllm_adaptor()
target = TransferEngine()
target.initialize(target_server_name,metadata_server, protocol, "")
while( True ):

5
scripts/build_wheel.sh
View File

@ -9,16 +9,17 @@ set -x
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib
echo "Creating directory structure..."
mkdir -p mooncake-wheel/mooncake/transfer/
echo "Copying Python modules..."
# Copy mooncake_vllm_adaptor to root level for backward compatibility
cp build/mooncake-integration/mooncake_vllm_adaptor.*.so mooncake-wheel/mooncake/mooncake_vllm_adaptor.so
cp build/mooncake-integration/mooncake_sglang_adaptor.*.so mooncake-wheel/mooncake/mooncake_sglang_adaptor.so
# Copy engine.so to mooncake directory (will be imported by transfer module)
cp build/mooncake-integration/engine.*.so mooncake-wheel/mooncake/engine.so
# Copy engine.so to mooncake directory (will be imported by transfer module)
cp build/mooncake-integration/store.*.so mooncake-wheel/mooncake/store.so
echo "Copying master binary and shared libraries..."
# Copy master binary and shared libraries
cp build/mooncake-store/src/mooncake_master mooncake-wheel/mooncake/

6
scripts/test_installation.sh
View File

@ -12,8 +12,8 @@ python -m venv test_env
source test_env/bin/activate
echo "Verifying that import fails before installation..."
# Verify that importing mooncake.transfer fails before installation
python -c "import mooncake.transfer" 2>/dev/null && { echo "ERROR: Import succeeded when it should have failed!"; exit 1; } || echo "Good: Import failed as expected before installation"
# Verify that importing mooncake.engine fails before installation
python -c "import mooncake.engine" 2>/dev/null && { echo "ERROR: Import succeeded when it should have failed!"; exit 1; } || echo "Good: Import failed as expected before installation"
echo "Installing the wheel package..."
# Install the wheel package
@ -31,7 +31,7 @@ sudo apt-get install -y $SYSTEM_PACKAGES
echo "Verifying that import succeeds after installation..."
python -c "import mooncake.transfer" && echo "Success: Import succeeded after installation" || { echo "ERROR: Import failed after installation!"; exit 1; }
python -c "import mooncake.engine" && echo "Success: Import succeeded after installation" || { echo "ERROR: Import failed after installation!"; exit 1; }
echo "Running import structure test..."
# Run the import structure test