diff --git a/docs/source/mooncake-store-api/python-binding.md b/docs/source/mooncake-store-api/python-binding.md index a1a51d09..1b568b00 100644 --- a/docs/source/mooncake-store-api/python-binding.md +++ b/docs/source/mooncake-store-api/python-binding.md @@ -351,7 +351,7 @@ print("Retrieved all keys successfully:", retrieved == values) ## get_buffer Buffer Protocol -The `get_buffer` method returns a `SliceBuffer` object that implements the Python buffer protocol: +The `get_buffer` method returns a `BufferHandle` object that implements the Python buffer protocol:
Click to expand: Buffer protocol usage example @@ -684,14 +684,14 @@ else: Get object data as a buffer that implements Python's buffer protocol. ```python -def get_buffer(self, key: str) -> SliceBuffer +def get_buffer(self, key: str) -> BufferHandle ``` **Parameters:** - `key` (str): Object identifier **Returns:** -- `SliceBuffer`: Buffer object or None if not found +- `BufferHandle`: Buffer object or None if not found **Example:** ```python diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index e89dd44b..2e3f167a 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -1,1529 +1,279 @@ -#include "store_py.h" -#include + #include // For GIL management #include -#include -#include -#include // for timing +#include "pybind_client.h" + #include // for atexit -#include // for std::setprecision -#include // for std::accumulate -#include -#include -#include "client_buffer.hpp" -#include "config.h" -#include "types.h" -#include "utils.h" #include "integration_utils.h" namespace py = pybind11; namespace mooncake { -// ResourceTracker implementation using singleton pattern -ResourceTracker &ResourceTracker::getInstance() { - static ResourceTracker instance; - return instance; -} +// Python-specific wrapper functions that handle GIL and return pybind11 types +class MooncakeStorePyWrapper { + public: + PyClient store_; -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(); + pybind11::bytes get(const std::string &key) { + if (!store_.client_) { + LOG(ERROR) << "Client is not initialized"; + return pybind11::bytes("\\0", 0); } - } -} -void ResourceTracker::signalHandler(int signal) { - LOG(INFO) << "Received signal " << signal << ", cleaning up resources"; - getInstance().cleanupAllResources(); + const auto kNullString = pybind11::bytes("\\0", 0); - // 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); -} + { + py::gil_scoped_release release_gil; + auto buffer_handle = store_.get_buffer(key); + if (!buffer_handle) { + py::gil_scoped_acquire acquire_gil; + return kNullString; + } -void ResourceTracker::exitHandler() { 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); - 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); -} - -tl::expected DistributedObjectStore::setup_internal( - 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 tl::unexpected(ErrorCode::INVALID_PARAMS); - } - // Combine hostname with port - this->local_hostname = hostname + ":" + std::to_string(port); - } else { - this->local_hostname = local_hostname; - } - - void **args = (protocol == "rdma") ? rdma_args(rdma_devices) : nullptr; - auto client_opt = - mooncake::Client::Create(this->local_hostname, metadata_server, - protocol, args, master_server_addr); - if (!client_opt) { - LOG(ERROR) << "Failed to create client"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - client_ = *client_opt; - - // Local_buffer_size is allowed to be 0, but we only register memory when - // local_buffer_size > 0. Invoke ibv_reg_mr() with size=0 is UB, and may - // fail in some rdma implementations. - client_buffer_allocator_ = ClientBufferAllocator::create(local_buffer_size); - if (local_buffer_size > 0) { - auto result = client_->RegisterLocalMemory( - client_buffer_allocator_->getBase(), local_buffer_size, - kWildcardLocation, false, true); - if (!result.has_value()) { - LOG(ERROR) << "Failed to register local memory: " - << toString(result.error()); - return tl::unexpected(result.error()); - } - } else { - LOG(INFO) << "Local buffer size is 0, skip registering local memory"; - } - - // If global_segment_size is 0, skip mount segment; - // If global_segment_size is larger than max_mr_size, split to multiple - // segments. - auto max_mr_size = globalConfig().max_mr_size; // Max segment size - uint64_t total_glbseg_size = global_segment_size; // For logging - uint64_t current_glbseg_size = 0; // For logging - while (global_segment_size > 0) { - size_t segment_size = std::min(global_segment_size, max_mr_size); - global_segment_size -= segment_size; - current_glbseg_size += segment_size; - LOG(INFO) << "Mounting segment: " << segment_size << " bytes, " - << current_glbseg_size << " of " << total_glbseg_size; - void *ptr = allocate_buffer_allocator_memory(segment_size); - if (!ptr) { - LOG(ERROR) << "Failed to allocate segment memory"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - segment_ptrs_.emplace_back(ptr); - auto mount_result = client_->MountSegment(ptr, segment_size); - if (!mount_result.has_value()) { - LOG(ERROR) << "Failed to mount segment: " - << toString(mount_result.error()); - return tl::unexpected(mount_result.error()); - } - } - if (total_glbseg_size == 0) { - LOG(INFO) << "Global segment size is 0, skip mounting segment"; - } - - return {}; -} - -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) { - return to_py_ret(setup_internal( - local_hostname, metadata_server, global_segment_size, local_buffer_size, - protocol, rdma_devices, master_server_addr)); -} - -tl::expected DistributedObjectStore::initAll_internal( - const std::string &protocol_, const std::string &device_name, - size_t mount_segment_size) { - if (client_) { - LOG(ERROR) << "Client is already initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - uint64_t buffer_allocator_size = 1024 * 1024 * 1024; - return setup_internal("localhost:12345", "127.0.0.1:2379", - mount_segment_size, buffer_allocator_size, protocol_, - device_name); -} - -int DistributedObjectStore::initAll(const std::string &protocol_, - const std::string &device_name, - size_t mount_segment_size) { - return to_py_ret( - initAll_internal(protocol_, device_name, mount_segment_size)); -} - -tl::expected DistributedObjectStore::tearDownAll_internal() { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - // Reset all resources - client_.reset(); - client_buffer_allocator_.reset(); - segment_ptrs_.clear(); - local_hostname = ""; - device_name = ""; - protocol = ""; - return {}; -} - -int DistributedObjectStore::tearDownAll() { - return to_py_ret(tearDownAll_internal()); -} - -tl::expected DistributedObjectStore::put_internal( - const std::string &key, std::span value, - const ReplicateConfig &config) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - auto alloc_result = client_buffer_allocator_->allocate(value.size_bytes()); - if (!alloc_result) { - LOG(ERROR) << "Failed to allocate buffer for put operation, key: " - << key << ", value size: " << value.size(); - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - auto &buffer_handle = *alloc_result; - memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); - - std::vector slices = split_into_slices(buffer_handle); - - auto put_result = client_->Put(key, slices, config); - if (!put_result) { - LOG(ERROR) << "Put operation failed with error: " - << toString(put_result.error()); - return tl::unexpected(put_result.error()); - } - - return {}; -} - -int DistributedObjectStore::put(const std::string &key, - std::span value, - const ReplicateConfig &config) { - return to_py_ret(put_internal(key, value, config)); -} - -tl::expected DistributedObjectStore::put_batch_internal( - const std::vector &keys, - const std::vector> &values, - const ReplicateConfig &config) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - if (keys.size() != values.size()) { - LOG(ERROR) << "Key and value size mismatch"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - std::vector buffer_handles; - std::unordered_map> batched_slices; - batched_slices.reserve(keys.size()); - - for (size_t i = 0; i < keys.size(); ++i) { - auto &key = keys[i]; - auto &value = values[i]; - auto alloc_result = - client_buffer_allocator_->allocate(value.size_bytes()); - if (!alloc_result) { - LOG(ERROR) - << "Failed to allocate buffer for put_batch operation, key: " - << key << ", value size: " << value.size(); - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - auto &buffer_handle = *alloc_result; - memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); - auto slices = split_into_slices(buffer_handle); - buffer_handles.emplace_back(std::move(*alloc_result)); - batched_slices.emplace(key, std::move(slices)); - } - - // Convert unordered_map to vector format expected by BatchPut - std::vector> ordered_batched_slices; - ordered_batched_slices.reserve(keys.size()); - for (const auto &key : keys) { - auto it = batched_slices.find(key); - if (it != batched_slices.end()) { - ordered_batched_slices.emplace_back(it->second); - } else { - LOG(ERROR) << "Missing slices for key: " << key; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - } - - auto results = client_->BatchPut(keys, ordered_batched_slices, config); - - // Check if any operations failed - for (size_t i = 0; i < results.size(); ++i) { - if (!results[i]) { - LOG(ERROR) << "BatchPut operation failed for key '" << keys[i] - << "' with error: " << toString(results[i].error()); - return tl::unexpected(results[i].error()); - } - } - return {}; -} - -int DistributedObjectStore::put_batch( - const std::vector &keys, - const std::vector> &values, - const ReplicateConfig &config) { - return to_py_ret(put_batch_internal(keys, values, config)); -} - -tl::expected DistributedObjectStore::put_parts_internal( - const std::string &key, std::vector> values, - const ReplicateConfig &config) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - // Calculate total size needed - size_t total_size = 0; - for (const auto &value : values) { - total_size += value.size_bytes(); - } - - if (total_size == 0) { - LOG(WARNING) << "Attempting to put empty data for key: " << key; - return {}; - } - - // Allocate buffer using the new allocator - auto alloc_result = client_buffer_allocator_->allocate(total_size); - if (!alloc_result) { - LOG(ERROR) << "Failed to allocate buffer for put_parts operation, key: " - << key << ", total size: " << total_size; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - auto &buffer_handle = *alloc_result; - - // Copy all parts into the contiguous buffer - size_t offset = 0; - for (const auto &value : values) { - memcpy(static_cast(buffer_handle.ptr()) + offset, value.data(), - value.size_bytes()); - offset += value.size_bytes(); - } - - // Split into slices - std::vector slices = split_into_slices(buffer_handle); - - // Perform the put operation - buffer_handle will be automatically released - auto put_result = client_->Put(key, slices, config); - if (!put_result) { - LOG(ERROR) << "Put operation failed with error: " - << toString(put_result.error()); - return tl::unexpected(put_result.error()); - } - - return {}; -} - -int DistributedObjectStore::put_parts(const std::string &key, - std::vector> values, - const ReplicateConfig &config) { - return to_py_ret(put_parts_internal(key, values, config)); -} - -pybind11::bytes DistributedObjectStore::get(const std::string &key) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return pybind11::bytes("\0", 0); - } - - const auto kNullString = pybind11::bytes("\0", 0); - - { - py::gil_scoped_release release_gil; - - auto query_result = client_->Query(key); - if (!query_result) { py::gil_scoped_acquire acquire_gil; - return kNullString; + return pybind11::bytes((char *)buffer_handle->ptr(), + buffer_handle->size()); } - - auto replica_list = query_result.value(); - if (replica_list.empty()) { - py::gil_scoped_acquire acquire_gil; - return kNullString; - } - - // Calculate total size - const auto &replica = replica_list[0]; - uint64_t total_size = calculate_total_size(replica); - - if (total_size == 0) { - py::gil_scoped_acquire acquire_gil; - return pybind11::bytes("", 0); - } - - // Allocate buffer using the new allocator - auto alloc_result = client_buffer_allocator_->allocate(total_size); - if (!alloc_result) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Failed to allocate buffer for get operation, key: " - << key << ", size: " << total_size; - return kNullString; - } - - auto &buffer_handle = *alloc_result; - - // Create slices for the allocated buffer based on memory descriptors - std::vector slices; - allocateSlices(slices, replica, buffer_handle); - - // Get the object data - auto get_result = client_->Get(key, replica_list, slices); - if (!get_result) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Get operation failed with error: " - << toString(get_result.error()); - return kNullString; - } - - py::gil_scoped_acquire acquire_gil; - - // Create Python bytes object - buffer_handle will be released - // automatically - return pybind11::bytes(static_cast(buffer_handle.ptr()), - total_size); - } -} - -std::vector DistributedObjectStore::get_batch( - const std::vector &keys) { - const auto kNullString = pybind11::bytes("\0", 0); - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - py::gil_scoped_acquire acquire_gil; - return {kNullString}; } - std::unordered_set seen; - for (const auto &key : keys) { - if (!seen.insert(key).second) { - LOG(ERROR) << "Duplicate key not supported for Batch API, key: " - << key; + std::vector get_batch( + const std::vector &keys) { + const auto kNullString = pybind11::bytes("\\0", 0); + if (!store_.client_) { + LOG(ERROR) << "Client is not initialized"; py::gil_scoped_acquire acquire_gil; return {kNullString}; } + + { + py::gil_scoped_release release_gil; + auto batch_data = store_.batch_get_buffer(keys); + if (batch_data.empty()) { + py::gil_scoped_acquire acquire_gil; + return {kNullString}; + } + + py::gil_scoped_acquire acquire_gil; + std::vector results; + results.reserve(batch_data.size()); + + for (const auto &data : batch_data) { + results.emplace_back( + data ? pybind11::bytes((char *)data->ptr(), data->size()) + : kNullString); + } + + return results; + } } - { - py::gil_scoped_release release_gil; - auto query_results = client_->BatchQuery(keys); - - // Extract successful replica lists - std::vector> replica_lists; - replica_lists.reserve(keys.size()); - for (size_t i = 0; i < query_results.size(); ++i) { - if (!query_results[i]) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Query failed for key '" << keys[i] - << "': " << toString(query_results[i].error()); - return {kNullString}; - } - replica_lists.emplace_back(query_results[i].value()); + pybind11::object get_tensor(const std::string &key) { + if (!store_.client_) { + LOG(ERROR) << "Client is not initialized"; + return pybind11::none(); } - // Prepare buffers and slices for each key - std::vector> buffer_handles; - std::vector> all_slices; - std::vector total_sizes; - - buffer_handles.reserve(keys.size()); - all_slices.reserve(keys.size()); - total_sizes.reserve(keys.size()); - - for (size_t i = 0; i < keys.size(); ++i) { - const auto &replica_list = replica_lists[i]; - if (replica_list.empty()) { + try { + // Query object info first + auto query_result = store_.client_->Query(key); + if (!query_result) { py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Empty replica list for key: " << keys[i]; - return {kNullString}; + LOG(ERROR) << "Query failed: " << query_result.error(); + return pybind11::none(); } - // Calculate total size - const auto &replica = replica_list[0]; - uint64_t total_size = calculate_total_size(replica); - ; + auto replica_list = query_result.value(); + if (replica_list.empty()) { + py::gil_scoped_acquire acquire_gil; + LOG(INFO) << "No replicas found for key: " << key; + return pybind11::none(); + } - // Allocate buffer - auto alloc_result = client_buffer_allocator_->allocate(total_size); + const auto &replica = replica_list[0]; + uint64_t total_length = calculate_total_size(replica); + + if (total_length == 0) { + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Failed to allocate slices for key: " << key; + return pybind11::none(); + } + + // Allocate buffer using the new allocator + auto alloc_result = + store_.client_buffer_allocator_->allocate(total_length); if (!alloc_result) { py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Failed to allocate buffer for key: " << keys[i]; - return {kNullString}; + return pybind11::none(); } auto &buffer_handle = *alloc_result; - // Create slices + // Create slices for the allocated buffer std::vector slices; allocateSlices(slices, replica, buffer_handle); - buffer_handles.emplace_back( - std::make_unique(std::move(buffer_handle))); - all_slices.emplace_back(std::move(slices)); - total_sizes.emplace_back(total_size); - } - - // Prepare batch transfer data structures - std::vector batch_keys = keys; - std::unordered_map> batch_slices; - - for (size_t i = 0; i < keys.size(); ++i) { - batch_slices[keys[i]] = all_slices[i]; - } - - // Execute batch transfer - auto batch_get_results = - client_->BatchGet(batch_keys, replica_lists, batch_slices); - - py::gil_scoped_acquire acquire_gil; - std::vector results; - results.reserve(keys.size()); - - for (size_t i = 0; i < keys.size(); ++i) { - if (!batch_get_results[i]) { - LOG(ERROR) << "BatchGet failed for key '" << keys[i] - << "': " << toString(batch_get_results[i].error()); - return {kNullString}; + // Get the object data + auto get_result = store_.client_->Get(key, replica_list, slices); + if (!get_result) { + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Get failed for key: " << key; + return pybind11::none(); } - // Create Python bytes object from buffer - results.emplace_back(pybind11::bytes( - static_cast(buffer_handles[i]->ptr()), total_sizes[i])); - } - - return results; - } -} - -tl::expected DistributedObjectStore::remove_internal( - const std::string &key) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - auto remove_result = client_->Remove(key); - if (!remove_result) { - return tl::unexpected(remove_result.error()); - } - return {}; -} - -int DistributedObjectStore::remove(const std::string &key) { - return to_py_ret(remove_internal(key)); -} - -tl::expected DistributedObjectStore::removeAll_internal() { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - return client_->RemoveAll(); -} - -long DistributedObjectStore::removeAll() { - return to_py_ret(removeAll_internal()); -} - -tl::expected DistributedObjectStore::isExist_internal( - const std::string &key) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - return client_->IsExist(key); -} - -int DistributedObjectStore::isExist(const std::string &key) { - auto result = isExist_internal(key); - - if (result.has_value()) { - return *result ? 1 : 0; // 1 if exists, 0 if not - } else { - return toInt(result.error()); - } -} - -std::vector DistributedObjectStore::batchIsExist( - const std::vector &keys) { - auto internal_results = batchIsExist_internal(keys); - std::vector results; - results.reserve(internal_results.size()); - - for (const auto &result : internal_results) { - if (result.has_value()) { - results.push_back(result.value() ? 1 : 0); // 1 if exists, 0 if not - } else { - results.push_back(toInt(result.error())); - } - } - - return results; -} - -tl::expected DistributedObjectStore::getSize_internal( - const std::string &key) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - auto query_result = client_->Query(key); - - if (!query_result) { - return tl::unexpected(query_result.error()); - } - - auto replica_list = query_result.value(); - - // Calculate total size from all replicas' handles - int64_t total_size = 0; - if (!replica_list.empty()) { - auto &replica = replica_list[0]; - total_size = calculate_total_size(replica); - } else { - LOG(ERROR) << "Internal error: replica_list is empty"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); // Internal error - } - - return total_size; -} - -int64_t DistributedObjectStore::getSize(const std::string &key) { - return to_py_ret(getSize_internal(key)); -} - -// SliceBuffer implementation -SliceBuffer::SliceBuffer(BufferHandle handle) : handle_(std::move(handle)) {} - -void *SliceBuffer::ptr() const { return handle_.ptr(); } - -uint64_t SliceBuffer::size() const { return handle_.size(); } - -// Implementation of get_buffer method -std::shared_ptr DistributedObjectStore::get_buffer( - const std::string &key) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return nullptr; - } - - // Query the object info - auto query_result = client_->Query(key); - if (!query_result) { - if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND) { - return nullptr; - } - LOG(ERROR) << "Query failed for key: " << key - << " with error: " << toString(query_result.error()); - return nullptr; - } - - auto replica_list = query_result.value(); - if (replica_list.empty()) { - LOG(ERROR) << "Empty replica list for key: " << key; - return nullptr; - } - - const auto &replica = replica_list[0]; - uint64_t total_length = calculate_total_size(replica); - - if (total_length == 0) { - return nullptr; - } - - // Allocate buffer using the new allocator - auto alloc_result = client_buffer_allocator_->allocate(total_length); - if (!alloc_result) { - LOG(ERROR) << "Failed to allocate buffer for get_buffer, key: " << key; - return nullptr; - } - - auto &buffer_handle = *alloc_result; - - // Create slices for the allocated buffer - std::vector slices; - allocateSlices(slices, replica, buffer_handle); - - // Get the object data - auto get_result = client_->Get(key, replica_list, slices); - if (!get_result) { - LOG(ERROR) << "Get failed for key: " << key - << " with error: " << toString(get_result.error()); - return nullptr; - } - - // Create SliceBuffer with the allocated memory - // The buffer will be managed by the SliceBuffer's shared_ptr - return std::make_shared(std::move(buffer_handle)); -} - -// Implementation of batch_get_buffer_internal method -std::vector> -DistributedObjectStore::batch_get_buffer_internal( - const std::vector &keys) { - std::vector> final_results(keys.size(), - nullptr); - - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return final_results; - } - - if (keys.empty()) { - return final_results; - } - - // 1. Query metadata for all keys - auto query_results = client_->BatchQuery(keys); - - // 2. Prepare for batch get: filter valid keys and prepare buffers - struct KeyOp { - size_t original_index; - std::string key; - std::vector replica_list; - std::unique_ptr buffer_handle; - std::vector slices; - }; - std::vector valid_ops; - valid_ops.reserve(keys.size()); - - for (size_t i = 0; i < keys.size(); ++i) { - const auto &key = keys[i]; - - if (!query_results[i]) { - if (query_results[i].error() != ErrorCode::OBJECT_NOT_FOUND) { - LOG(ERROR) << "Query failed for key '" << key - << "': " << toString(query_results[i].error()); + // Create contiguous buffer and copy data + char *exported_data = new char[total_length]; + if (!exported_data) { + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Invalid data format: insufficient data for " + "metadata"; + return pybind11::none(); } - continue; - } + TensorMetadata metadata; - auto replica_list = query_results[i].value(); - if (replica_list.empty()) { - LOG(ERROR) << "Empty replica list for key: " << key; - continue; - } + // Copy data from buffer to contiguous memory + memcpy(exported_data, buffer_handle.ptr(), total_length); + memcpy(&metadata, exported_data, sizeof(TensorMetadata)); - const auto &replica = replica_list[0]; - uint64_t total_size = calculate_total_size(replica); - if (total_size == 0) { - continue; - } - - auto alloc_result = client_buffer_allocator_->allocate(total_size); - if (!alloc_result) { - LOG(ERROR) << "Failed to allocate buffer for key: " << key; - continue; - } - - auto buffer_handle = - std::make_unique(std::move(*alloc_result)); - std::vector slices; - allocateSlices(slices, replica, *buffer_handle); - - valid_ops.emplace_back(KeyOp{.original_index = i, - .key = key, - .replica_list = std::move(replica_list), - .buffer_handle = std::move(buffer_handle), - .slices = std::move(slices)}); - } - - if (valid_ops.empty()) { - return final_results; - } - - // 3. Execute batch get - std::vector batch_keys; - std::vector> batch_replica_lists; - std::unordered_map> batch_slices; - batch_keys.reserve(valid_ops.size()); - batch_replica_lists.reserve(valid_ops.size()); - - for (auto &op : valid_ops) { - batch_keys.push_back(op.key); - batch_replica_lists.push_back(op.replica_list); - batch_slices[op.key] = op.slices; - } - - auto batch_get_results = - client_->BatchGet(batch_keys, batch_replica_lists, batch_slices); - - // 4. Process results and create SliceBuffers - for (size_t i = 0; i < valid_ops.size(); ++i) { - if (batch_get_results[i]) { - auto &op = valid_ops[i]; - final_results[op.original_index] = - std::make_shared(std::move(*op.buffer_handle)); - } else { - LOG(ERROR) << "BatchGet failed for key '" << valid_ops[i].key - << "': " << toString(batch_get_results[i].error()); - } - } - - return final_results; -} - -// Implementation of batch_get_buffer method -std::vector> -DistributedObjectStore::batch_get_buffer(const std::vector &keys) { - return batch_get_buffer_internal(keys); -} - -tl::expected DistributedObjectStore::register_buffer_internal( - void *buffer, size_t size) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - return client_->RegisterLocalMemory(buffer, size, kWildcardLocation, false, - true); -} - -int DistributedObjectStore::register_buffer(void *buffer, size_t size) { - return to_py_ret(register_buffer_internal(buffer, size)); -} - -tl::expected -DistributedObjectStore::unregister_buffer_internal(void *buffer) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - auto unregister_result = client_->unregisterLocalMemory(buffer, true); - if (!unregister_result) { - LOG(ERROR) << "Unregister buffer failed with error: " - << toString(unregister_result.error()); - return tl::unexpected(unregister_result.error()); - } - return {}; -} - -int DistributedObjectStore::unregister_buffer(void *buffer) { - return to_py_ret(unregister_buffer_internal(buffer)); -} - -tl::expected DistributedObjectStore::get_into_internal( - const std::string &key, void *buffer, size_t size) { - // NOTE: The buffer address must be previously registered with - // register_buffer() for zero-copy RDMA operations to work correctly - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - // Step 1: Get object info - auto query_result = client_->Query(key); - if (!query_result) { - if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND) { - VLOG(1) << "Object not found for key: " << key; - return tl::unexpected(query_result.error()); - } - LOG(ERROR) << "Query failed for key: " << key - << " with error: " << toString(query_result.error()); - return tl::unexpected(query_result.error()); - } - - auto replica_list = query_result.value(); - - // Calculate total size from replica list - if (replica_list.empty()) { - LOG(ERROR) << "Internal error: replica_list is empty"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - auto &replica = replica_list[0]; - uint64_t total_size = calculate_total_size(replica); - - // Check if user buffer is large enough - if (size < total_size) { - LOG(ERROR) << "User buffer too small. Required: " << total_size - << ", provided: " << size; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - // Step 2: Split user buffer according to object info and create - // slices - std::vector slices; - uint64_t offset = 0; - - if (replica.is_memory_replica() == false) { - while (offset < total_size) { - auto chunk_size = std::min(total_size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - } else { - for (auto &handle : - replica.get_memory_descriptor().buffer_descriptors) { - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, handle.size_}); - offset += handle.size_; - } - } - - // Step 3: Read data directly into user buffer - auto get_result = client_->Get(key, replica_list, slices); - if (!get_result) { - LOG(ERROR) << "Get failed for key: " << key - << " with error: " << toString(get_result.error()); - return tl::unexpected(get_result.error()); - } - - return static_cast(total_size); -} - -int DistributedObjectStore::get_into(const std::string &key, void *buffer, - size_t size) { - return to_py_ret(get_into_internal(key, buffer, size)); -} - -std::string DistributedObjectStore::get_hostname() const { - return local_hostname; -} - -std::vector DistributedObjectStore::batch_put_from( - const std::vector &keys, const std::vector &buffers, - const std::vector &sizes, const ReplicateConfig &config) { - auto internal_results = - batch_put_from_internal(keys, buffers, sizes, config); - std::vector results; - results.reserve(internal_results.size()); - - for (const auto &result : internal_results) { - results.push_back(to_py_ret(result)); - } - - return results; -} - -std::vector DistributedObjectStore::batch_get_into( - const std::vector &keys, const std::vector &buffers, - const std::vector &sizes) { - auto internal_results = batch_get_into_internal(keys, buffers, sizes); - std::vector results; - results.reserve(internal_results.size()); - - for (const auto &result : internal_results) { - results.push_back(to_py_ret(result)); - } - - return results; -} - -tl::expected DistributedObjectStore::put_from_internal( - const std::string &key, void *buffer, size_t size, - const ReplicateConfig &config) { - // NOTE: The buffer address must be previously registered with - // register_buffer() for zero-copy RDMA operations to work correctly - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - if (size == 0) { - LOG(WARNING) << "Attempting to put empty data for key: " << key; - return {}; - } - - // Create slices directly from the user buffer - std::vector slices; - uint64_t offset = 0; - - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - - auto put_result = client_->Put(key, slices, config); - if (!put_result) { - LOG(ERROR) << "Put operation failed with error: " - << toString(put_result.error()); - return tl::unexpected(put_result.error()); - } - - return {}; -} - -int DistributedObjectStore::put_from(const std::string &key, void *buffer, - size_t size, - const ReplicateConfig &config) { - return to_py_ret(put_from_internal(key, buffer, size, config)); -} - -std::vector> -DistributedObjectStore::batch_get_into_internal( - const std::vector &keys, const std::vector &buffers, - const std::vector &sizes) { - // Validate preconditions - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } - - if (keys.size() != buffers.size() || keys.size() != sizes.size()) { - LOG(ERROR) << "Input vector sizes mismatch: keys=" << keys.size() - << ", buffers=" << buffers.size() - << ", sizes=" << sizes.size(); - return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } - - const size_t num_keys = keys.size(); - std::vector> results; - results.reserve(num_keys); - - if (num_keys == 0) { - return results; - } - - // Query metadata for all keys - const auto query_results = client_->BatchQuery(keys); - - // Process each key individually and prepare for batch transfer - struct ValidKeyInfo { - std::string key; - size_t original_index; - std::vector replica_list; - std::vector slices; - uint64_t total_size; - }; - - std::vector valid_operations; - valid_operations.reserve(num_keys); - - for (size_t i = 0; i < num_keys; ++i) { - const auto &key = keys[i]; - - // Handle query failures - if (!query_results[i]) { - const auto error = query_results[i].error(); - results.emplace_back(tl::unexpected(error)); - if (error != ErrorCode::OBJECT_NOT_FOUND) { - LOG(ERROR) << "Query failed for key '" << key - << "': " << toString(error); + if (metadata.ndim < 0 || metadata.ndim > 4) { + delete[] exported_data; + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Invalid tensor metadata: ndim=" << metadata.ndim; + return pybind11::none(); } - continue; - } - // Validate replica list - auto replica_list = query_results[i].value(); - if (replica_list.empty()) { - LOG(ERROR) << "Empty replica list for key: " << key; - results.emplace_back(tl::unexpected(ErrorCode::INVALID_REPLICA)); - continue; - } - - // Calculate required buffer size - const auto &replica = replica_list[0]; - uint64_t total_size = calculate_total_size(replica); - - // Validate buffer capacity - if (sizes[i] < total_size) { - LOG(ERROR) << "Buffer too small for key '" << key - << "': required=" << total_size - << ", available=" << sizes[i]; - results.emplace_back(tl::unexpected(ErrorCode::INVALID_PARAMS)); - continue; - } - - // Create slices for this key's buffer - std::vector key_slices; - uint64_t offset = 0; - if (replica.is_memory_replica() == false) { - while (offset < total_size) { - auto chunk_size = std::min(total_size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffers[i]) + offset; - key_slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; + TensorDtype dtype_enum = static_cast(metadata.dtype); + if (dtype_enum == TensorDtype::UNKNOWN) { + delete[] exported_data; + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Unknown tensor dtype!"; + return pybind11::none(); } - } else { - for (auto &handle : - replica.get_memory_descriptor().buffer_descriptors) { - void *chunk_ptr = static_cast(buffers[i]) + offset; - key_slices.emplace_back(Slice{chunk_ptr, handle.size_}); - offset += handle.size_; + + size_t tensor_size = total_length - sizeof(TensorMetadata); + if (tensor_size == 0) { + delete[] exported_data; + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Invalid data format: no tensor data found"; + return pybind11::none(); } - } - // Store operation info for batch processing - valid_operations.push_back({.key = key, - .original_index = i, - .replica_list = std::move(replica_list), - .slices = std::move(key_slices), - .total_size = total_size}); - - // Set success result (actual bytes transferred) - results.emplace_back(static_cast(total_size)); - } - - // Early return if no valid operations - if (valid_operations.empty()) { - return results; - } - - // Prepare batch transfer data structures - std::vector batch_keys; - std::vector> batch_replica_lists; - std::unordered_map> batch_slices; - - batch_keys.reserve(valid_operations.size()); - batch_replica_lists.reserve(valid_operations.size()); - - for (const auto &op : valid_operations) { - batch_keys.push_back(op.key); - batch_replica_lists.push_back(op.replica_list); - batch_slices[op.key] = op.slices; - } - - // Execute batch transfer - const auto batch_get_results = - client_->BatchGet(batch_keys, batch_replica_lists, batch_slices); - - // Process transfer results - for (size_t j = 0; j < batch_get_results.size(); ++j) { - const auto &op = valid_operations[j]; - - if (!batch_get_results[j]) { - const auto error = batch_get_results[j].error(); - LOG(ERROR) << "BatchGet failed for key '" << op.key - << "': " << toString(error); - results[op.original_index] = tl::unexpected(error); - } - } - - return results; -} - -std::vector> -DistributedObjectStore::batch_put_from_internal( - const std::vector &keys, const std::vector &buffers, - const std::vector &sizes, const ReplicateConfig &config) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } - - if (keys.size() != buffers.size() || keys.size() != sizes.size()) { - LOG(ERROR) << "Mismatched sizes for keys, buffers, and sizes"; - return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } - - std::unordered_map> all_slices; - - // Create slices from user buffers - for (size_t i = 0; i < keys.size(); ++i) { - const std::string &key = keys[i]; - void *buffer = buffers[i]; - size_t size = sizes[i]; - - std::vector slices; - uint64_t offset = 0; - - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - - all_slices[key] = std::move(slices); - } - - std::vector> ordered_batched_slices; - ordered_batched_slices.reserve(keys.size()); - for (const auto &key : keys) { - auto it = all_slices.find(key); - if (it != all_slices.end()) { - ordered_batched_slices.emplace_back(it->second); - } else { - LOG(ERROR) << "Missing slices for key: " << key; - return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } - } - - // Call client BatchPut and return the vector directly - return client_->BatchPut(keys, ordered_batched_slices, config); -} - -std::vector> -DistributedObjectStore::batchIsExist_internal( - const std::vector &keys) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } - - if (keys.empty()) { - LOG(WARNING) << "Empty keys vector provided to batchIsExist_internal"; - return std::vector>(); - } - - // Call client BatchIsExist and return the vector directly - return client_->BatchIsExist(keys); -} - -int DistributedObjectStore::put_from_with_metadata( - const std::string &key, void *buffer, void *metadata_buffer, size_t size, - size_t metadata_size, const ReplicateConfig &config) { - // NOTE: The buffer address must be previously registered with - // register_buffer() for zero-copy RDMA operations to work correctly - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return -1; - } - - if (size == 0) { - LOG(WARNING) << "Attempting to put empty data for key: " << key; - return 0; - } - - // Create slices directly from the user buffer - std::vector slices; - // Add metadata slice - uint64_t metadata_offset = 0; - while (metadata_offset < metadata_size) { - auto metadata_chunk_size = - std::min(metadata_size - metadata_offset, kMaxSliceSize); - void *metadata_chunk_ptr = - static_cast(metadata_buffer) + metadata_offset; - slices.emplace_back(Slice{metadata_chunk_ptr, metadata_chunk_size}); - metadata_offset += metadata_chunk_size; - } - - uint64_t offset = 0; - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - auto put_result = client_->Put(key, slices, config); - if (!put_result) { - LOG(ERROR) << "Put operation failed with error: " - << toString(put_result.error()); - return -toInt(put_result.error()); - } - return 0; -} - -pybind11::object DistributedObjectStore::get_tensor(const std::string &key) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return pybind11::none(); - } - - try { - // Query object info first - auto query_result = client_->Query(key); - if (!query_result) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Query failed: " << query_result.error(); - return pybind11::none(); - } - - auto replica_list = query_result.value(); - if (replica_list.empty()) { - py::gil_scoped_acquire acquire_gil; - LOG(INFO) << "No replicas found for key: " << key; - return pybind11::none(); - } - - const auto &replica = replica_list[0]; - uint64_t total_length = calculate_total_size(replica); - - if (total_length == 0) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Failed to allocate slices for key: " << key; - return pybind11::none(); - } - - // Allocate buffer using the new allocator - auto alloc_result = client_buffer_allocator_->allocate(total_length); - if (!alloc_result) { - py::gil_scoped_acquire acquire_gil; - return pybind11::none(); - } - - auto &buffer_handle = *alloc_result; - - // Create slices for the allocated buffer - std::vector slices; - allocateSlices(slices, replica, buffer_handle); - - // Get the object data - auto get_result = client_->Get(key, replica_list, slices); - if (!get_result) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Get failed for key: " << key; - return pybind11::none(); - } - - // Create contiguous buffer and copy data - char *exported_data = new char[total_length]; - if (!exported_data) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Invalid data format: insufficient data for " - "metadata"; - return pybind11::none(); - } - TensorMetadata metadata; - - // Copy data from buffer to contiguous memory - memcpy(exported_data, buffer_handle.ptr(), total_length); - memcpy(&metadata, exported_data, sizeof(TensorMetadata)); - - if (metadata.ndim < 0 || metadata.ndim > 4) { - delete[] exported_data; - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Invalid tensor metadata: ndim=" << metadata.ndim; - return pybind11::none(); - } - - TensorDtype dtype_enum = static_cast(metadata.dtype); - if (dtype_enum == TensorDtype::UNKNOWN) { - delete[] exported_data; - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Unknown tensor dtype!"; - return pybind11::none(); - } - - size_t tensor_size = total_length - sizeof(TensorMetadata); - if (tensor_size == 0) { - delete[] exported_data; - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Invalid data format: no tensor data found"; - return pybind11::none(); - } - - // Convert bytes to tensor using torch.from_numpy - pybind11::object np_array; - int dtype_index = static_cast(dtype_enum); - if (dtype_index >= 0 && - dtype_index < static_cast(array_creators.size())) { - np_array = array_creators[dtype_index]( - exported_data, sizeof(TensorMetadata), tensor_size); - } else { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Unsupported dtype enum: " << dtype_index; - return pybind11::none(); - } - - if (metadata.ndim > 0) { - std::vector shape_vec; - for (int i = 0; i < metadata.ndim; i++) { - shape_vec.push_back(metadata.shape[i]); - } - py::tuple shape_tuple = py::cast(shape_vec); - np_array = np_array.attr("reshape")(shape_tuple); - } - py::gil_scoped_acquire acquire_gil; - pybind11::object tensor = torch.attr("from_numpy")(np_array); - return tensor; - - } catch (const pybind11::error_already_set &e) { - py::gil_scoped_acquire acquire_gil; - LOG(ERROR) << "Failed to get tensor data: " << e.what(); - return pybind11::none(); - } -} - -tl::expected DistributedObjectStore::put_tensor_internal( - const std::string &key, pybind11::object tensor) { - if (!client_) { - LOG(ERROR) << "Client is not initialized"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - try { - if (!(tensor.attr("__class__") - .attr("__name__") - .cast() - .find("Tensor") != std::string::npos)) { - LOG(ERROR) << "Input is not a PyTorch tensor"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - uintptr_t data_ptr = tensor.attr("data_ptr")().cast(); - size_t numel = tensor.attr("numel")().cast(); - size_t element_size = tensor.attr("element_size")().cast(); - size_t tensor_size = numel * element_size; - - pybind11::object shape_obj = tensor.attr("shape"); - pybind11::object dtype_obj = tensor.attr("dtype"); - - TensorDtype dtype_enum = get_tensor_dtype(dtype_obj); - if (dtype_enum == TensorDtype::UNKNOWN) { - LOG(ERROR) << "Unsupported tensor dtype!"; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - pybind11::tuple shape_tuple = - pybind11::cast(shape_obj); - int32_t ndim = static_cast(shape_tuple.size()); - if (ndim > 4) { - LOG(ERROR) << "Tensor has more than 4 dimensions: " << ndim; - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - TensorMetadata metadata; - metadata.dtype = static_cast(dtype_enum); - metadata.ndim = ndim; - - for (int i = 0; i < 4; i++) { - if (i < ndim) { - metadata.shape[i] = shape_tuple[i].cast(); + // Convert bytes to tensor using torch.from_numpy + pybind11::object np_array; + int dtype_index = static_cast(dtype_enum); + if (dtype_index >= 0 && + dtype_index < static_cast(array_creators.size())) { + np_array = array_creators[dtype_index]( + exported_data, sizeof(TensorMetadata), tensor_size); } else { - metadata.shape[i] = -1; + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Unsupported dtype enum: " << dtype_index; + return pybind11::none(); } + + if (metadata.ndim > 0) { + std::vector shape_vec; + for (int i = 0; i < metadata.ndim; i++) { + shape_vec.push_back(metadata.shape[i]); + } + py::tuple shape_tuple = py::cast(shape_vec); + np_array = np_array.attr("reshape")(shape_tuple); + } + py::gil_scoped_acquire acquire_gil; + pybind11::object tensor = torch.attr("from_numpy")(np_array); + return tensor; + + } catch (const pybind11::error_already_set &e) { + py::gil_scoped_acquire acquire_gil; + LOG(ERROR) << "Failed to get tensor data: " << e.what(); + return pybind11::none(); } - - char *buffer = reinterpret_cast(data_ptr); - char *metadata_buffer = reinterpret_cast(&metadata); - std::vector> values; - values.emplace_back( - std::span(metadata_buffer, sizeof(TensorMetadata))); - values.emplace_back(std::span(buffer, tensor_size)); - - auto register_result = register_buffer_internal( - reinterpret_cast(data_ptr), tensor_size); - if (!register_result) { - return tl::unexpected(register_result.error()); - } - - // Use put_parts to put metadata and tensor together - auto put_result = this->put_parts_internal(key, values); - - auto unregister_result = - unregister_buffer_internal(reinterpret_cast(data_ptr)); - if (!unregister_result) { - LOG(WARNING) << "Failed to unregister buffer after put_tensor"; - } - - if (!put_result) { - return tl::unexpected(put_result.error()); - } - - return {}; - } catch (const pybind11::error_already_set &e) { - LOG(ERROR) << "Failed to access tensor data: " << e.what(); - return tl::unexpected(ErrorCode::INVALID_PARAMS); } -} -int DistributedObjectStore::put_tensor(const std::string &key, - pybind11::object tensor) { - return to_py_ret(put_tensor_internal(key, tensor)); -} + int put_tensor(const std::string &key, pybind11::object tensor) { + if (!store_.client_) { + LOG(ERROR) << "Client is not initialized"; + return -static_cast(ErrorCode::INVALID_PARAMS); + } + try { + if (!(tensor.attr("__class__") + .attr("__name__") + .cast() + .find("Tensor") != std::string::npos)) { + LOG(ERROR) << "Input is not a PyTorch tensor"; + return -static_cast(ErrorCode::INVALID_PARAMS); + } + + uintptr_t data_ptr = tensor.attr("data_ptr")().cast(); + size_t numel = tensor.attr("numel")().cast(); + size_t element_size = tensor.attr("element_size")().cast(); + size_t tensor_size = numel * element_size; + + pybind11::object shape_obj = tensor.attr("shape"); + pybind11::object dtype_obj = tensor.attr("dtype"); + + TensorDtype dtype_enum = get_tensor_dtype(dtype_obj); + if (dtype_enum == TensorDtype::UNKNOWN) { + LOG(ERROR) << "Unsupported tensor dtype!"; + return -static_cast(ErrorCode::INVALID_PARAMS); + } + + pybind11::tuple shape_tuple = + pybind11::cast(shape_obj); + int32_t ndim = static_cast(shape_tuple.size()); + if (ndim > 4) { + LOG(ERROR) << "Tensor has more than 4 dimensions: " << ndim; + return -static_cast(ErrorCode::INVALID_PARAMS); + } + + TensorMetadata metadata; + metadata.dtype = static_cast(dtype_enum); + metadata.ndim = ndim; + + for (int i = 0; i < 4; i++) { + if (i < ndim) { + metadata.shape[i] = shape_tuple[i].cast(); + } else { + metadata.shape[i] = -1; + } + } + + char *buffer = reinterpret_cast(data_ptr); + char *metadata_buffer = reinterpret_cast(&metadata); + std::vector> values; + values.emplace_back( + std::span(metadata_buffer, sizeof(TensorMetadata))); + values.emplace_back(std::span(buffer, tensor_size)); + + auto register_result = store_.register_buffer_internal( + reinterpret_cast(data_ptr), tensor_size); + if (!register_result) { + return -static_cast(register_result.error()); + } + + // Use put_parts to put metadata and tensor together + auto put_result = store_.put_parts_internal(key, values); + + auto unregister_result = store_.unregister_buffer_internal( + reinterpret_cast(data_ptr)); + if (!unregister_result) { + LOG(WARNING) << "Failed to unregister buffer after put_tensor"; + } + + if (!put_result) { + return -static_cast(put_result.error()); + } + + return 0; + } catch (const pybind11::error_already_set &e) { + LOG(ERROR) << "Failed to access tensor data: " << e.what(); + return -static_cast(ErrorCode::INVALID_PARAMS); + } + } +}; PYBIND11_MODULE(store, m) { // Define the ReplicateConfig class @@ -1538,18 +288,18 @@ PYBIND11_MODULE(store, m) { return oss.str(); }); - // Define the SliceBuffer class - py::class_>(m, "SliceBuffer", - py::buffer_protocol()) + // Define the BufferHandle class + py::class_>( + m, "BufferHandle", py::buffer_protocol()) .def("ptr", - [](const SliceBuffer &self) { + [](const BufferHandle &self) { // Return the pointer as an integer for Python return reinterpret_cast(self.ptr()); }) - .def("size", &SliceBuffer::size) - .def("__len__", &SliceBuffer::size) - .def_buffer([](SliceBuffer &self) -> py::buffer_info { - // SliceBuffer now always contains contiguous memory + .def("size", &BufferHandle::size) + .def("__len__", &BufferHandle::size) + .def_buffer([](BufferHandle &self) -> py::buffer_info { + // BufferHandle now always contains contiguous memory if (self.size() > 0) { return py::buffer_info( self.ptr(), /* Pointer to buffer */ @@ -1576,73 +326,121 @@ PYBIND11_MODULE(store, m) { } }); - // Define the DistributedObjectStore class - py::class_(m, "MooncakeDistributedStore") + // Create a wrapper that exposes DistributedObjectStore with Python-specific + // methods + py::class_(m, "MooncakeDistributedStore") .def(py::init<>()) - .def("setup", &DistributedObjectStore::setup) - .def("init_all", &DistributedObjectStore::initAll) - .def("get", &DistributedObjectStore::get) - .def("get_batch", &DistributedObjectStore::get_batch) - .def("get_buffer", &DistributedObjectStore::get_buffer, - py::call_guard(), - py::return_value_policy::take_ownership) - .def("batch_get_buffer", &DistributedObjectStore::batch_get_buffer, - py::call_guard(), - py::return_value_policy::take_ownership) - .def("remove", &DistributedObjectStore::remove, - py::call_guard()) - .def("remove_all", &DistributedObjectStore::removeAll, - py::call_guard()) - .def("is_exist", &DistributedObjectStore::isExist, - py::call_guard()) - .def("batch_is_exist", &DistributedObjectStore::batchIsExist, - py::call_guard(), py::arg("keys"), - "Check if multiple objects exist. Returns list of " - "results: 1 if " - "exists, 0 if not exists, -1 if error") - .def("close", &DistributedObjectStore::tearDownAll) - .def("get_size", &DistributedObjectStore::getSize, - py::call_guard()) - .def("get_tensor", &DistributedObjectStore::get_tensor, py::arg("key"), + .def("setup", + [](MooncakeStorePyWrapper &self, 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") { + return self.store_.setup(local_hostname, metadata_server, + global_segment_size, + local_buffer_size, protocol, + rdma_devices, master_server_addr); + }) + .def("init_all", + [](MooncakeStorePyWrapper &self, const std::string &protocol, + const std::string &device_name, + size_t mount_segment_size = 1024 * 1024 * 16) { + return self.store_.initAll(protocol, device_name, + mount_segment_size); + }) + .def("get", &MooncakeStorePyWrapper::get) + .def("get_batch", &MooncakeStorePyWrapper::get_batch) + .def( + "get_buffer", + [](MooncakeStorePyWrapper &self, const std::string &key) { + py::gil_scoped_release release; + return self.store_.get_buffer(key); + }, + py::return_value_policy::take_ownership) + .def( + "batch_get_buffer", + [](MooncakeStorePyWrapper &self, + const std::vector &keys) { + py::gil_scoped_release release; + return self.store_.batch_get_buffer(keys); + }, + py::return_value_policy::take_ownership) + .def("remove", + [](MooncakeStorePyWrapper &self, const std::string &key) { + py::gil_scoped_release release; + return self.store_.remove(key); + }) + .def("remove_all", + [](MooncakeStorePyWrapper &self) { + py::gil_scoped_release release; + return self.store_.removeAll(); + }) + .def("is_exist", + [](MooncakeStorePyWrapper &self, const std::string &key) { + py::gil_scoped_release release; + return self.store_.isExist(key); + }) + .def( + "batch_is_exist", + [](MooncakeStorePyWrapper &self, + const std::vector &keys) { + py::gil_scoped_release release; + return self.store_.batchIsExist(keys); + }, + py::arg("keys"), + "Check if multiple objects exist. Returns list of results: 1 if " + "exists, 0 if not exists, -1 if error") + .def("close", + [](MooncakeStorePyWrapper &self) { + return self.store_.tearDownAll(); + }) + .def("get_size", + [](MooncakeStorePyWrapper &self, const std::string &key) { + py::gil_scoped_release release; + return self.store_.getSize(key); + }) + .def("get_tensor", &MooncakeStorePyWrapper::get_tensor, py::arg("key"), "Get a PyTorch tensor from the store") - .def("put_tensor", &DistributedObjectStore::put_tensor, py::arg("key"), + .def("put_tensor", &MooncakeStorePyWrapper::put_tensor, py::arg("key"), py::arg("tensor"), "Put a PyTorch tensor into the store") .def( "register_buffer", - [](DistributedObjectStore &self, uintptr_t buffer_ptr, + [](MooncakeStorePyWrapper &self, uintptr_t buffer_ptr, size_t size) { // Register memory buffer for RDMA operations void *buffer = reinterpret_cast(buffer_ptr); py::gil_scoped_release release; - return self.register_buffer(buffer, size); + return self.store_.register_buffer(buffer, size); }, py::arg("buffer_ptr"), py::arg("size"), "Register a memory buffer for direct access operations") .def( "unregister_buffer", - [](DistributedObjectStore &self, uintptr_t buffer_ptr) { + [](MooncakeStorePyWrapper &self, uintptr_t buffer_ptr) { // Unregister memory buffer void *buffer = reinterpret_cast(buffer_ptr); py::gil_scoped_release release; - return self.unregister_buffer(buffer); + return self.store_.unregister_buffer(buffer); }, py::arg("buffer_ptr"), "Unregister a previously registered memory " "buffer for direct access operations") .def( "get_into", - [](DistributedObjectStore &self, const std::string &key, + [](MooncakeStorePyWrapper &self, const std::string &key, uintptr_t buffer_ptr, size_t size) { // Get data directly into user-provided buffer void *buffer = reinterpret_cast(buffer_ptr); py::gil_scoped_release release; - return self.get_into(key, buffer, size); + return self.store_.get_into(key, buffer, size); }, py::arg("key"), py::arg("buffer_ptr"), py::arg("size"), "Get object data directly into a pre-allocated buffer") .def( "batch_get_into", - [](DistributedObjectStore &self, + [](MooncakeStorePyWrapper &self, const std::vector &keys, const std::vector &buffer_ptrs, const std::vector &sizes) { @@ -1652,7 +450,7 @@ PYBIND11_MODULE(store, m) { buffers.push_back(reinterpret_cast(ptr)); } py::gil_scoped_release release; - return self.batch_get_into(keys, buffers, sizes); + return self.store_.batch_get_into(keys, buffers, sizes); }, py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"), "Get object data directly into pre-allocated buffers for " @@ -1660,20 +458,20 @@ PYBIND11_MODULE(store, m) { "keys") .def( "put_from", - [](DistributedObjectStore &self, const std::string &key, + [](MooncakeStorePyWrapper &self, const std::string &key, uintptr_t buffer_ptr, size_t size, const ReplicateConfig &config = ReplicateConfig{}) { // Put data directly from user-provided buffer void *buffer = reinterpret_cast(buffer_ptr); py::gil_scoped_release release; - return self.put_from(key, buffer, size, config); + return self.store_.put_from(key, buffer, size, config); }, py::arg("key"), py::arg("buffer_ptr"), py::arg("size"), py::arg("config") = ReplicateConfig{}, "Put object data directly from a pre-allocated buffer") .def( "put_from_with_metadata", - [](DistributedObjectStore &self, const std::string &key, + [](MooncakeStorePyWrapper &self, const std::string &key, uintptr_t buffer_ptr, uintptr_t metadata_buffer_ptr, size_t size, size_t metadata_size, const ReplicateConfig &config = ReplicateConfig{}) { @@ -1683,8 +481,8 @@ PYBIND11_MODULE(store, m) { void *metadata_buffer = reinterpret_cast(metadata_buffer_ptr); py::gil_scoped_release release; - return self.put_from_with_metadata(key, buffer, metadata_buffer, - size, metadata_size, config); + return self.store_.put_from_with_metadata( + key, buffer, metadata_buffer, size, metadata_size, config); }, py::arg("key"), py::arg("buffer_ptr"), py::arg("metadata_buffer_ptr"), py::arg("size"), @@ -1693,7 +491,7 @@ PYBIND11_MODULE(store, m) { "metadata") .def( "batch_put_from", - [](DistributedObjectStore &self, + [](MooncakeStorePyWrapper &self, const std::vector &keys, const std::vector &buffer_ptrs, const std::vector &sizes, @@ -1704,7 +502,7 @@ PYBIND11_MODULE(store, m) { buffers.push_back(reinterpret_cast(ptr)); } py::gil_scoped_release release; - return self.batch_put_from(keys, buffers, sizes, config); + return self.store_.batch_put_from(keys, buffers, sizes, config); }, py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"), py::arg("config") = ReplicateConfig{}, @@ -1713,12 +511,12 @@ PYBIND11_MODULE(store, m) { "keys") .def( "put", - [](DistributedObjectStore &self, const std::string &key, + [](MooncakeStorePyWrapper &self, const std::string &key, py::buffer buf, const ReplicateConfig &config = ReplicateConfig{}) { py::buffer_info info = buf.request(/*writable=*/false); py::gil_scoped_release release; - return self.put( + return self.store_.put( key, std::span(static_cast(info.ptr), static_cast(info.size)), @@ -1728,7 +526,7 @@ PYBIND11_MODULE(store, m) { py::arg("config") = ReplicateConfig{}) .def( "put_parts", - [](DistributedObjectStore &self, const std::string &key, + [](MooncakeStorePyWrapper &self, const std::string &key, py::args parts, const ReplicateConfig &config = ReplicateConfig{}) { // 1) Python buffer → span @@ -1751,12 +549,12 @@ PYBIND11_MODULE(store, m) { // 2) Call C++ function py::gil_scoped_release unlock; - return self.put_parts(key, spans, config); + return self.store_.put_parts(key, spans, config); }, py::arg("key"), py::arg("config") = ReplicateConfig{}) .def( "put_batch", - [](DistributedObjectStore &self, + [](MooncakeStorePyWrapper &self, const std::vector &keys, const std::vector &buffers, const ReplicateConfig &config = ReplicateConfig{}) { @@ -1774,11 +572,13 @@ PYBIND11_MODULE(store, m) { } py::gil_scoped_release release; - return self.put_batch(keys, spans, config); + return self.store_.put_batch(keys, spans, config); }, py::arg("keys"), py::arg("values"), py::arg("config") = ReplicateConfig{}) - .def("get_hostname", &DistributedObjectStore::get_hostname); + .def("get_hostname", [](MooncakeStorePyWrapper &self) { + return self.store_.get_hostname(); + }); } -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-integration/store/store_py.h b/mooncake-store/include/pybind_client.h similarity index 82% rename from mooncake-integration/store/store_py.h rename to mooncake-store/include/pybind_client.h index f5df34e7..85bd2af5 100644 --- a/mooncake-integration/store/store_py.h +++ b/mooncake-store/include/pybind_client.h @@ -1,8 +1,5 @@ #pragma once -#include -#include - #include #include #include @@ -13,10 +10,7 @@ namespace mooncake { -class DistributedObjectStore; - -// Forward declarations -class SliceBuffer; +class PyClient; template constexpr bool is_supported_return_type_v = @@ -45,10 +39,10 @@ class ResourceTracker { static ResourceTracker &getInstance(); // Register a DistributedObjectStore instance for cleanup - void registerInstance(DistributedObjectStore *instance); + void registerInstance(PyClient *instance); // Unregister a DistributedObjectStore instance - void unregisterInstance(DistributedObjectStore *instance); + void unregisterInstance(PyClient *instance); private: ResourceTracker(); @@ -68,29 +62,13 @@ class ResourceTracker { static void exitHandler(); std::mutex mutex_; - std::unordered_set instances_; + std::unordered_set instances_; }; -/** - * @brief A class that holds a contiguous buffer of data - * This class is responsible for freeing the buffer when it's destroyed (RAII) - */ -class SliceBuffer { +class PyClient { public: - SliceBuffer(BufferHandle handle); - - void *ptr() const; - uint64_t size() const; - - private: - BufferHandle handle_; -}; - -class DistributedObjectStore { - public: - friend class SliceBuffer; // Allow SliceBuffer to access private members - DistributedObjectStore(); - ~DistributedObjectStore(); + PyClient(); + ~PyClient(); int setup(const std::string &local_hostname, const std::string &metadata_server, @@ -198,26 +176,21 @@ class DistributedObjectStore { [[nodiscard]] std::string get_hostname() const; - pybind11::bytes get(const std::string &key); - - std::vector get_batch( - const std::vector &keys); - /** * @brief Get a buffer containing the data for a key * @param key Key to get data for - * @return std::shared_ptr Buffer containing the data, or + * @return std::shared_ptr Buffer containing the data, or * nullptr if error */ - std::shared_ptr get_buffer(const std::string &key); + std::shared_ptr get_buffer(const std::string &key); /** * @brief Get buffers containing the data for multiple keys (batch version) * @param keys Vector of keys to get data for - * @return Vector of std::shared_ptr buffers containing the + * @return Vector of std::shared_ptr buffers containing the * data, or nullptr for each key if error */ - std::vector> batch_get_buffer( + std::vector> batch_get_buffer( const std::vector &keys); int remove(const std::string &key); @@ -249,20 +222,6 @@ class DistributedObjectStore { */ int64_t getSize(const std::string &key); - /** - * @brief Get a PyTorch tensor from the store - * @param key Key of the tensor to get - * @return PyTorch tensor, or nullptr if error or tensor doesn't exist - */ - pybind11::object get_tensor(const std::string &key); - /** - * @brief Put a PyTorch tensor into the store - * @param key Key for the tensor - * @param tensor PyTorch tensor to store - * @return 0 on success, negative value on error - */ - int put_tensor(const std::string &key, pybind11::object tensor); - // Internal versions that return tl::expected tl::expected setup_internal( const std::string &local_hostname, const std::string &metadata_server, @@ -311,11 +270,6 @@ class DistributedObjectStore { const std::vector> &values, const ReplicateConfig &config = ReplicateConfig{}); - tl::expected put_batch_internal( - const std::vector &keys, - const std::vector &buffers, - const ReplicateConfig &config = ReplicateConfig{}); - tl::expected remove_internal(const std::string &key); tl::expected removeAll_internal(); @@ -329,10 +283,7 @@ class DistributedObjectStore { tl::expected getSize_internal(const std::string &key); - tl::expected put_tensor_internal(const std::string &key, - pybind11::object tensor); - - std::vector> batch_get_buffer_internal( + std::vector> batch_get_buffer_internal( const std::vector &keys); std::shared_ptr client_ = nullptr; @@ -352,4 +303,4 @@ class DistributedObjectStore { std::string local_hostname; }; -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 4cc1da7c..6618e76e 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -111,4 +111,21 @@ void** rdma_args(const std::string& device_name); return oss.str(); } -} // namespace mooncake \ No newline at end of file +// Network utility functions + +/** + * @brief Check if a TCP port is available for binding + * @param port The port number to check + * @return true if port is available, false otherwise + */ +bool isPortAvailable(int port); + +/** + * @brief Get a random available port in the specified range + * @param min_port Minimum port number (default: 12300) + * @param max_port Maximum port number (default: 14300) + * @return Available port number, or -1 if no port found after 10 attempts + */ +int getRandomAvailablePort(int min_port = 12300, int max_port = 14300); + +} // namespace mooncake diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 3f5569cd..fab7df13 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -21,6 +21,7 @@ set(MOONCAKE_STORE_SOURCES offset_allocator.cpp posix_file.cpp client_buffer.cpp + pybind_client.cpp ) set(EXTRA_LIBS "") diff --git a/mooncake-store/src/pybind_client.cpp b/mooncake-store/src/pybind_client.cpp new file mode 100644 index 00000000..d6f6beae --- /dev/null +++ b/mooncake-store/src/pybind_client.cpp @@ -0,0 +1,1080 @@ +#include "pybind_client.h" + +#include +#include +#include + +#include // for atexit +#include + +#include "client_buffer.hpp" +#include "config.h" +#include "types.h" +#include "utils.h" + +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(PyClient *instance) { + std::lock_guard lock(mutex_); + instances_.insert(instance); +} + +void ResourceTracker::unregisterInstance(PyClient *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_) { + PyClient *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() { getInstance().cleanupAllResources(); } + +PyClient::PyClient() { + // Register this instance with the global tracker + easylog::set_min_severity(easylog::Severity::WARN); + ResourceTracker::getInstance().registerInstance(this); +} + +PyClient::~PyClient() { + // Unregister from the tracker before cleanup + ResourceTracker::getInstance().unregisterInstance(this); +} + +tl::expected PyClient::setup_internal( + 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 tl::unexpected(ErrorCode::INVALID_PARAMS); + } + // Combine hostname with port + this->local_hostname = hostname + ":" + std::to_string(port); + } else { + this->local_hostname = local_hostname; + } + + void **args = (protocol == "rdma") ? rdma_args(rdma_devices) : nullptr; + auto client_opt = + mooncake::Client::Create(this->local_hostname, metadata_server, + protocol, args, master_server_addr); + if (!client_opt) { + LOG(ERROR) << "Failed to create client"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + client_ = *client_opt; + + // Local_buffer_size is allowed to be 0, but we only register memory when + // local_buffer_size > 0. Invoke ibv_reg_mr() with size=0 is UB, and may + // fail in some rdma implementations. + client_buffer_allocator_ = ClientBufferAllocator::create(local_buffer_size); + if (local_buffer_size > 0) { + auto result = client_->RegisterLocalMemory( + client_buffer_allocator_->getBase(), local_buffer_size, + kWildcardLocation, false, true); + if (!result.has_value()) { + LOG(ERROR) << "Failed to register local memory: " + << toString(result.error()); + return tl::unexpected(result.error()); + } + } else { + LOG(INFO) << "Local buffer size is 0, skip registering local memory"; + } + + // If global_segment_size is 0, skip mount segment; + // If global_segment_size is larger than max_mr_size, split to multiple + // segments. + auto max_mr_size = globalConfig().max_mr_size; // Max segment size + uint64_t total_glbseg_size = global_segment_size; // For logging + uint64_t current_glbseg_size = 0; // For logging + while (global_segment_size > 0) { + size_t segment_size = std::min(global_segment_size, max_mr_size); + global_segment_size -= segment_size; + current_glbseg_size += segment_size; + LOG(INFO) << "Mounting segment: " << segment_size << " bytes, " + << current_glbseg_size << " of " << total_glbseg_size; + void *ptr = allocate_buffer_allocator_memory(segment_size); + if (!ptr) { + LOG(ERROR) << "Failed to allocate segment memory"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + segment_ptrs_.emplace_back(ptr); + auto mount_result = client_->MountSegment(ptr, segment_size); + if (!mount_result.has_value()) { + LOG(ERROR) << "Failed to mount segment: " + << toString(mount_result.error()); + return tl::unexpected(mount_result.error()); + } + } + if (total_glbseg_size == 0) { + LOG(INFO) << "Global segment size is 0, skip mounting segment"; + } + + return {}; +} + +int PyClient::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) { + return to_py_ret(setup_internal( + local_hostname, metadata_server, global_segment_size, local_buffer_size, + protocol, rdma_devices, master_server_addr)); +} + +tl::expected PyClient::initAll_internal( + const std::string &protocol_, const std::string &device_name, + size_t mount_segment_size) { + if (client_) { + LOG(ERROR) << "Client is already initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + uint64_t buffer_allocator_size = 1024 * 1024 * 1024; + return setup_internal("localhost:12345", "127.0.0.1:2379", + mount_segment_size, buffer_allocator_size, protocol_, + device_name); +} + +int PyClient::initAll(const std::string &protocol_, + const std::string &device_name, + size_t mount_segment_size) { + return to_py_ret( + initAll_internal(protocol_, device_name, mount_segment_size)); +} + +tl::expected PyClient::tearDownAll_internal() { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + // Reset all resources + client_.reset(); + client_buffer_allocator_.reset(); + segment_ptrs_.clear(); + local_hostname = ""; + device_name = ""; + protocol = ""; + return {}; +} + +int PyClient::tearDownAll() { return to_py_ret(tearDownAll_internal()); } + +tl::expected PyClient::put_internal( + const std::string &key, std::span value, + const ReplicateConfig &config) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto alloc_result = client_buffer_allocator_->allocate(value.size_bytes()); + if (!alloc_result) { + LOG(ERROR) << "Failed to allocate buffer for put operation, key: " + << key << ", value size: " << value.size(); + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto &buffer_handle = *alloc_result; + memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); + + std::vector slices = split_into_slices(buffer_handle); + + auto put_result = client_->Put(key, slices, config); + if (!put_result) { + LOG(ERROR) << "Put operation failed with error: " + << toString(put_result.error()); + return tl::unexpected(put_result.error()); + } + + return {}; +} + +int PyClient::put(const std::string &key, std::span value, + const ReplicateConfig &config) { + return to_py_ret(put_internal(key, value, config)); +} + +tl::expected PyClient::put_batch_internal( + const std::vector &keys, + const std::vector> &values, + const ReplicateConfig &config) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + if (keys.size() != values.size()) { + LOG(ERROR) << "Key and value size mismatch"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + std::vector buffer_handles; + std::unordered_map> batched_slices; + batched_slices.reserve(keys.size()); + + for (size_t i = 0; i < keys.size(); ++i) { + auto &key = keys[i]; + auto &value = values[i]; + auto alloc_result = + client_buffer_allocator_->allocate(value.size_bytes()); + if (!alloc_result) { + LOG(ERROR) + << "Failed to allocate buffer for put_batch operation, key: " + << key << ", value size: " << value.size(); + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto &buffer_handle = *alloc_result; + memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); + auto slices = split_into_slices(buffer_handle); + buffer_handles.emplace_back(std::move(*alloc_result)); + batched_slices.emplace(key, std::move(slices)); + } + + // Convert unordered_map to vector format expected by BatchPut + std::vector> ordered_batched_slices; + ordered_batched_slices.reserve(keys.size()); + for (const auto &key : keys) { + auto it = batched_slices.find(key); + if (it != batched_slices.end()) { + ordered_batched_slices.emplace_back(it->second); + } else { + LOG(ERROR) << "Missing slices for key: " << key; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + } + + auto results = client_->BatchPut(keys, ordered_batched_slices, config); + + // Check if any operations failed + for (size_t i = 0; i < results.size(); ++i) { + if (!results[i]) { + LOG(ERROR) << "BatchPut operation failed for key '" << keys[i] + << "' with error: " << toString(results[i].error()); + return tl::unexpected(results[i].error()); + } + } + return {}; +} + +int PyClient::put_batch(const std::vector &keys, + const std::vector> &values, + const ReplicateConfig &config) { + return to_py_ret(put_batch_internal(keys, values, config)); +} + +tl::expected PyClient::put_parts_internal( + const std::string &key, std::vector> values, + const ReplicateConfig &config) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + // Calculate total size needed + size_t total_size = 0; + for (const auto &value : values) { + total_size += value.size_bytes(); + } + + if (total_size == 0) { + LOG(WARNING) << "Attempting to put empty data for key: " << key; + return {}; + } + + // Allocate buffer using the new allocator + auto alloc_result = client_buffer_allocator_->allocate(total_size); + if (!alloc_result) { + LOG(ERROR) << "Failed to allocate buffer for put_parts operation, key: " + << key << ", total size: " << total_size; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + auto &buffer_handle = *alloc_result; + + // Copy all parts into the contiguous buffer + size_t offset = 0; + for (const auto &value : values) { + memcpy(static_cast(buffer_handle.ptr()) + offset, value.data(), + value.size_bytes()); + offset += value.size_bytes(); + } + + // Split into slices + std::vector slices = split_into_slices(buffer_handle); + + // Perform the put operation - buffer_handle will be automatically released + auto put_result = client_->Put(key, slices, config); + if (!put_result) { + LOG(ERROR) << "Put operation failed with error: " + << toString(put_result.error()); + return tl::unexpected(put_result.error()); + } + + return {}; +} + +int PyClient::put_parts(const std::string &key, + std::vector> values, + const ReplicateConfig &config) { + return to_py_ret(put_parts_internal(key, values, config)); +} + +tl::expected PyClient::remove_internal( + const std::string &key) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto remove_result = client_->Remove(key); + if (!remove_result) { + return tl::unexpected(remove_result.error()); + } + return {}; +} + +int PyClient::remove(const std::string &key) { + return to_py_ret(remove_internal(key)); +} + +tl::expected PyClient::removeAll_internal() { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + return client_->RemoveAll(); +} + +long PyClient::removeAll() { return to_py_ret(removeAll_internal()); } + +tl::expected PyClient::isExist_internal( + const std::string &key) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + return client_->IsExist(key); +} + +int PyClient::isExist(const std::string &key) { + auto result = isExist_internal(key); + + if (result.has_value()) { + return *result ? 1 : 0; // 1 if exists, 0 if not + } else { + return toInt(result.error()); + } +} + +std::vector PyClient::batchIsExist(const std::vector &keys) { + auto internal_results = batchIsExist_internal(keys); + std::vector results; + results.reserve(internal_results.size()); + + for (const auto &result : internal_results) { + if (result.has_value()) { + results.push_back(result.value() ? 1 : 0); // 1 if exists, 0 if not + } else { + results.push_back(toInt(result.error())); + } + } + + return results; +} + +tl::expected PyClient::getSize_internal( + const std::string &key) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + auto query_result = client_->Query(key); + + if (!query_result) { + return tl::unexpected(query_result.error()); + } + + auto replica_list = query_result.value(); + + // Calculate total size from all replicas' handles + int64_t total_size = 0; + if (!replica_list.empty()) { + auto &replica = replica_list[0]; + total_size = calculate_total_size(replica); + } else { + LOG(ERROR) << "Internal error: replica_list is empty"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); // Internal error + } + + return total_size; +} + +int64_t PyClient::getSize(const std::string &key) { + return to_py_ret(getSize_internal(key)); +} + +// Implementation of get_buffer method +std::shared_ptr PyClient::get_buffer(const std::string &key) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return nullptr; + } + + // Query the object info + auto query_result = client_->Query(key); + if (!query_result) { + if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND) { + return nullptr; + } + LOG(ERROR) << "Query failed for key: " << key + << " with error: " << toString(query_result.error()); + return nullptr; + } + + auto replica_list = query_result.value(); + if (replica_list.empty()) { + LOG(ERROR) << "Empty replica list for key: " << key; + return nullptr; + } + + const auto &replica = replica_list[0]; + uint64_t total_length = calculate_total_size(replica); + + if (total_length == 0) { + return nullptr; + } + + // Allocate buffer using the new allocator + auto alloc_result = client_buffer_allocator_->allocate(total_length); + if (!alloc_result) { + LOG(ERROR) << "Failed to allocate buffer for get_buffer, key: " << key; + return nullptr; + } + + auto &buffer_handle = *alloc_result; + + // Create slices for the allocated buffer + std::vector slices; + allocateSlices(slices, replica, buffer_handle); + + // Get the object data + auto get_result = client_->Get(key, replica_list, slices); + if (!get_result) { + LOG(ERROR) << "Get failed for key: " << key + << " with error: " << toString(get_result.error()); + return nullptr; + } + + // Create BufferHandle with the allocated memory + // The buffer will be managed by the BufferHandle's shared_ptr + return std::make_shared(std::move(buffer_handle)); +} + +// Implementation of batch_get_buffer_internal method +std::vector> PyClient::batch_get_buffer_internal( + const std::vector &keys) { + std::vector> final_results(keys.size(), + nullptr); + + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return final_results; + } + + if (keys.empty()) { + return final_results; + } + + // 1. Query metadata for all keys + auto query_results = client_->BatchQuery(keys); + + // 2. Prepare for batch get: filter valid keys and prepare buffers + struct KeyOp { + size_t original_index; + std::string key; + std::vector replica_list; + std::unique_ptr buffer_handle; + std::vector slices; + }; + std::vector valid_ops; + valid_ops.reserve(keys.size()); + + for (size_t i = 0; i < keys.size(); ++i) { + const auto &key = keys[i]; + + if (!query_results[i]) { + if (query_results[i].error() != ErrorCode::OBJECT_NOT_FOUND) { + LOG(ERROR) << "Query failed for key '" << key + << "': " << toString(query_results[i].error()); + } + continue; + } + + auto replica_list = query_results[i].value(); + if (replica_list.empty()) { + LOG(ERROR) << "Empty replica list for key: " << key; + continue; + } + + const auto &replica = replica_list[0]; + uint64_t total_size = calculate_total_size(replica); + if (total_size == 0) { + continue; + } + + auto alloc_result = client_buffer_allocator_->allocate(total_size); + if (!alloc_result) { + LOG(ERROR) << "Failed to allocate buffer for key: " << key; + continue; + } + + auto buffer_handle = + std::make_unique(std::move(*alloc_result)); + std::vector slices; + allocateSlices(slices, replica, *buffer_handle); + + valid_ops.emplace_back(KeyOp{.original_index = i, + .key = key, + .replica_list = std::move(replica_list), + .buffer_handle = std::move(buffer_handle), + .slices = std::move(slices)}); + } + + if (valid_ops.empty()) { + return final_results; + } + + // 3. Execute batch get + std::vector batch_keys; + std::vector> batch_replica_lists; + std::unordered_map> batch_slices; + batch_keys.reserve(valid_ops.size()); + batch_replica_lists.reserve(valid_ops.size()); + + for (auto &op : valid_ops) { + batch_keys.push_back(op.key); + batch_replica_lists.push_back(op.replica_list); + batch_slices[op.key] = op.slices; + } + + auto batch_get_results = + client_->BatchGet(batch_keys, batch_replica_lists, batch_slices); + + // 4. Process results and create BufferHandles + for (size_t i = 0; i < valid_ops.size(); ++i) { + if (batch_get_results[i]) { + auto &op = valid_ops[i]; + final_results[op.original_index] = + std::make_shared(std::move(*op.buffer_handle)); + } else { + LOG(ERROR) << "BatchGet failed for key '" << valid_ops[i].key + << "': " << toString(batch_get_results[i].error()); + } + } + + return final_results; +} + +// Implementation of batch_get_buffer method +std::vector> PyClient::batch_get_buffer( + const std::vector &keys) { + return batch_get_buffer_internal(keys); +} + +tl::expected PyClient::register_buffer_internal(void *buffer, + size_t size) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + return client_->RegisterLocalMemory(buffer, size, kWildcardLocation, false, + true); +} + +int PyClient::register_buffer(void *buffer, size_t size) { + return to_py_ret(register_buffer_internal(buffer, size)); +} + +tl::expected PyClient::unregister_buffer_internal( + void *buffer) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto unregister_result = client_->unregisterLocalMemory(buffer, true); + if (!unregister_result) { + LOG(ERROR) << "Unregister buffer failed with error: " + << toString(unregister_result.error()); + return tl::unexpected(unregister_result.error()); + } + return {}; +} + +int PyClient::unregister_buffer(void *buffer) { + return to_py_ret(unregister_buffer_internal(buffer)); +} + +tl::expected PyClient::get_into_internal( + const std::string &key, void *buffer, size_t size) { + // NOTE: The buffer address must be previously registered with + // register_buffer() for zero-copy RDMA operations to work correctly + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + // Step 1: Get object info + auto query_result = client_->Query(key); + if (!query_result) { + if (query_result.error() == ErrorCode::OBJECT_NOT_FOUND) { + VLOG(1) << "Object not found for key: " << key; + return tl::unexpected(query_result.error()); + } + LOG(ERROR) << "Query failed for key: " << key + << " with error: " << toString(query_result.error()); + return tl::unexpected(query_result.error()); + } + + auto replica_list = query_result.value(); + + // Calculate total size from replica list + if (replica_list.empty()) { + LOG(ERROR) << "Internal error: replica_list is empty"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + auto &replica = replica_list[0]; + uint64_t total_size = calculate_total_size(replica); + + // Check if user buffer is large enough + if (size < total_size) { + LOG(ERROR) << "User buffer too small. Required: " << total_size + << ", provided: " << size; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + // Step 2: Split user buffer according to object info and create + // slices + std::vector slices; + uint64_t offset = 0; + + if (replica.is_memory_replica() == false) { + while (offset < total_size) { + auto chunk_size = std::min(total_size - offset, kMaxSliceSize); + void *chunk_ptr = static_cast(buffer) + offset; + slices.emplace_back(Slice{chunk_ptr, chunk_size}); + offset += chunk_size; + } + } else { + for (auto &handle : + replica.get_memory_descriptor().buffer_descriptors) { + void *chunk_ptr = static_cast(buffer) + offset; + slices.emplace_back(Slice{chunk_ptr, handle.size_}); + offset += handle.size_; + } + } + + // Step 3: Read data directly into user buffer + auto get_result = client_->Get(key, replica_list, slices); + if (!get_result) { + LOG(ERROR) << "Get failed for key: " << key + << " with error: " << toString(get_result.error()); + return tl::unexpected(get_result.error()); + } + + return static_cast(total_size); +} + +int PyClient::get_into(const std::string &key, void *buffer, size_t size) { + return to_py_ret(get_into_internal(key, buffer, size)); +} + +std::string PyClient::get_hostname() const { return local_hostname; } + +std::vector PyClient::batch_put_from(const std::vector &keys, + const std::vector &buffers, + const std::vector &sizes, + const ReplicateConfig &config) { + auto internal_results = + batch_put_from_internal(keys, buffers, sizes, config); + std::vector results; + results.reserve(internal_results.size()); + + for (const auto &result : internal_results) { + results.push_back(to_py_ret(result)); + } + + return results; +} + +std::vector PyClient::batch_get_into(const std::vector &keys, + const std::vector &buffers, + const std::vector &sizes) { + auto internal_results = batch_get_into_internal(keys, buffers, sizes); + std::vector results; + results.reserve(internal_results.size()); + + for (const auto &result : internal_results) { + results.push_back(to_py_ret(result)); + } + + return results; +} + +tl::expected PyClient::put_from_internal( + const std::string &key, void *buffer, size_t size, + const ReplicateConfig &config) { + // NOTE: The buffer address must be previously registered with + // register_buffer() for zero-copy RDMA operations to work correctly + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + if (size == 0) { + LOG(WARNING) << "Attempting to put empty data for key: " << key; + return {}; + } + + // Create slices directly from the user buffer + std::vector slices; + uint64_t offset = 0; + + while (offset < size) { + auto chunk_size = std::min(size - offset, kMaxSliceSize); + void *chunk_ptr = static_cast(buffer) + offset; + slices.emplace_back(Slice{chunk_ptr, chunk_size}); + offset += chunk_size; + } + + auto put_result = client_->Put(key, slices, config); + if (!put_result) { + LOG(ERROR) << "Put operation failed with error: " + << toString(put_result.error()); + return tl::unexpected(put_result.error()); + } + + return {}; +} + +int PyClient::put_from(const std::string &key, void *buffer, size_t size, + const ReplicateConfig &config) { + return to_py_ret(put_from_internal(key, buffer, size, config)); +} + +std::vector> PyClient::batch_get_into_internal( + const std::vector &keys, const std::vector &buffers, + const std::vector &sizes) { + // Validate preconditions + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + + if (keys.size() != buffers.size() || keys.size() != sizes.size()) { + LOG(ERROR) << "Input vector sizes mismatch: keys=" << keys.size() + << ", buffers=" << buffers.size() + << ", sizes=" << sizes.size(); + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + + const size_t num_keys = keys.size(); + std::vector> results; + results.reserve(num_keys); + + if (num_keys == 0) { + return results; + } + + // Query metadata for all keys + const auto query_results = client_->BatchQuery(keys); + + // Process each key individually and prepare for batch transfer + struct ValidKeyInfo { + std::string key; + size_t original_index; + std::vector replica_list; + std::vector slices; + uint64_t total_size; + }; + + std::vector valid_operations; + valid_operations.reserve(num_keys); + + for (size_t i = 0; i < num_keys; ++i) { + const auto &key = keys[i]; + + // Handle query failures + if (!query_results[i]) { + const auto error = query_results[i].error(); + results.emplace_back(tl::unexpected(error)); + if (error != ErrorCode::OBJECT_NOT_FOUND) { + LOG(ERROR) << "Query failed for key '" << key + << "': " << toString(error); + } + continue; + } + + // Validate replica list + auto replica_list = query_results[i].value(); + if (replica_list.empty()) { + LOG(ERROR) << "Empty replica list for key: " << key; + results.emplace_back(tl::unexpected(ErrorCode::INVALID_REPLICA)); + continue; + } + + // Calculate required buffer size + const auto &replica = replica_list[0]; + uint64_t total_size = calculate_total_size(replica); + + // Validate buffer capacity + if (sizes[i] < total_size) { + LOG(ERROR) << "Buffer too small for key '" << key + << "': required=" << total_size + << ", available=" << sizes[i]; + results.emplace_back(tl::unexpected(ErrorCode::INVALID_PARAMS)); + continue; + } + + // Create slices for this key's buffer + std::vector key_slices; + uint64_t offset = 0; + if (replica.is_memory_replica() == false) { + while (offset < total_size) { + auto chunk_size = std::min(total_size - offset, kMaxSliceSize); + void *chunk_ptr = static_cast(buffers[i]) + offset; + key_slices.emplace_back(Slice{chunk_ptr, chunk_size}); + offset += chunk_size; + } + } else { + for (auto &handle : + replica.get_memory_descriptor().buffer_descriptors) { + void *chunk_ptr = static_cast(buffers[i]) + offset; + key_slices.emplace_back(Slice{chunk_ptr, handle.size_}); + offset += handle.size_; + } + } + + // Store operation info for batch processing + valid_operations.push_back({.key = key, + .original_index = i, + .replica_list = std::move(replica_list), + .slices = std::move(key_slices), + .total_size = total_size}); + + // Set success result (actual bytes transferred) + results.emplace_back(static_cast(total_size)); + } + + // Early return if no valid operations + if (valid_operations.empty()) { + return results; + } + + // Prepare batch transfer data structures + std::vector batch_keys; + std::vector> batch_replica_lists; + std::unordered_map> batch_slices; + + batch_keys.reserve(valid_operations.size()); + batch_replica_lists.reserve(valid_operations.size()); + + for (const auto &op : valid_operations) { + batch_keys.push_back(op.key); + batch_replica_lists.push_back(op.replica_list); + batch_slices[op.key] = op.slices; + } + + // Execute batch transfer + const auto batch_get_results = + client_->BatchGet(batch_keys, batch_replica_lists, batch_slices); + + // Process transfer results + for (size_t j = 0; j < batch_get_results.size(); ++j) { + const auto &op = valid_operations[j]; + + if (!batch_get_results[j]) { + const auto error = batch_get_results[j].error(); + LOG(ERROR) << "BatchGet failed for key '" << op.key + << "': " << toString(error); + results[op.original_index] = tl::unexpected(error); + } + } + + return results; +} + +std::vector> PyClient::batch_put_from_internal( + const std::vector &keys, const std::vector &buffers, + const std::vector &sizes, const ReplicateConfig &config) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + + if (keys.size() != buffers.size() || keys.size() != sizes.size()) { + LOG(ERROR) << "Mismatched sizes for keys, buffers, and sizes"; + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + + std::unordered_map> all_slices; + + // Create slices from user buffers + for (size_t i = 0; i < keys.size(); ++i) { + const std::string &key = keys[i]; + void *buffer = buffers[i]; + size_t size = sizes[i]; + + std::vector slices; + uint64_t offset = 0; + + while (offset < size) { + auto chunk_size = std::min(size - offset, kMaxSliceSize); + void *chunk_ptr = static_cast(buffer) + offset; + slices.emplace_back(Slice{chunk_ptr, chunk_size}); + offset += chunk_size; + } + + all_slices[key] = std::move(slices); + } + + std::vector> ordered_batched_slices; + ordered_batched_slices.reserve(keys.size()); + for (const auto &key : keys) { + auto it = all_slices.find(key); + if (it != all_slices.end()) { + ordered_batched_slices.emplace_back(it->second); + } else { + LOG(ERROR) << "Missing slices for key: " << key; + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + } + + // Call client BatchPut and return the vector directly + return client_->BatchPut(keys, ordered_batched_slices, config); +} + +std::vector> PyClient::batchIsExist_internal( + const std::vector &keys) { + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + + if (keys.empty()) { + LOG(WARNING) << "Empty keys vector provided to batchIsExist_internal"; + return std::vector>(); + } + + // Call client BatchIsExist and return the vector directly + return client_->BatchIsExist(keys); +} + +int PyClient::put_from_with_metadata(const std::string &key, void *buffer, + void *metadata_buffer, size_t size, + size_t metadata_size, + const ReplicateConfig &config) { + // NOTE: The buffer address must be previously registered with + // register_buffer() for zero-copy RDMA operations to work correctly + if (!client_) { + LOG(ERROR) << "Client is not initialized"; + return -1; + } + + if (size == 0) { + LOG(WARNING) << "Attempting to put empty data for key: " << key; + return 0; + } + + // Create slices directly from the user buffer + std::vector slices; + // Add metadata slice + uint64_t metadata_offset = 0; + while (metadata_offset < metadata_size) { + auto metadata_chunk_size = + std::min(metadata_size - metadata_offset, kMaxSliceSize); + void *metadata_chunk_ptr = + static_cast(metadata_buffer) + metadata_offset; + slices.emplace_back(Slice{metadata_chunk_ptr, metadata_chunk_size}); + metadata_offset += metadata_chunk_size; + } + + uint64_t offset = 0; + while (offset < size) { + auto chunk_size = std::min(size - offset, kMaxSliceSize); + void *chunk_ptr = static_cast(buffer) + offset; + slices.emplace_back(Slice{chunk_ptr, chunk_size}); + offset += chunk_size; + } + auto put_result = client_->Put(key, slices, config); + if (!put_result) { + LOG(ERROR) << "Put operation failed with error: " + << toString(put_result.error()); + return -toInt(put_result.error()); + } + return 0; +} + +} // namespace mooncake diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index d9fbb7c6..9be97d78 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -2,6 +2,11 @@ #include #include +#include +#include +#include + +#include namespace mooncake { void *allocate_buffer_allocator_memory(size_t total_size) { @@ -46,4 +51,41 @@ void **rdma_args(const std::string &device_name) { return args; } +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; +} + +int getRandomAvailablePort(int min_port, int max_port) { + // Handle invalid range + if (min_port > max_port) { + std::swap(min_port, max_port); + } + + 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); + if (isPortAvailable(port)) { + return port; + } + } + return -1; // Failed to find available port +} + } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 43cdd625..f85b5b28 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -124,4 +124,16 @@ target_link_libraries(client_buffer_test PUBLIC ) add_test(NAME client_buffer_test COMMAND client_buffer_test) +add_executable(pybind_client_test pybind_client_test.cpp) +target_link_libraries(pybind_client_test PUBLIC + mooncake_store + cachelib_memory_allocator + ${ETCD_WRAPPER_LIB} + glog + gtest + gtest_main + pthread +) +add_test(NAME pybind_client_test COMMAND pybind_client_test) + add_subdirectory(e2e) diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp new file mode 100644 index 00000000..598a558b --- /dev/null +++ b/mooncake-store/tests/pybind_client_test.cpp @@ -0,0 +1,115 @@ +#include +#include +#include + +#include +#include + +#include "pybind_client.h" + +DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); +DEFINE_string(device_name, "ibp6s0", + "Device name to use, valid if protocol=rdma"); +DEFINE_string(transfer_engine_metadata_url, "http://localhost:8080/metadata", + "Metadata connection string for transfer engine"); + +namespace mooncake { +namespace testing { + +class PyClientTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + google::InitGoogleLogging("PyClientTest"); + FLAGS_logtostderr = 1; + + // Override flags from environment variables if present + if (getenv("PROTOCOL")) FLAGS_protocol = getenv("PROTOCOL"); + if (getenv("DEVICE_NAME")) FLAGS_device_name = getenv("DEVICE_NAME"); + if (getenv("MC_METADATA_SERVER")) + FLAGS_transfer_engine_metadata_url = getenv("MC_METADATA_SERVER"); + + LOG(INFO) << "Protocol: " << FLAGS_protocol + << ", Device name: " << FLAGS_device_name + << ", Metadata URL: " << FLAGS_transfer_engine_metadata_url; + } + + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } + + void SetUp() override { py_client_ = std::make_unique(); } + + void TearDown() override { + if (py_client_) { + py_client_->tearDownAll(); + } + } + + std::unique_ptr py_client_; +}; + +// Test PyClient construction and setup +TEST_F(PyClientTest, ConstructorAndSetup) { + ASSERT_TRUE(py_client_ != nullptr); + + // Test setup + int setup_result = py_client_->setup( + "localhost:17813", // local_hostname + FLAGS_transfer_engine_metadata_url, // metadata_server + 16 * 1024 * 1024, // global_segment_size (16MB) + 16 * 1024 * 1024, // local_buffer_size (16MB) + FLAGS_protocol, // protocol + FLAGS_device_name, // rdma_devices + "localhost:50051" // master_server_addr + ); + EXPECT_EQ(setup_result, 0) << "Setup should succeed"; + + // Verify hostname + std::string hostname = py_client_->get_hostname(); + EXPECT_EQ(hostname, "localhost:17813"); +} + +// Test basic Put and Get operations +TEST_F(PyClientTest, BasicPutGetOperations) { + // Setup the client + ASSERT_EQ( + py_client_->setup("localhost:17813", FLAGS_transfer_engine_metadata_url, + 16 * 1024 * 1024, 16 * 1024 * 1024, FLAGS_protocol, + FLAGS_device_name, "localhost:50051"), + 0); + + const std::string test_data = "Hello, PyClient!"; + const std::string key = "test_key_pyclient"; + + // Test Put operation using span + std::span data_span(test_data.data(), test_data.size()); + ReplicateConfig config; + config.replica_num = 1; + + int put_result = py_client_->put(key, data_span, config); + EXPECT_EQ(put_result, 0) << "Put operation should succeed"; + + // Test Get operation using buffer handle + auto buffer_handle = py_client_->get_buffer(key); + ASSERT_TRUE(buffer_handle != nullptr) << "Get buffer should succeed"; + EXPECT_EQ(buffer_handle->size(), test_data.size()) + << "Buffer size should match"; + + // Verify the data + std::string retrieved_data(static_cast(buffer_handle->ptr()), + buffer_handle->size()); + EXPECT_EQ(retrieved_data, test_data) + << "Retrieved data should match original"; + + // Test isExist + int exist_result = py_client_->isExist(key); + EXPECT_EQ(exist_result, 1) << "Key should exist"; +} + +} // namespace testing + +} // namespace mooncake + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, false); + return RUN_ALL_TESTS(); +} \ No newline at end of file diff --git a/mooncake-store/tests/utils_test.cpp b/mooncake-store/tests/utils_test.cpp index ad55b6ce..259bef9e 100644 --- a/mooncake-store/tests/utils_test.cpp +++ b/mooncake-store/tests/utils_test.cpp @@ -1,6 +1,10 @@ #include "utils.h" #include +#include +#include +#include +#include using namespace mooncake; @@ -18,3 +22,121 @@ TEST(UtilsTest, ByteSizeToString) { EXPECT_EQ(byte_size_to_string(15 * 1024 + 134), "15.13 KB"); EXPECT_EQ(byte_size_to_string(15 * 1024 * 1024 + 44048), "15.04 MB"); } + +TEST(UtilsTest, IsPortAvailable) { + // Find an available port + int test_port = -1; + for (int port = 50000; port < 50010; ++port) { + if (isPortAvailable(port)) { + test_port = port; + break; + } + } + ASSERT_NE(test_port, -1) << "Could not find available port for testing"; + + // Initially the port should be available + EXPECT_TRUE(isPortAvailable(test_port)); + + // Bind to the port to make it unavailable + int sock = socket(AF_INET, SOCK_STREAM, 0); + ASSERT_NE(sock, -1); + + 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(test_port); + + ASSERT_EQ(bind(sock, (struct sockaddr*)&addr, sizeof(addr)), 0); + ASSERT_EQ(listen(sock, 1), 0); + + // Now the port should be unavailable + EXPECT_FALSE(isPortAvailable(test_port)); + + // Clean up + close(sock); + + // Port should be available again + EXPECT_TRUE(isPortAvailable(test_port)); +} + +TEST(UtilsTest, IsPortAvailableWithBindingConflict) { + // Find an available port + int test_port = -1; + for (int port = 40000; port < 40010; ++port) { + if (isPortAvailable(port)) { + test_port = port; + break; + } + } + + ASSERT_NE(test_port, -1) << "Could not find available port for testing"; + + // Initially the port should be available + EXPECT_TRUE(isPortAvailable(test_port)); + + // Bind to the port to make it unavailable + int sock = socket(AF_INET, SOCK_STREAM, 0); + ASSERT_NE(sock, -1); + + 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(test_port); + + ASSERT_EQ(bind(sock, (struct sockaddr*)&addr, sizeof(addr)), 0); + ASSERT_EQ(listen(sock, 1), 0); + + // Now the port should be unavailable + EXPECT_FALSE(isPortAvailable(test_port)); + + // Clean up + close(sock); + + // Port should be available again + EXPECT_TRUE(isPortAvailable(test_port)); +} + +TEST(UtilsTest, GetRandomAvailablePort) { + // Test with default parameters + int port = getRandomAvailablePort(); + EXPECT_GE(port, 12300); + EXPECT_LE(port, 14300); + EXPECT_TRUE(isPortAvailable(port)); +} + +TEST(UtilsTest, GetRandomAvailablePortCustomRange) { + // Test with custom range + int port = getRandomAvailablePort(20000, 20010); + if (port != -1) { // If we found a port + EXPECT_GE(port, 20000); + EXPECT_LE(port, 20010); + EXPECT_TRUE(isPortAvailable(port)); + } +} + +TEST(UtilsTest, GetRandomAvailablePortInvalidRange) { + // Test with invalid range (min > max) + int port = getRandomAvailablePort(14300, 12300); + // Should still work by swapping or handling gracefully + EXPECT_TRUE(port == -1 || (port >= 12300 && port <= 14300)); +} + +TEST(UtilsTest, GetRandomAvailablePortReturnsValidPorts) { + // Test that multiple calls return valid, potentially different ports + std::set ports_found; + for (int i = 0; i < 5; ++i) { + int port = getRandomAvailablePort(30000, 30100); + if (port != -1) { + EXPECT_TRUE(isPortAvailable(port)); + ports_found.insert(port); + } + } + + // We should find at least some available ports in this range + EXPECT_GT(ports_found.size(), 0u); +}