Skip to content

Commit 142696e

Browse files
nickyc975nickyc975
committed
[Store][bind]: Add setup_with_files to setup file-based client
Signed-off-by: Jinlong Chen <[email protected]>
1 parent a116150 commit 142696e

File tree

3 files changed

+101
-5
lines changed

3 files changed

+101
-5
lines changed

‎mooncake-integration/store/store_py.cpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -343,6 +343,18 @@ PYBIND11_MODULE(store, m) {
343343
local_buffer_size, protocol,
344344
rdma_devices, master_server_addr);
345345
})
346+
.def("setup_with_files",
347+
[](MooncakeStorePyWrapper &self, const std::string &local_hostname,
348+
const std::string &metadata_server,
349+
const std::vector<std::string> &files,
350+
size_t local_buffer_size = 1024 * 1024 * 16,
351+
const std::string &protocol = "nvmeof_generic",
352+
const std::string &protocol_arg = "",
353+
const std::string &master_server_addr = "127.0.0.1:50051") {
354+
return self.store_.setup_with_files(
355+
local_hostname, metadata_server, files, local_buffer_size,
356+
protocol, protocol_arg, master_server_addr);
357+
})
346358
.def("init_all",
347359
[](MooncakeStorePyWrapper &self, const std::string &protocol,
348360
const std::string &device_name,

‎mooncake-store/include/pybind_client.h‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,14 @@ class PyClient {
7979
const std::string &rdma_devices = "",
8080
const std::string &master_server_addr = "127.0.0.1:50051");
8181

82+
int setup_with_files(
83+
const std::string &local_hostname, const std::string &metadata_server,
84+
const std::vector<std::string> &files,
85+
size_t local_buffer_size = 1024 * 1024 * 16,
86+
const std::string &protocol = "nvmeof_generic",
87+
const std::string &protocol_arg = "",
88+
const std::string &master_server_addr = "127.0.0.1:50051");
89+
8290
int initAll(const std::string &protocol, const std::string &device_name,
8391
size_t mount_segment_size = 1024 * 1024 * 16); // Default 16MB
8492

@@ -226,6 +234,13 @@ class PyClient {
226234
int64_t getSize(const std::string &key);
227235

228236
// Internal versions that return tl::expected
237+
238+
tl::expected<void, ErrorCode> common_setup_internal(
239+
const std::string &local_hostname, const std::string &metadata_server,
240+
size_t local_buffer_size, const std::string &protocol,
241+
const std::string &protocol_args,
242+
const std::string &master_server_addr);
243+
229244
tl::expected<void, ErrorCode> setup_internal(
230245
const std::string &local_hostname, const std::string &metadata_server,
231246
size_t global_segment_size = 1024 * 1024 * 16,
@@ -234,6 +249,14 @@ class PyClient {
234249
const std::string &rdma_devices = "",
235250
const std::string &master_server_addr = "127.0.0.1:50051");
236251

252+
tl::expected<void, ErrorCode> setup_with_files_internal(
253+
const std::string &local_hostname, const std::string &metadata_server,
254+
const std::vector<std::string> &files,
255+
size_t local_buffer_size = 1024 * 1024 * 16,
256+
const std::string &protocol = "nvmeof_generic",
257+
const std::string &protocol_arg = "",
258+
const std::string &master_server_addr = "127.0.0.1:50051");
259+
237260
tl::expected<void, ErrorCode> initAll_internal(
238261
const std::string &protocol, const std::string &device_name,
239262
size_t mount_segment_size = 1024 * 1024 * 16);

‎mooncake-store/src/pybind_client.cpp‎

Lines changed: 66 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -89,11 +89,10 @@ PyClient::~PyClient() {
8989
ResourceTracker::getInstance().unregisterInstance(this);
9090
}
9191

92-
tl::expected<void, ErrorCode> PyClient::setup_internal(
92+
tl::expected<void, ErrorCode> PyClient::common_setup_internal(
9393
const std::string &local_hostname, const std::string &metadata_server,
94-
size_t global_segment_size, size_t local_buffer_size,
95-
const std::string &protocol, const std::string &rdma_devices,
96-
const std::string &master_server_addr) {
94+
size_t local_buffer_size, const std::string &protocol,
95+
const std::string &protocol_args, const std::string &master_server_addr) {
9796
this->protocol = protocol;
9897

9998
// Remove port if hostname already contains one
@@ -112,7 +111,13 @@ tl::expected<void, ErrorCode> PyClient::setup_internal(
112111
this->local_hostname = local_hostname;
113112
}
114113

115-
void **args = (protocol == "rdma") ? rdma_args(rdma_devices) : nullptr;
114+
void **args = nullptr;
115+
if (protocol == "rdma") {
116+
args = rdma_args(protocol_args);
117+
} else if (protocol == "nvmeof_generic" && !protocol_args.empty()) {
118+
args = (void **)calloc(2, sizeof(void *));
119+
args[0] = (void *)protocol_args.c_str();
120+
}
116121
auto client_opt =
117122
mooncake::Client::Create(this->local_hostname, metadata_server,
118123
protocol, args, master_server_addr);
@@ -139,6 +144,23 @@ tl::expected<void, ErrorCode> PyClient::setup_internal(
139144
LOG(INFO) << "Local buffer size is 0, skip registering local memory";
140145
}
141146

147+
return {};
148+
}
149+
150+
tl::expected<void, ErrorCode> PyClient::setup_internal(
151+
const std::string &local_hostname, const std::string &metadata_server,
152+
size_t global_segment_size, size_t local_buffer_size,
153+
const std::string &protocol, const std::string &rdma_devices,
154+
const std::string &master_server_addr) {
155+
// Common setups.
156+
auto result = common_setup_internal(local_hostname, metadata_server,
157+
local_buffer_size, protocol,
158+
rdma_devices, master_server_addr);
159+
if (!result.has_value()) {
160+
LOG(ERROR) << "Failed to setup PyClient";
161+
return tl::unexpected(result.error());
162+
}
163+
142164
// If global_segment_size is 0, skip mount segment;
143165
// If global_segment_size is larger than max_mr_size, split to multiple
144166
// segments.
@@ -182,6 +204,45 @@ int PyClient::setup(const std::string &local_hostname,
182204
protocol, rdma_devices, master_server_addr));
183205
}
184206

207+
tl::expected<void, ErrorCode> PyClient::setup_with_files_internal(
208+
const std::string &local_hostname, const std::string &metadata_server,
209+
const std::vector<std::string> &files, size_t local_buffer_size,
210+
const std::string &protocol, const std::string &protocol_arg,
211+
const std::string &master_server_addr) {
212+
// Common setups.
213+
auto result = common_setup_internal(local_hostname, metadata_server,
214+
local_buffer_size, protocol,
215+
protocol_arg, master_server_addr);
216+
if (!result.has_value()) {
217+
LOG(ERROR) << "Failed to setup PyClient";
218+
return tl::unexpected(result.error());
219+
}
220+
221+
// Mount file segments.
222+
for (auto &file : files) {
223+
auto result = client_->MountFileSegment(file);
224+
if (!result.has_value()) {
225+
LOG(ERROR) << "Failed to mount file " << file
226+
<< ", error=" << result.error();
227+
return tl::unexpected(result.error());
228+
}
229+
}
230+
231+
return {};
232+
}
233+
234+
int PyClient::setup_with_files(const std::string &local_hostname,
235+
const std::string &metadata_server,
236+
const std::vector<std::string> &files,
237+
size_t local_buffer_size,
238+
const std::string &protocol,
239+
const std::string &protocol_arg,
240+
const std::string &master_server_addr) {
241+
return to_py_ret(setup_with_files_internal(
242+
local_hostname, metadata_server, files, local_buffer_size, protocol,
243+
protocol_arg, master_server_addr));
244+
}
245+
185246
tl::expected<void, ErrorCode> PyClient::initAll_internal(
186247
const std::string &protocol_, const std::string &device_name,
187248
size_t mount_segment_size) {

0 commit comments

Comments
 (0)