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

refactor(store): move python bidning to `pybind_client` (#723) 

* refactor(store): replace SliceBuffer with BufferHandle in Python bindings

* refactor(store_py): introduce PyClient and wrapper class for Python bindings

* fix

Signed-off-by: Jinyang Su <751080330@qq.com>

---------

Signed-off-by: Jinyang Su <751080330@qq.com>
This commit is contained in:
JinYan Su 2025-08-11 14:09:42 +08:00 committed by GitHub
parent 723276feaa
commit d2fa93dfb2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 1735 additions and 1595 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
docs/source/mooncake-store-api/python-binding.md
View File

@ -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:
<details>
<summary>Click to expand: Buffer protocol usage example</summary>
@ -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

1858
mooncake-integration/store/store_py.cpp
View File

File diff suppressed because it is too large Load Diff

75
mooncake-integration/store/store_py.h → mooncake-store/include/pybind_client.h
View File

@ -1,8 +1,5 @@
#pragma once
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <csignal>
#include <mutex>
#include <string>
@ -13,10 +10,7 @@
namespace mooncake {
class DistributedObjectStore;
// Forward declarations
class SliceBuffer;
class PyClient;
template <class T>
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<DistributedObjectStore *> instances_;
std::unordered_set<PyClient *> 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<pybind11::bytes> get_batch(
const std::vector<std::string> &keys);
/**
* @brief Get a buffer containing the data for a key
* @param key Key to get data for
* @return std::shared_ptr<SliceBuffer> Buffer containing the data, or
* @return std::shared_ptr<BufferHandle> Buffer containing the data, or
* nullptr if error
*/
std::shared_ptr<SliceBuffer> get_buffer(const std::string &key);
std::shared_ptr<BufferHandle> 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<SliceBuffer> buffers containing the
* @return Vector of std::shared_ptr<BufferHandle> buffers containing the
* data, or nullptr for each key if error
*/
std::vector<std::shared_ptr<SliceBuffer>> batch_get_buffer(
std::vector<std::shared_ptr<BufferHandle>> batch_get_buffer(
const std::vector<std::string> &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<void, ErrorCode> setup_internal(
const std::string &local_hostname, const std::string &metadata_server,
@ -311,11 +270,6 @@ class DistributedObjectStore {
const std::vector<std::span<const char>> &values,
const ReplicateConfig &config = ReplicateConfig{});
tl::expected<void, ErrorCode> put_batch_internal(
const std::vector<std::string> &keys,
const std::vector<pybind11::buffer> &buffers,
const ReplicateConfig &config = ReplicateConfig{});
tl::expected<void, ErrorCode> remove_internal(const std::string &key);
tl::expected<int64_t, ErrorCode> removeAll_internal();
@ -329,10 +283,7 @@ class DistributedObjectStore {
tl::expected<int64_t, ErrorCode> getSize_internal(const std::string &key);
tl::expected<void, ErrorCode> put_tensor_internal(const std::string &key,
pybind11::object tensor);
std::vector<std::shared_ptr<SliceBuffer>> batch_get_buffer_internal(
std::vector<std::shared_ptr<BufferHandle>> batch_get_buffer_internal(
const std::vector<std::string> &keys);
std::shared_ptr<mooncake::Client> client_ = nullptr;
@ -352,4 +303,4 @@ class DistributedObjectStore {
std::string local_hostname;
};
} // namespace mooncake
} // namespace mooncake

19
mooncake-store/include/utils.h
View File

@ -111,4 +111,21 @@ void** rdma_args(const std::string& device_name);
return oss.str();
}
} // namespace mooncake
// 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

1
mooncake-store/src/CMakeLists.txt
View File

@ -21,6 +21,7 @@ set(MOONCAKE_STORE_SOURCES
offset_allocator.cpp
posix_file.cpp
client_buffer.cpp
pybind_client.cpp
)
set(EXTRA_LIBS "")

1080
mooncake-store/src/pybind_client.cpp Normal file
View File

File diff suppressed because it is too large Load Diff

42
mooncake-store/src/utils.cpp
View File

@ -2,6 +2,11 @@
#include <Slab.h>
#include <glog/logging.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
#include <random>
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

12
mooncake-store/tests/CMakeLists.txt
View File

@ -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)

115
mooncake-store/tests/pybind_client_test.cpp Normal file
View File

@ -0,0 +1,115 @@
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <memory>
#include <string>
#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<PyClient>(); }
void TearDown() override {
if (py_client_) {
py_client_->tearDownAll();
}
}
std::unique_ptr<PyClient> 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<const char> 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<const char*>(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();
}

122
mooncake-store/tests/utils_test.cpp
View File

@ -1,6 +1,10 @@
#include "utils.h"
#include <gtest/gtest.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
#include <set>
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<int> 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);
}