Skip to content

Commit 1fc3688

Browse files
YiXRYiXR
committed
[Store]: Change dummy client setup into a new func
Signed-off-by: Xingrui Yi <[email protected]>
1 parent f73b712 commit 1fc3688

File tree

8 files changed

+86
-76
lines changed

8 files changed

+86
-76
lines changed

‎mooncake-integration/store/store_py.cpp‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -580,31 +580,38 @@ PYBIND11_MODULE(store, m) {
580580
const std::string &protocol = "tcp",
581581
const std::string &rdma_devices = "",
582582
const std::string &master_server_addr = "127.0.0.1:50051",
583-
const py::object &engine = py::none(),
584-
bool use_dummy_client = false) {
585-
self.use_dummy_client_ = use_dummy_client;
586-
if (use_dummy_client) {
587-
self.store_ = std::make_shared<DummyClient>();
588-
} else {
589-
self.store_ = std::make_shared<RealClient>();
590-
}
583+
const py::object &engine = py::none()) {
584+
self.use_dummy_client_ = false;
585+
self.store_ = std::make_shared<RealClient>();
591586
ResourceTracker::getInstance().registerInstance(
592587
std::dynamic_pointer_cast<PyClient>(self.store_));
593588
std::shared_ptr<TransferEngine> transfer_engine = nullptr;
594589
if (!engine.is_none()) {
595590
transfer_engine =
596591
engine.cast<std::shared_ptr<TransferEngine>>();
597592
}
598-
return self.store_->setup(
593+
return self.store_->setup_real(
599594
local_hostname, metadata_server, global_segment_size,
600595
local_buffer_size, protocol, rdma_devices,
601596
master_server_addr, transfer_engine);
602597
},
603598
py::arg("local_hostname"), py::arg("metadata_server"),
604599
py::arg("global_segment_size"), py::arg("local_buffer_size"),
605600
py::arg("protocol"), py::arg("rdma_devices"),
606-
py::arg("master_server_addr"), py::arg("engine") = py::none(),
607-
py::arg("use_dummy_client") = false)
601+
py::arg("master_server_addr"), py::arg("engine") = py::none())
602+
.def(
603+
"setup_dummy",
604+
[](MooncakeStorePyWrapper &self, size_t mem_pool_size,
605+
size_t local_buffer_size, const std::string &server_address) {
606+
self.use_dummy_client_ = true;
607+
self.store_ = std::make_shared<DummyClient>();
608+
ResourceTracker::getInstance().registerInstance(
609+
std::dynamic_pointer_cast<PyClient>(self.store_));
610+
return self.store_->setup_dummy(
611+
mem_pool_size, local_buffer_size, server_address);
612+
},
613+
py::arg("mem_pool_size"), py::arg("local_buffer_size"),
614+
py::arg("server_address"))
608615
.def("init_all",
609616
[](MooncakeStorePyWrapper &self, const std::string &protocol,
610617
const std::string &device_name,

‎mooncake-store/include/dummy_client.h‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,18 @@ class DummyClient : public PyClient {
1717

1818
int64_t unregister_shm();
1919

20-
int setup(const std::string &local_hostname,
21-
const std::string &metadata_server, size_t global_segment_size,
22-
size_t local_buffer_size, const std::string &protocol,
23-
const std::string &rdma_devices,
24-
const std::string &master_server_addr,
25-
const std::shared_ptr<TransferEngine> &transfer_engine);
20+
int setup_real(const std::string &local_hostname,
21+
const std::string &metadata_server,
22+
size_t global_segment_size, size_t local_buffer_size,
23+
const std::string &protocol, const std::string &rdma_devices,
24+
const std::string &master_server_addr,
25+
const std::shared_ptr<TransferEngine> &transfer_engine) {
26+
// Dummy client does not support real setup
27+
return -1;
28+
};
29+
30+
int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
31+
const std::string &server_address);
2632

2733
int initAll(const std::string &protocol, const std::string &device_name,
2834
size_t mount_segment_size);

‎mooncake-store/include/pyclient.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,16 @@ namespace mooncake {
1818
class PyClient {
1919
public:
2020
virtual ~PyClient() = 0;
21-
virtual int setup(
21+
virtual int setup_real(
2222
const std::string &local_hostname, const std::string &metadata_server,
2323
size_t global_segment_size, size_t local_buffer_size,
2424
const std::string &protocol, const std::string &rdma_devices,
2525
const std::string &master_server_addr,
2626
const std::shared_ptr<TransferEngine> &transfer_engine) = 0;
2727

28+
virtual int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
29+
const std::string &server_address) = 0;
30+
2831
virtual int initAll(const std::string &protocol,
2932
const std::string &device_name,
3033
size_t mount_segment_size) = 0;

‎mooncake-store/include/real_client.h‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -66,14 +66,20 @@ class RealClient : public PyClient {
6666
// Factory to create shared instances and auto-register to ResourceTracker
6767
static std::shared_ptr<RealClient> create();
6868

69-
int setup(const std::string &local_hostname,
70-
const std::string &metadata_server,
71-
size_t global_segment_size = 1024 * 1024 * 16,
72-
size_t local_buffer_size = 1024 * 1024 * 16,
73-
const std::string &protocol = "tcp",
74-
const std::string &rdma_devices = "",
75-
const std::string &master_server_addr = "127.0.0.1:50051",
76-
const std::shared_ptr<TransferEngine> &transfer_engine = nullptr);
69+
int setup_real(
70+
const std::string &local_hostname, const std::string &metadata_server,
71+
size_t global_segment_size = 1024 * 1024 * 16,
72+
size_t local_buffer_size = 1024 * 1024 * 16,
73+
const std::string &protocol = "tcp",
74+
const std::string &rdma_devices = "",
75+
const std::string &master_server_addr = "127.0.0.1:50051",
76+
const std::shared_ptr<TransferEngine> &transfer_engine = nullptr);
77+
78+
int setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
79+
const std::string &server_address) {
80+
// Real client does not support dummy setup
81+
return -1;
82+
};
7783

7884
int initAll(const std::string &protocol, const std::string &device_name,
7985
size_t mount_segment_size = 1024 * 1024 * 16); // Default 16MB

‎mooncake-store/src/dummy_client.cpp‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -115,13 +115,8 @@ ErrorCode DummyClient::connect(const std::string& server_address) {
115115
return ErrorCode::OK;
116116
}
117117

118-
int DummyClient::setup(const std::string& local_hostname,
119-
const std::string& metadata_server,
120-
size_t global_segment_size, size_t local_buffer_size,
121-
const std::string& protocol,
122-
const std::string& rdma_devices,
123-
const std::string& server_address,
124-
const std::shared_ptr<TransferEngine>& transfer_engine) {
118+
int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size,
119+
const std::string& server_address) {
125120
int64_t ret;
126121
ErrorCode err = connect(server_address);
127122
if (err != ErrorCode::OK) {
@@ -131,7 +126,7 @@ int DummyClient::setup(const std::string& local_hostname,
131126

132127
shm_name_ = "/dummy_client_shm_" + std::to_string(client_id_.first) + "_" +
133128
std::to_string(client_id_.second);
134-
shm_size_ = local_buffer_size + global_segment_size;
129+
shm_size_ = local_buffer_size + mem_pool_size;
135130
local_buffer_size_ = local_buffer_size;
136131

137132
// Open or create shared memory object
@@ -195,8 +190,8 @@ int DummyClient::initAll(const std::string& protocol,
195190
const std::string& device_name,
196191
size_t mount_segment_size) {
197192
uint64_t buffer_allocator_size = 1024 * 1024 * 1024;
198-
return setup("", "", 0, buffer_allocator_size, protocol, device_name,
199-
"127.0.0.1:50052", nullptr);
193+
return setup_real("", "", 0, buffer_allocator_size, protocol, device_name,
194+
"127.0.0.1:50052", nullptr);
200195
}
201196

202197
int DummyClient::tearDownAll() {

‎mooncake-store/src/real_client.cpp‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -248,13 +248,12 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
248248
return {};
249249
}
250250

251-
int RealClient::setup(const std::string &local_hostname,
252-
const std::string &metadata_server,
253-
size_t global_segment_size, size_t local_buffer_size,
254-
const std::string &protocol,
255-
const std::string &rdma_devices,
256-
const std::string &master_server_addr,
257-
const std::shared_ptr<TransferEngine> &transfer_engine) {
251+
int RealClient::setup_real(
252+
const std::string &local_hostname, const std::string &metadata_server,
253+
size_t global_segment_size, size_t local_buffer_size,
254+
const std::string &protocol, const std::string &rdma_devices,
255+
const std::string &master_server_addr,
256+
const std::shared_ptr<TransferEngine> &transfer_engine) {
258257
return to_py_ret(setup_internal(
259258
local_hostname, metadata_server, global_segment_size, local_buffer_size,
260259
protocol, rdma_devices, master_server_addr, transfer_engine));

‎mooncake-store/tests/pybind_client_test.cpp‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -77,10 +77,11 @@ TEST_F(RealClientTest, BasicPutGetOperations) {
7777
const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
7878
? FLAGS_device_name
7979
: std::string("");
80-
ASSERT_EQ(py_client_->setup("localhost:17813", "P2PHANDSHAKE",
81-
16 * 1024 * 1024, 16 * 1024 * 1024,
82-
FLAGS_protocol, rdma_devices, master_address_),
83-
0);
80+
ASSERT_EQ(
81+
py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
82+
16 * 1024 * 1024, 16 * 1024 * 1024,
83+
FLAGS_protocol, rdma_devices, master_address_),
84+
0);
8485

8586
const std::string test_data = "Hello, RealClient!";
8687
const std::string key = "test_key_realclient";
@@ -126,10 +127,11 @@ TEST_F(RealClientTest, GetWithLeaseTimeOut) {
126127
const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
127128
? FLAGS_device_name
128129
: std::string("");
129-
ASSERT_EQ(py_client_->setup("localhost:17813", "P2PHANDSHAKE",
130-
512 * 1024 * 1024, 512 * 1024 * 1024,
131-
FLAGS_protocol, rdma_devices, master_address_),
132-
0);
130+
ASSERT_EQ(
131+
py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
132+
512 * 1024 * 1024, 512 * 1024 * 1024,
133+
FLAGS_protocol, rdma_devices, master_address_),
134+
0);
133135

134136
const size_t data_size = 256 * 1024 * 1024; // 256MB
135137
std::string test_data(data_size, 'A'); // Fill with 'A' characters
@@ -257,9 +259,9 @@ TEST_F(RealClientTest, ConcurrentPutGetWithLeaseTimeOut) {
257259
const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
258260
? FLAGS_device_name
259261
: std::string("");
260-
ASSERT_EQ(py_client_->setup("localhost:17813", "P2PHANDSHAKE", segment_size,
261-
segment_size, FLAGS_protocol, rdma_devices,
262-
master_address_),
262+
ASSERT_EQ(py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
263+
segment_size, segment_size, FLAGS_protocol,
264+
rdma_devices, master_address_),
263265
0);
264266

265267
// Test Single Get Operation with Concurrent Put
@@ -512,11 +514,11 @@ TEST_F(RealClientTest, TestSetupExistTransferEngine) {
512514
} else {
513515
ASSERT_TRUE(false) << "Unsupported protocol: " << FLAGS_protocol;
514516
}
515-
ASSERT_EQ(
516-
py_client_->setup("localhost:17813", "P2PHANDSHAKE", 16 * 1024 * 1024,
517-
16 * 1024 * 1024, FLAGS_protocol, rdma_devices,
518-
master_address_, transfer_engine),
519-
0);
517+
ASSERT_EQ(py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
518+
16 * 1024 * 1024, 16 * 1024 * 1024,
519+
FLAGS_protocol, rdma_devices,
520+
master_address_, transfer_engine),
521+
0);
520522

521523
const std::string test_data = "Hello, RealClient!";
522524
const std::string key = "test_key_realclient";
@@ -541,10 +543,11 @@ TEST_F(RealClientTest, TestBatchPutAndGetMultiBuffers) {
541543
const std::string rdma_devices = (FLAGS_protocol == std::string("rdma"))
542544
? FLAGS_device_name
543545
: std::string("");
544-
ASSERT_EQ(py_client_->setup("localhost:17813", "P2PHANDSHAKE",
545-
16 * 1024 * 1024, 16 * 1024 * 1024,
546-
FLAGS_protocol, rdma_devices, master_address_),
547-
0);
546+
ASSERT_EQ(
547+
py_client_->setup_real("localhost:17813", "P2PHANDSHAKE",
548+
16 * 1024 * 1024, 16 * 1024 * 1024,
549+
FLAGS_protocol, rdma_devices, master_address_),
550+
0);
548551

549552
std::string test_data(1000, '1');
550553
std::string dst_data(1000, '0');

‎mooncake-wheel/tests/test_dummy_client.py‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,26 +14,17 @@
1414

1515
def get_client(store, local_buffer_size_param=None):
1616
"""Initialize and setup the distributed store client."""
17-
protocol = os.getenv("PROTOCOL", "tcp")
18-
device_name = os.getenv("DEVICE_NAME", "ibp6s0")
19-
local_hostname = os.getenv("LOCAL_HOSTNAME", "localhost")
20-
metadata_server = os.getenv("MC_METADATA_SERVER", "http://127.0.0.1:8080/metadata")
21-
global_segment_size = 3200 * 1024 * 1024 # 3200 MB
17+
mem_pool_size = 3200 * 1024 * 1024 # 3200 MB
2218
local_buffer_size = (
2319
local_buffer_size_param if local_buffer_size_param is not None
2420
else 512 * 1024 * 1024 # 512 MB
2521
)
2622
real_client_address = "127.0.0.1:50052"
2723

28-
retcode = store.setup(
29-
local_hostname,
30-
metadata_server,
31-
global_segment_size,
24+
retcode = store.setup_dummy(
25+
mem_pool_size,
3226
local_buffer_size,
33-
protocol,
34-
device_name,
35-
real_client_address,
36-
use_dummy_client=True
27+
real_client_address
3728
)
3829

3930
if retcode:

0 commit comments

Comments
 (0)