From abad4f758bae754f43f83cfc3916cd6ae615c6b0 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Fri, 11 Apr 2025 17:54:02 +0800 Subject: [PATCH] [MooncakeAdaptor] reconstruct the adaptor arch (#232) --- .gitignore | 2 + mooncake-integration/CMakeLists.txt | 23 +- .../sglang/sglang_adaptor.cpp | 327 -------------- mooncake-integration/sglang/sglang_adaptor.h | 106 ----- mooncake-integration/store/store_py.cpp | 406 ++++++++++++++++++ mooncake-integration/store/store_py.h | 117 +++++ .../tests/stress_cluster_benchmark.py | 2 +- .../tests/test_distributed_object_store.py | 2 +- mooncake-wheel/mooncake/__init__.py | 4 - mooncake-wheel/mooncake/transfer/__init__.py | 6 - mooncake-wheel/setup.py | 3 - .../tests/test_distributed_object_store.py | 18 +- mooncake-wheel/tests/test_import_structure.py | 37 +- .../tests/transfer_engine_initiator_test.py | 22 +- .../tests/transfer_engine_target.py | 4 +- scripts/build_wheel.sh | 5 +- scripts/test_installation.sh | 6 +- 17 files changed, 589 insertions(+), 501 deletions(-) delete mode 100644 mooncake-integration/sglang/sglang_adaptor.cpp delete mode 100644 mooncake-integration/sglang/sglang_adaptor.h create mode 100644 mooncake-integration/store/store_py.cpp create mode 100644 mooncake-integration/store/store_py.h delete mode 100644 mooncake-wheel/mooncake/transfer/__init__.py diff --git a/.gitignore b/.gitignore index dbb6259a..9aba795e 100644 --- a/.gitignore +++ b/.gitignore @@ -191,3 +191,5 @@ cmake-build-minsizerel cmake-build libetcd_wrapper.h + +mooncake-wheel/mooncake/mooncake_master diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 2d38db29..c1a047d9 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -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}) \ No newline at end of file diff --git a/mooncake-integration/sglang/sglang_adaptor.cpp b/mooncake-integration/sglang/sglang_adaptor.cpp deleted file mode 100644 index 8fbedf9f..00000000 --- a/mooncake-integration/sglang/sglang_adaptor.cpp +++ /dev/null @@ -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 - -#ifdef USE_CUDA -#include -#include -#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 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 parseConnectionString( - const std::string &conn_string) { - std::pair 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(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 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 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(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(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_ 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_(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; -} diff --git a/mooncake-integration/sglang/sglang_adaptor.h b/mooncake-integration/sglang/sglang_adaptor.h deleted file mode 100644 index d9cc544c..00000000 --- a/mooncake-integration/sglang/sglang_adaptor.h +++ /dev/null @@ -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 -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#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(reinterpret_cast(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 engine_; - Transport *xport_; - - std::mutex mutex_; - std::vector> free_list_; - std::vector buffer_list_; - std::unordered_set large_buffer_list_; - std::unordered_map handle_map_; -}; diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp new file mode 100644 index 00000000..4878836b --- /dev/null +++ b/mooncake-integration/store/store_py.cpp @@ -0,0 +1,406 @@ +#include "store_py.h" + +#include +#include +#include + +#include // for atexit +#include + +#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 lock(mutex_); + instances_.insert(instance); +} + +void ResourceTracker::unregisterInstance(DistributedObjectStore *instance) { + std::lock_guard lock(mutex_); + instances_.erase(instance); +} + +void ResourceTracker::cleanupAllResources() { + std::lock_guard lock(mutex_); + + // Perform cleanup outside the lock to avoid potential deadlocks + for (void *instance : instances_) { + DistributedObjectStore *store = + static_cast(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(); + + 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(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 &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 &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 &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 &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 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 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_(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); + +} + diff --git a/mooncake-integration/store/store_py.h b/mooncake-integration/store/store_py.h new file mode 100644 index 00000000..037c8ba8 --- /dev/null +++ b/mooncake-integration/store/store_py.h @@ -0,0 +1,117 @@ +#pragma once + +#include + +#include +#include +#include +#include + +#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 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 &slices, + const std::string &value); + + int allocateSlices(std::vector &slices, + const mooncake::Client::ObjectInfo &object_info, + uint64_t &length); + + char *exportSlices(const std::vector &slices, + uint64_t length); + + int freeSlices(const std::vector &slices); + + public: + std::unique_ptr client_ = nullptr; + std::unique_ptr client_buffer_allocator_ = + nullptr; + struct SegmentDeleter { + void operator()(void *ptr) { + if (ptr) { + free(ptr); + } + } + }; + + std::unique_ptr segment_ptr_; + std::string protocol; + std::string device_name; + std::string local_hostname; +}; diff --git a/mooncake-store/tests/stress_cluster_benchmark.py b/mooncake-store/tests/stress_cluster_benchmark.py index 3e8a3485..b14cd2b8 100644 --- a/mooncake-store/tests/stress_cluster_benchmark.py +++ b/mooncake-store/tests/stress_cluster_benchmark.py @@ -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 diff --git a/mooncake-store/tests/test_distributed_object_store.py b/mooncake-store/tests/test_distributed_object_store.py index 23be11aa..fd783975 100644 --- a/mooncake-store/tests/test_distributed_object_store.py +++ b/mooncake-store/tests/test_distributed_object_store.py @@ -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): diff --git a/mooncake-wheel/mooncake/__init__.py b/mooncake-wheel/mooncake/__init__.py index 5351cf87..ebc08cad 100644 --- a/mooncake-wheel/mooncake/__init__.py +++ b/mooncake-wheel/mooncake/__init__.py @@ -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 diff --git a/mooncake-wheel/mooncake/transfer/__init__.py b/mooncake-wheel/mooncake/transfer/__init__.py deleted file mode 100644 index 44e33c67..00000000 --- a/mooncake-wheel/mooncake/transfer/__init__.py +++ /dev/null @@ -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'] diff --git a/mooncake-wheel/setup.py b/mooncake-wheel/setup.py index 2e52d313..765a6707 100644 --- a/mooncake-wheel/setup.py +++ b/mooncake-wheel/setup.py @@ -28,9 +28,6 @@ setup( "*.so", "mooncake_master", ], - "mooncake.transfer": [ - "*.so", - ], }, include_package_data=True, zip_safe=False, diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index 6ee75821..6ce4293f 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -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.""" diff --git a/mooncake-wheel/tests/test_import_structure.py b/mooncake-wheel/tests/test_import_structure.py index 5b731492..8200e2f7 100644 --- a/mooncake-wheel/tests/test_import_structure.py +++ b/mooncake-wheel/tests/test_import_structure.py @@ -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) diff --git a/mooncake-wheel/tests/transfer_engine_initiator_test.py b/mooncake-wheel/tests/transfer_engine_initiator_test.py index 7b069c5a..58fe7685 100644 --- a/mooncake-wheel/tests/transfer_engine_initiator_test.py +++ b/mooncake-wheel/tests/transfer_engine_initiator_test.py @@ -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.") diff --git a/mooncake-wheel/tests/transfer_engine_target.py b/mooncake-wheel/tests/transfer_engine_target.py index 64f7d1ad..f77e26e0 100644 --- a/mooncake-wheel/tests/transfer_engine_target.py +++ b/mooncake-wheel/tests/transfer_engine_target.py @@ -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 ): diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index 1da61afa..ea300aa8 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -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/ diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index 8d2e5d9e..21662c64 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -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