GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 46.9% 267 / 0 / 569
Functions: 59.4% 19 / 0 / 32
Branches: 20.0% 226 / 0 / 1129
ps/brpc/brpc_ps_client.cpp
Line Branch Exec Source
1 #include "brpc_ps_client.h"
2
3 #include <brpc/channel.h>
4 #include <fmt/format.h>
5 #include <gflags/gflags.h>
6
7 #include <cstdint>
8 #include <cstring>
9 #include <future>
10 #include <string>
11 #include <vector>
12
13 #include <google/protobuf/arena.h>
14
15 #include "base/array.h"
16 #include "base/factory.h"
17 #include "base/flatc.h"
18 #include "base/log.h"
19 #include "base/timer.h"
20 #include "ps/base/parameters.h"
21 #include "ps_brpc.pb.h"
22
23 #ifdef ENABLE_PERF_REPORT
24 # include <chrono>
25 # include "base/report/report_client.h"
26 #endif
27
28 using recstoreps_brpc::CommandRequest;
29 using recstoreps_brpc::CommandResponse;
30 using recstoreps_brpc::GetParameterRequest;
31 using recstoreps_brpc::GetParameterResponse;
32 using recstoreps_brpc::InitEmbeddingTableRequest;
33 using recstoreps_brpc::InitEmbeddingTableResponse;
34 using recstoreps_brpc::PSCommand;
35 using recstoreps_brpc::PutParameterRequest;
36 using recstoreps_brpc::PutParameterResponse;
37 using recstoreps_brpc::UpdateParameterRequest;
38 using recstoreps_brpc::UpdateParameterResponse;
39
40 namespace {
41
42 50 const ParameterCompressReader* ExtractGetResponseReader(
43 const brpc::Controller& cntl,
44 const GetParameterResponse& response,
45 std::string* payload_storage,
46 int* payload_size) {
47
1/2
✓ Branch 2 taken 50 times.
✗ Branch 3 not taken.
 50 if (!cntl.response_attachment().empty()) {
48 50 payload_storage->clear();
49 50 cntl.response_attachment().copy_to(payload_storage);
50 50 *payload_size = payload_storage->size();
51 return reinterpret_cast<const ParameterCompressReader*>(
52 50 payload_storage->data());
53 }
54
55 ✗ *payload_size = response.parameter_value().size();
56 return reinterpret_cast<const ParameterCompressReader*>(
57 ✗ response.parameter_value().data());
58 }
59
60 } // namespace
61
62 namespace {
63
64 ✗ int BuildUpdateBlocksFromFlat(
65 const base::ConstArray<uint64_t>& keys,
66 const float* grads,
67 int64_t num_rows,
68 int64_t embedding_dim,
69 ParameterCompressor* compressor) {
70 ✗ if (grads == nullptr) {
71 ✗ LOG(ERROR) << "UpdateParameterFlat grads pointer is null";
72 ✗ return -1;
73 }
74 ✗ if (num_rows < 0 || embedding_dim <= 0) {
75 ✗ LOG(ERROR) << "UpdateParameterFlat invalid shape: rows=" << num_rows
76 ✗ << " dim=" << embedding_dim;
77 ✗ return -1;
78 }
79 ✗ if (keys.Size() != static_cast<size_t>(num_rows)) {
80 ✗ LOG(ERROR) << "UpdateParameterFlat keys/grads size mismatch: "
81 ✗ << keys.Size() << " vs " << num_rows;
82 ✗ return -1;
83 }
84
85 ✗ for (int64_t i = 0; i < num_rows; ++i) {
86 ✗ ParameterPack pack;
87 ✗ pack.key = keys[static_cast<size_t>(i)];
88 ✗ pack.dim = embedding_dim;
89 ✗ pack.emb_data = grads + i * embedding_dim;
90 ✗ compressor->AddItem(pack, nullptr);
91 }
92 ✗ return 0;
93 }
94
95 } // namespace
96
97 DEFINE_int32(brpc_timeout_ms, 5000, "brpc request timeout in milliseconds");
98 DEFINE_int32(brpc_max_retry, 3, "brpc max retry times");
99 DEFINE_bool(parameter_client_random_init_brpc, false, "");
100 DEFINE_bool(brpc_ps_use_rdma,
101 false,
102 "Use RDMA transport for the brpc PS channel (requires brpc built "
103 "with WITH_RDMA=ON).");
104
105 namespace {
106
107 // The brpc PS client is loaded as a shared library from Python and never calls
108 // gflags::ParseCommandLineFlags, so RDMA options are configured from env vars
109 // (consistent with the RECSTORE_RDMA_* convention used by the raw-verbs path).
110 // RECSTORE_BRPC_USE_RDMA=1 -> enable RDMA transport
111 // RECSTORE_BRPC_RDMA_DEVICE=mlx5_0 -> select the HCA (maps to brpc -rdma_device)
112 30 bool ResolveBrpcUseRdmaFromEnv(bool fallback) {
113 30 const char* value = std::getenv("RECSTORE_BRPC_USE_RDMA");
114
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 30 if (value == nullptr || *value == '\0') {
115 30 return fallback;
116 }
117 ✗ return std::string(value) != "0";
118 }
119
120 ✗ void ApplyBrpcRdmaDeviceFromEnv() {
121 ✗ const char* dev = std::getenv("RECSTORE_BRPC_RDMA_DEVICE");
122 ✗ if (dev != nullptr && *dev != '\0') {
123 ✗ google::SetCommandLineOption("rdma_device", dev);
124 }
125 ✗ }
126
127 } // namespace
128
129 // New constructor that takes JSON config
130 24 BRPCParameterClient::BRPCParameterClient(json config)
131
1/2
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
 24 : recstore::BasePSClient(config) {
132
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 host_ = config.value("host", "localhost");
133
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 port_ = config.value("port", 15000);
134
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 shard_ = config.value("shard", 0);
135
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 timeout_ms_ = config.value("timeout_ms", FLAGS_brpc_timeout_ms);
136
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 max_retry_ = config.value("max_retry", FLAGS_brpc_max_retry);
137
138
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 Initialize();
139
140 // Initialize bRPC channel
141
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 channel_ = std::make_shared<brpc::Channel>();
142
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 brpc::ChannelOptions options;
143 24 options.timeout_ms = timeout_ms_;
144 24 options.max_retry = max_retry_;
145 48 bool use_rdma_enabled = ResolveBrpcUseRdmaFromEnv(
146
2/4
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 24 times.
✗ Branch 5 not taken.
 24 config.value("use_rdma", FLAGS_brpc_ps_use_rdma));
147
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
 24 if (use_rdma_enabled) {
148 ✗ ApplyBrpcRdmaDeviceFromEnv();
149 }
150 #if BRPC_WITH_RDMA
151 options.use_rdma = use_rdma_enabled;
152 #endif
153
154
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 std::string server_addr = fmt::format("{}:{}", host_, port_);
155
2/4
✓ Branch 3 taken 24 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 24 times.
 24 if (channel_->Init(server_addr.c_str(), &options) != 0) {
156 ✗ LOG(ERROR) << "Failed to initialize bRPC channel to " << server_addr;
157 } else {
158
3/6
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 24 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 24 times.
✗ Branch 8 not taken.
 48 LOG(INFO) << "Initialized bRPC PS Client Shard " << shard_ << " at "
159
6/12
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 24 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 24 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 24 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 24 times.
✗ Branch 14 not taken.
✓ Branch 16 taken 24 times.
✗ Branch 17 not taken.
 24 << server_addr << " (use_rdma=" << use_rdma_enabled << ")";
160 }
161 24 }
162
163 // Legacy constructor for backward compatibility
164 6 BRPCParameterClient::BRPCParameterClient(
165 6 const std::string& host, int port, int shard)
166 : recstore::BasePSClient(
167
14/28
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 6 times.
✗ Branch 6 not taken.
✓ Branch 9 taken 6 times.
✗ Branch 10 not taken.
✓ Branch 12 taken 6 times.
✗ Branch 13 not taken.
✓ Branch 16 taken 6 times.
✗ Branch 17 not taken.
✓ Branch 19 taken 6 times.
✗ Branch 20 not taken.
✓ Branch 22 taken 18 times.
✓ Branch 23 taken 6 times.
✓ Branch 25 taken 12 times.
✓ Branch 26 taken 6 times.
✓ Branch 28 taken 12 times.
✓ Branch 29 taken 6 times.
✓ Branch 31 taken 12 times.
✓ Branch 32 taken 6 times.
✗ Branch 35 not taken.
✗ Branch 36 not taken.
✗ Branch 38 not taken.
✗ Branch 39 not taken.
✗ Branch 41 not taken.
✗ Branch 42 not taken.
✗ Branch 44 not taken.
✗ Branch 45 not taken.
 66 json{{"host", host}, {"port", port}, {"shard", shard}}),
168 6 host_(host),
169 6 port_(port),
170 6 shard_(shard),
171 6 timeout_ms_(FLAGS_brpc_timeout_ms),
172
2/4
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 6 times.
✗ Branch 6 not taken.
 18 max_retry_(FLAGS_brpc_max_retry) {
173
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 Initialize();
174
175
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 channel_ = std::make_shared<brpc::Channel>();
176
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 brpc::ChannelOptions options;
177 6 options.timeout_ms = timeout_ms_;
178 6 options.max_retry = max_retry_;
179
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 bool use_rdma_enabled = ResolveBrpcUseRdmaFromEnv(FLAGS_brpc_ps_use_rdma);
180
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (use_rdma_enabled) {
181 ✗ ApplyBrpcRdmaDeviceFromEnv();
182 }
183 #if BRPC_WITH_RDMA
184 options.use_rdma = use_rdma_enabled;
185 #endif
186
187 ✗ std::string server_addr = fmt::format("{}:{}", host, port);
188
2/4
✓ Branch 3 taken 6 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
 6 if (channel_->Init(server_addr.c_str(), &options) != 0) {
189 ✗ LOG(ERROR) << "Failed to initialize bRPC channel to " << server_addr;
190 } else {
191
3/6
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 6 times.
✗ Branch 8 not taken.
 12 LOG(INFO) << "Initialized bRPC PS Client Shard " << shard_ << " at "
192
6/12
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 6 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 6 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 6 times.
✗ Branch 14 not taken.
✓ Branch 16 taken 6 times.
✗ Branch 17 not taken.
 6 << server_addr << " (use_rdma=" << use_rdma_enabled << ")";
193 }
194 6 }
195
196 30 bool BRPCParameterClient::Initialize() { return true; }
197
198 ✗ int BRPCParameterClient::GetParameter(const base::ConstArray<uint64_t>& keys,
199 float* values) {
200 #ifdef ENABLE_PERF_REPORT
201 auto start_time = std::chrono::high_resolution_clock::now();
202 #endif
203
204 ✗ if (FLAGS_parameter_client_random_init_brpc) {
205 ✗ CHECK(0) << "todo implement";
206 return true;
207 }
208
209 int request_num =
210 ✗ (keys.Size() + MAX_PARAMETER_BATCH_BRPC - 1) / MAX_PARAMETER_BATCH_BRPC;
211 ✗ google::protobuf::Arena arena;
212 ✗ std::vector<GetParameterRequest*> requests(request_num);
213 ✗ std::vector<GetParameterResponse*> responses(request_num);
214 ✗ for (int i = 0; i < request_num; ++i) {
215 ✗ requests[i] = google::protobuf::Arena::Create<GetParameterRequest>(&arena);
216 ✗ responses[i] = google::protobuf::Arena::Create<GetParameterResponse>(&arena);
217 }
218 ✗ std::vector<brpc::Controller> controllers(request_num);
219 ✗ std::vector<int> key_sizes;
220
221 // Create stub
222 ✗ recstoreps_brpc::ParameterService_Stub stub(channel_.get());
223
224 #ifdef ENABLE_PERF_REPORT
225 auto wait_start_time = std::chrono::high_resolution_clock::now();
226 #endif
227
228 // Send async RPC requests
229 ✗ for (int start = 0, index = 0; start < keys.Size();
230 ✗ start += MAX_PARAMETER_BATCH_BRPC, ++index) {
231 int key_size =
232 ✗ std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH_BRPC);
233 ✗ key_sizes.push_back(key_size);
234
235 ✗ controllers[index].request_attachment().append(
236 ✗ reinterpret_cast<const char*>(&keys[start]),
237 ✗ sizeof(uint64_t) * key_size);
238
239 ✗ google::protobuf::Closure* done = brpc::NewCallback([]() { /* no-op */ });
240 ✗ stub.GetParameter(
241 ✗ &controllers[index], requests[index], responses[index], done);
242 }
243
244 // Wait for all RPCs to complete
245 ✗ for (int i = 0; i < request_num; ++i) {
246 ✗ brpc::Join(controllers[i].call_id());
247 ✗ if (controllers[i].Failed()) {
248 ✗ LOG(ERROR) << "bRPC GetParameter failed: " << controllers[i].ErrorText();
249 ✗ return false;
250 }
251 }
252
253 #ifdef ENABLE_PERF_REPORT
254 auto wait_end_time = std::chrono::high_resolution_clock::now();
255 auto wait_duration =
256 std::chrono::duration_cast<std::chrono::microseconds>(
257 wait_end_time - wait_start_time)
258 .count();
259 double wait_start_us =
260 std::chrono::duration_cast<std::chrono::microseconds>(
261 wait_start_time.time_since_epoch())
262 .count();
263 std::string wait_label =
264 "brpc_client::RPC_Call_And_Wait_Shard" + std::to_string(shard_);
265 FlameGraphData wait_fg = {
266 wait_label,
267 wait_start_us,
268 2, // level
269 static_cast<double>(wait_duration),
270 static_cast<double>(wait_duration)};
271 std::string unique_id =
272 "embread_debug|" + std::to_string(static_cast<uint64_t>(wait_start_us));
273 report_flame_graph("emb_read_flame_map", unique_id.c_str(), wait_fg);
274
275 double start_us_for_rpc =
276 std::chrono::duration_cast<std::chrono::microseconds>(
277 start_time.time_since_epoch())
278 .count();
279 std::string report_id_for_rpc =
280 "brpc_client::GetParameter|" +
281 std::to_string(static_cast<uint64_t>(start_us_for_rpc));
282 report("embread_stages",
283 report_id_for_rpc.c_str(),
284 "rpc_duration_us",
285 static_cast<double>(wait_duration));
286
287 auto deserialize_start_time = std::chrono::high_resolution_clock::now();
288 #endif
289
290 // Parse responses
291 ✗ size_t get_embedding_acc = 0;
292 ✗ int old_dimension = -1;
293 ✗ std::string payload_storage;
294
295 ✗ for (int i = 0; i < responses.size(); ++i) {
296 ✗ auto* response = responses[i];
297 ✗ int key_size = key_sizes[i];
298 ✗ int payload_size = 0;
299 ✗ auto parameters = ExtractGetResponseReader(
300 ✗ controllers[i], *response, &payload_storage, &payload_size);
301
302 ✗ if (parameters == nullptr || !parameters->Valid(payload_size)) {
303 ✗ LOG(ERROR) << "GetParameter invalid payload: " << payload_size;
304 ✗ return false;
305 }
306
307 ✗ if (parameters->size != key_size) {
308 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
309 ✗ << key_size;
310 ✗ return false;
311 }
312
313 ✗ for (int index = 0; index < parameters->item_size(); ++index) {
314 ✗ auto item = parameters->item(index);
315 ✗ if (item->dim != 0) {
316 ✗ if (old_dimension == -1)
317 ✗ old_dimension = item->dim;
318 ✗ CHECK_EQ(item->dim, old_dimension);
319 ✗ std::copy_n(
320 ✗ item->embedding, item->dim, values + item->dim * get_embedding_acc);
321 } else {
322 ✗ RECSTORE_LOG_EVERY_MS(ERROR, 2000)
323 ✗ << "error; not find key " << keys[get_embedding_acc] << " in ps";
324 }
325 ✗ get_embedding_acc++;
326 }
327 }
328
329 #ifdef ENABLE_PERF_REPORT
330 auto deserialize_end_time = std::chrono::high_resolution_clock::now();
331 auto deserialize_duration =
332 std::chrono::duration_cast<std::chrono::microseconds>(
333 deserialize_end_time - deserialize_start_time)
334 .count();
335 double deserialize_start_us =
336 std::chrono::duration_cast<std::chrono::microseconds>(
337 deserialize_start_time.time_since_epoch())
338 .count();
339 std::string des_label =
340 "brpc_client::Deserialize_Shard" + std::to_string(shard_);
341 FlameGraphData des_fg = {
342 des_label,
343 deserialize_start_us,
344 2, // level
345 static_cast<double>(deserialize_duration),
346 static_cast<double>(deserialize_duration)};
347 std::string des_unique_id =
348 "embread_debug|" +
349 std::to_string(static_cast<uint64_t>(deserialize_start_us));
350 report_flame_graph("emb_read_flame_map", des_unique_id.c_str(), des_fg);
351
352 double start_us_for_des =
353 std::chrono::duration_cast<std::chrono::microseconds>(
354 start_time.time_since_epoch())
355 .count();
356 std::string report_id_for_des =
357 "brpc_client::GetParameter|" +
358 std::to_string(static_cast<uint64_t>(start_us_for_des));
359 report("embread_stages",
360 report_id_for_des.c_str(),
361 "deserialize_duration_us",
362 static_cast<double>(deserialize_duration));
363 #endif
364
365 #ifdef ENABLE_PERF_REPORT
366 auto end_time = std::chrono::high_resolution_clock::now();
367 auto duration =
368 std::chrono::duration_cast<std::chrono::microseconds>(
369 end_time - start_time)
370 .count();
371 report("ps_client_latency",
372 "GetParameter",
373 "latency_us",
374 static_cast<double>(duration));
375
376 double start_us =
377 std::chrono::duration_cast<std::chrono::microseconds>(
378 start_time.time_since_epoch())
379 .count();
380 FlameGraphData fg_data = {
381 "brpc_client::GetParameter",
382 start_us,
383 1, // level
384 static_cast<double>(duration),
385 static_cast<double>(duration)};
386
387 std::string report_id = "brpc_client::GetParameter|" +
388 std::to_string(static_cast<uint64_t>(start_us));
389
390 report("embread_stages",
391 report_id.c_str(),
392 "duration_us",
393 static_cast<double>(duration));
394
395 report("embread_stages",
396 report_id.c_str(),
397 "request_size",
398 static_cast<double>(keys.Size()));
399
400 std::string final_unique_id =
401 "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us));
402 report_flame_graph("emb_read_flame_map", final_unique_id.c_str(), fg_data);
403 #endif
404
405 ✗ return true;
406 ✗ }
407
408 40 int BRPCParameterClient::GetParameter(const base::ConstArray<uint64_t>& keys,
409 std::vector<std::vector<float>>* values) {
410 #ifdef ENABLE_PERF_REPORT
411 auto start_time = std::chrono::high_resolution_clock::now();
412 #endif
413
414
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 40 times.
 40 if (FLAGS_parameter_client_random_init_brpc) {
415 ✗ values->clear();
416 ✗ values->reserve(keys.Size());
417 ✗ for (size_t i = 0; i < keys.Size(); i++)
418 ✗ values->emplace_back(std::vector<float>(128, 0.1));
419 ✗ return true;
420 }
421
422 40 values->clear();
423
1/2
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
 40 values->reserve(keys.Size());
424
425 int request_num =
426 40 (keys.Size() + MAX_PARAMETER_BATCH_BRPC - 1) / MAX_PARAMETER_BATCH_BRPC;
427
1/2
✓ Branch 1 taken 40 times.
✗ Branch 2 not taken.
 40 google::protobuf::Arena arena;
428
1/2
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
 40 std::vector<GetParameterRequest*> requests(request_num);
429
1/2
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
 40 std::vector<GetParameterResponse*> responses(request_num);
430
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 40 times.
 80 for (int i = 0; i < request_num; ++i) {
431
1/2
✓ Branch 1 taken 40 times.
✗ Branch 2 not taken.
 40 requests[i] = google::protobuf::Arena::Create<GetParameterRequest>(&arena);
432
1/2
✓ Branch 1 taken 40 times.
✗ Branch 2 not taken.
 40 responses[i] = google::protobuf::Arena::Create<GetParameterResponse>(&arena);
433 }
434
1/2
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
 40 std::vector<brpc::Controller> controllers(request_num);
435 40 std::vector<int> key_sizes;
436
437
1/2
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
 40 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
438
439 #ifdef ENABLE_PERF_REPORT
440 auto wait_start_time = std::chrono::high_resolution_clock::now();
441 #endif
442
443 // Send async RPC requests
444
2/2
✓ Branch 1 taken 40 times.
✓ Branch 2 taken 40 times.
 80 for (int start = 0, index = 0; start < keys.Size();
445 40 start += MAX_PARAMETER_BATCH_BRPC, ++index) {
446 int key_size =
447 40 std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH_BRPC);
448
1/2
✓ Branch 1 taken 40 times.
✗ Branch 2 not taken.
 40 key_sizes.push_back(key_size);
449
450
1/2
✓ Branch 3 taken 40 times.
✗ Branch 4 not taken.
 80 controllers[index].request_attachment().append(
451 40 reinterpret_cast<const char*>(&keys[start]),
452 40 sizeof(uint64_t) * key_size);
453
454
1/2
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
 40 google::protobuf::Closure* done = brpc::NewCallback([]() { /* no-op */ });
455
1/2
✓ Branch 1 taken 40 times.
✗ Branch 2 not taken.
 40 stub.GetParameter(
456 40 &controllers[index], requests[index], responses[index], done);
457 }
458
459 // Wait for all RPCs to complete
460
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 40 times.
 80 for (int i = 0; i < request_num; ++i) {
461
2/4
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 40 times.
✗ Branch 6 not taken.
 40 brpc::Join(controllers[i].call_id());
462
2/4
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 40 times.
 40 if (controllers[i].Failed()) {
463 ✗ LOG(ERROR) << "bRPC GetParameter failed: " << controllers[i].ErrorText();
464 ✗ return false;
465 }
466 }
467
468 #ifdef ENABLE_PERF_REPORT
469 auto wait_end_time = std::chrono::high_resolution_clock::now();
470 auto wait_duration =
471 std::chrono::duration_cast<std::chrono::microseconds>(
472 wait_end_time - wait_start_time)
473 .count();
474 double wait_start_us =
475 std::chrono::duration_cast<std::chrono::microseconds>(
476 wait_start_time.time_since_epoch())
477 .count();
478 std::string wait_label =
479 "brpc_client::RPC_Call_And_Wait_Shard" + std::to_string(shard_);
480 FlameGraphData wait_fg = {
481 wait_label,
482 wait_start_us,
483 2, // level
484 static_cast<double>(wait_duration),
485 static_cast<double>(wait_duration)};
486 std::string unique_id =
487 "embread_debug|" + std::to_string(static_cast<uint64_t>(wait_start_us));
488 report_flame_graph("emb_read_flame_map", unique_id.c_str(), wait_fg);
489
490 double start_us_for_rpc =
491 std::chrono::duration_cast<std::chrono::microseconds>(
492 start_time.time_since_epoch())
493 .count();
494 std::string report_id_for_rpc =
495 "brpc_client::GetParameter_Vec|" +
496 std::to_string(static_cast<uint64_t>(start_us_for_rpc));
497
498 report("embread_stages",
499 report_id_for_rpc.c_str(),
500 "rpc_duration_us",
501 static_cast<double>(wait_duration));
502
503 auto deserialize_start_time = std::chrono::high_resolution_clock::now();
504 #endif
505
506 // Parse responses
507 40 std::string payload_storage;
508
2/2
✓ Branch 1 taken 40 times.
✓ Branch 2 taken 40 times.
 80 for (int i = 0; i < responses.size(); ++i) {
509 40 auto* response = responses[i];
510 40 int key_size = key_sizes[i];
511 40 int payload_size = 0;
512
1/2
✓ Branch 1 taken 40 times.
✗ Branch 2 not taken.
 40 auto parameters = ExtractGetResponseReader(
513 40 controllers[i], *response, &payload_storage, &payload_size);
514
515
4/8
✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 40 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 40 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 40 times.
 40 if (parameters == nullptr || !parameters->Valid(payload_size)) {
516 ✗ LOG(ERROR) << "GetParameter(vector) invalid payload: " << payload_size;
517 ✗ return false;
518 }
519
520
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 40 times.
 40 if (unlikely(parameters->size != key_size)) {
521 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
522 ✗ << key_size;
523 ✗ return false;
524 }
525
526
2/2
✓ Branch 1 taken 266 times.
✓ Branch 2 taken 40 times.
 306 for (int index = 0; index < parameters->item_size(); ++index) {
527
1/2
✓ Branch 1 taken 266 times.
✗ Branch 2 not taken.
 266 auto item = parameters->item(index);
528
2/2
✓ Branch 0 taken 254 times.
✓ Branch 1 taken 12 times.
 266 if (item->dim != 0) {
529
1/2
✓ Branch 1 taken 254 times.
✗ Branch 2 not taken.
 254 values->emplace_back(
530
1/2
✓ Branch 2 taken 254 times.
✗ Branch 3 not taken.
 508 std::vector<float>(item->embedding, item->embedding + item->dim));
531 } else {
532
2/4
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 12 times.
✗ Branch 6 not taken.
 12 values->emplace_back(std::vector<float>(0));
533 }
534 }
535 }
536
537 #ifdef ENABLE_PERF_REPORT
538 auto deserialize_end_time = std::chrono::high_resolution_clock::now();
539 auto deserialize_duration =
540 std::chrono::duration_cast<std::chrono::microseconds>(
541 deserialize_end_time - deserialize_start_time)
542 .count();
543 double deserialize_start_us =
544 std::chrono::duration_cast<std::chrono::microseconds>(
545 deserialize_start_time.time_since_epoch())
546 .count();
547 std::string des_label =
548 "brpc_client::Deserialize_Shard" + std::to_string(shard_);
549 FlameGraphData des_fg = {
550 des_label,
551 deserialize_start_us,
552 2, // level
553 static_cast<double>(deserialize_duration),
554 static_cast<double>(deserialize_duration)};
555 std::string des_unique_id =
556 "embread_debug|" +
557 std::to_string(static_cast<uint64_t>(deserialize_start_us));
558 report_flame_graph("emb_read_flame_map", des_unique_id.c_str(), des_fg);
559
560 double start_us_for_des =
561 std::chrono::duration_cast<std::chrono::microseconds>(
562 start_time.time_since_epoch())
563 .count();
564 std::string report_id_for_des =
565 "brpc_client::GetParameter_Vec|" +
566 std::to_string(static_cast<uint64_t>(start_us_for_des));
567 report("embread_stages",
568 report_id_for_des.c_str(),
569 "deserialize_duration_us",
570 static_cast<double>(deserialize_duration));
571 #endif
572
573 #ifdef ENABLE_PERF_REPORT
574 auto end_time = std::chrono::high_resolution_clock::now();
575 auto duration =
576 std::chrono::duration_cast<std::chrono::microseconds>(
577 end_time - start_time)
578 .count();
579 report("ps_client_latency",
580 "GetParameter",
581 "latency_us",
582 static_cast<double>(duration));
583
584 double start_us =
585 std::chrono::duration_cast<std::chrono::microseconds>(
586 start_time.time_since_epoch())
587 .count();
588 FlameGraphData fg_data = {
589 "brpc_client::GetParameter_Vec",
590 start_us,
591 1, // level
592 static_cast<double>(duration),
593 static_cast<double>(duration)};
594
595 std::string report_id = "brpc_client::GetParameter_Vec|" +
596 std::to_string(static_cast<uint64_t>(start_us));
597
598 report("embread_stages",
599 report_id.c_str(),
600 "duration_us",
601 static_cast<double>(duration));
602
603 report("embread_stages",
604 report_id.c_str(),
605 "request_size",
606 static_cast<double>(keys.Size()));
607
608 std::string final_unique_id =
609 "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us));
610 report_flame_graph("emb_read_flame_map", final_unique_id.c_str(), fg_data);
611 #endif
612
613 40 return true;
614 40 }
615
616 10 static void OnPrefetchDone(BrpcPrefetchBatch* batch) {
617 10 batch->completed_count_++;
618 10 }
619
620 8 static void OnPrewriteDone(BrpcPrewriteBatch* batch) {
621 8 batch->completed_count_++;
622 8 }
623
624 uint64_t
625 10 BRPCParameterClient::PrefetchParameter(const base::ConstArray<uint64_t>& keys) {
626 10 uint64_t prefetch_id = next_prefetch_id_++;
627 int request_num =
628 10 (keys.Size() + MAX_PARAMETER_BATCH_BRPC - 1) / MAX_PARAMETER_BATCH_BRPC;
629
630 // Construct in map so batch pointers stay valid
631
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 auto it = prefetch_batches_.emplace(prefetch_id, request_num).first;
632 10 struct BrpcPrefetchBatch* pb = &it->second;
633
634
1/2
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
635
636
2/2
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 10 times.
 20 for (int start = 0, index = 0; start < keys.Size();
637 10 start += MAX_PARAMETER_BATCH_BRPC, ++index) {
638 int key_size =
639 10 std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH_BRPC);
640 10 pb->key_sizes_[index] = key_size;
641
642
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 GetParameterRequest request;
643
644
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 pb->controllers_[index] = std::make_unique<brpc::Controller>();
645
1/2
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
 20 pb->controllers_[index]->request_attachment().append(
646 10 reinterpret_cast<const char*>(&keys[start]),
647 10 sizeof(uint64_t) * key_size);
648
649
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 google::protobuf::Closure* done = brpc::NewCallback(OnPrefetchDone, pb);
650
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 stub.GetParameter(
651 10 pb->controllers_[index].get(), &request, &pb->responses_[index], done);
652 10 }
653
654 10 return prefetch_id;
655 10 }
656
657 ✗ bool BRPCParameterClient::IsPrefetchDone(uint64_t prefetch_id) {
658 ✗ auto it = prefetch_batches_.find(prefetch_id);
659 ✗ if (it == prefetch_batches_.end()) {
660 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
661 ✗ return false;
662 }
663
664 ✗ auto& pb = it->second;
665
666 ✗ return pb.completed_count_ == pb.batch_size_;
667 }
668
669 10 void BRPCParameterClient::WaitForPrefetch(uint64_t prefetch_id) {
670
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 auto it = prefetch_batches_.find(prefetch_id);
671
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
 10 if (it == prefetch_batches_.end()) {
672 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
673 ✗ return;
674 }
675 10 auto& pb = it->second;
676
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 10 times.
 20 for (int i = 0; i < pb.batch_size_; ++i) {
677
1/2
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 if (pb.controllers_[i]) {
678
2/4
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
✓ Branch 6 taken 10 times.
✗ Branch 7 not taken.
 10 brpc::Join(pb.controllers_[i]->call_id());
679 }
680 }
681 10 pb.completed_count_ = pb.batch_size_;
682 }
683
684 10 bool BRPCParameterClient::GetPrefetchResult(
685 uint64_t prefetch_id, std::vector<std::vector<float>>* values) {
686
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 auto it = prefetch_batches_.find(prefetch_id);
687
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
 10 if (it == prefetch_batches_.end()) {
688 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
689 ✗ return false;
690 }
691
692 10 auto& pb = it->second;
693 10 int request_num = pb.batch_size_;
694
695 10 values->clear();
696 10 int keys_size = 0;
697
2/2
✓ Branch 4 taken 10 times.
✓ Branch 5 taken 10 times.
 20 for (const auto& size : pb.key_sizes_) {
698 10 keys_size += size;
699 }
700
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 values->reserve(keys_size);
701
702
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 10 times.
 20 for (int i = 0; i < request_num; ++i) {
703
2/4
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 10 times.
 10 if (pb.controllers_[i]->Failed()) {
704 ✗ LOG(ERROR) << "Prefetch request failed: "
705 ✗ << pb.controllers_[i]->ErrorText();
706 ✗ return false;
707 }
708
709 10 auto& response = pb.responses_[i];
710 10 int key_size = pb.key_sizes_[i];
711 10 std::string payload_storage;
712 10 int payload_size = 0;
713
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 auto parameters = ExtractGetResponseReader(
714 10 *pb.controllers_[i], response, &payload_storage, &payload_size);
715
716
4/8
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 10 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 10 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 10 times.
 10 if (parameters == nullptr || !parameters->Valid(payload_size)) {
717 ✗ LOG(ERROR) << "Prefetch invalid payload: " << payload_size;
718 ✗ return false;
719 }
720
721
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
 10 if (unlikely(parameters->size != key_size)) {
722 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
723 ✗ << key_size;
724 ✗ return false;
725 }
726
727
2/2
✓ Branch 1 taken 106 times.
✓ Branch 2 taken 10 times.
 116 for (int index = 0; index < parameters->item_size(); ++index) {
728
1/2
✓ Branch 1 taken 106 times.
✗ Branch 2 not taken.
 106 auto item = parameters->item(index);
729
1/2
✓ Branch 0 taken 106 times.
✗ Branch 1 not taken.
 106 if (item->dim != 0) {
730
1/2
✓ Branch 1 taken 106 times.
✗ Branch 2 not taken.
 106 values->emplace_back(
731
1/2
✓ Branch 2 taken 106 times.
✗ Branch 3 not taken.
 212 std::vector<float>(item->embedding, item->embedding + item->dim));
732 } else {
733 ✗ values->emplace_back(std::vector<float>(0));
734 }
735 }
736
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 }
737
738 // Remove completed batch
739
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 prefetch_batches_.erase(it);
740
741 10 return true;
742 }
743
744 ✗ bool BRPCParameterClient::GetPrefetchResultFlat(
745 uint64_t prefetch_id,
746 std::vector<float>* values,
747 int64_t* num_rows,
748 int64_t embedding_dim) {
749 ✗ auto it = prefetch_batches_.find(prefetch_id);
750 ✗ if (it == prefetch_batches_.end()) {
751 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
752 ✗ return false;
753 }
754 ✗ if (values == nullptr || num_rows == nullptr) {
755 ✗ LOG(ERROR) << "GetPrefetchResultFlat output pointer is null";
756 ✗ return false;
757 }
758
759 ✗ auto& pb = it->second;
760 ✗ int request_num = pb.batch_size_;
761 ✗ int total_keys = 0;
762 ✗ for (const auto& size : pb.key_sizes_) {
763 ✗ total_keys += size;
764 }
765
766 ✗ *num_rows = static_cast<int64_t>(total_keys);
767 ✗ values->assign(
768 ✗ static_cast<size_t>(*num_rows) * static_cast<size_t>(embedding_dim),
769 ✗ 0.0f);
770
771 ✗ size_t row_offset = 0;
772 ✗ for (int i = 0; i < request_num; ++i) {
773 ✗ if (pb.controllers_[i]->Failed()) {
774 ✗ LOG(ERROR) << "Prefetch request failed: "
775 ✗ << pb.controllers_[i]->ErrorText();
776 ✗ return false;
777 }
778
779 ✗ auto& response = pb.responses_[i];
780 ✗ int key_size = pb.key_sizes_[i];
781 ✗ std::string payload_storage;
782 ✗ int payload_size = 0;
783 ✗ auto parameters = ExtractGetResponseReader(
784 ✗ *pb.controllers_[i], response, &payload_storage, &payload_size);
785
786 ✗ if (parameters == nullptr || !parameters->Valid(payload_size)) {
787 ✗ LOG(ERROR) << "Prefetch invalid payload: " << payload_size;
788 ✗ return false;
789 }
790
791 ✗ if (unlikely(parameters->size != key_size)) {
792 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
793 ✗ << key_size;
794 ✗ return false;
795 }
796
797 ✗ for (int index = 0; index < parameters->item_size();
798 ✗ ++index, ++row_offset) {
799 ✗ auto item = parameters->item(index);
800 ✗ if (item->dim != 0) {
801 const int64_t copy_d =
802 ✗ std::min<int64_t>(embedding_dim, static_cast<int64_t>(item->dim));
803 ✗ std::memcpy(values->data() + row_offset * embedding_dim,
804 ✗ item->embedding,
805 ✗ static_cast<size_t>(copy_d) * sizeof(float));
806 }
807 }
808 ✗ }
809
810 ✗ prefetch_batches_.erase(it);
811 ✗ return true;
812 }
813
814 34 bool BRPCParameterClient::ClearPS() {
815
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 CommandRequest request;
816
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 CommandResponse response;
817
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 request.set_command(PSCommand::CLEAR_PS);
818
819
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 brpc::Controller cntl;
820
1/2
✓ Branch 2 taken 34 times.
✗ Branch 3 not taken.
 34 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
821
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 stub.Command(&cntl, &request, &response, nullptr);
822
823
2/4
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 34 times.
 34 if (cntl.Failed()) {
824 ✗ LOG(ERROR) << "bRPC Command failed: " << cntl.ErrorText();
825 ✗ return false;
826 }
827 34 return true;
828 34 }
829
830 6 bool BRPCParameterClient::LoadFakeData(int64_t data) {
831
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 CommandRequest request;
832
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 CommandResponse response;
833
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 request.set_command(PSCommand::LOAD_FAKE_DATA);
834
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 request.add_arg1(&data, sizeof(int64_t));
835
836
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 brpc::Controller cntl;
837
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
838
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 stub.Command(&cntl, &request, &response, nullptr);
839
840
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
 6 if (cntl.Failed()) {
841 ✗ LOG(ERROR) << "bRPC LoadFakeData failed: " << cntl.ErrorText();
842 ✗ return false;
843 }
844
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (response.reply().size() != static_cast<size_t>(data)) {
845 ✗ LOG(ERROR) << "bRPC LoadFakeData reply size mismatch: expected " << data
846 ✗ << ", got " << response.reply().size();
847 ✗ return false;
848 }
849 6 return true;
850 6 }
851
852 6 bool BRPCParameterClient::DumpFakeData(int64_t n) {
853
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 CommandRequest request;
854
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 CommandResponse response;
855
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 request.set_command(PSCommand::DUMP_FAKE_DATA);
856
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 request.add_arg1(&n, sizeof(int64_t));
857
858
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 brpc::Controller cntl;
859
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
860
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 stub.Command(&cntl, &request, &response, nullptr);
861
862
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
 6 if (cntl.Failed()) {
863 ✗ LOG(ERROR) << "bRPC DumpFakeData failed: " << cntl.ErrorText();
864 ✗ return false;
865 }
866
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (response.reply() != "ok") {
867 ✗ LOG(ERROR) << "bRPC DumpFakeData unexpected reply: " << response.reply();
868 ✗ return false;
869 }
870 6 return true;
871 6 }
872
873 ✗ bool BRPCParameterClient::LoadCkpt(
874 const std::vector<std::string>& model_config_path,
875 const std::vector<std::string>& emb_file_path) {
876 ✗ CommandRequest request;
877 ✗ CommandResponse response;
878 ✗ request.set_command(PSCommand::RELOAD_PS);
879
880 ✗ for (auto& each : model_config_path) {
881 ✗ request.add_arg1(each);
882 }
883 ✗ for (auto& each : emb_file_path) {
884 ✗ request.add_arg2(each);
885 }
886
887 ✗ brpc::Controller cntl;
888 ✗ recstoreps_brpc::ParameterService_Stub stub(channel_.get());
889 ✗ stub.Command(&cntl, &request, &response, nullptr);
890
891 ✗ if (cntl.Failed()) {
892 ✗ LOG(ERROR) << "bRPC LoadCkpt failed: " << cntl.ErrorText();
893 ✗ return false;
894 }
895 ✗ return true;
896 ✗ }
897
898 20 bool BRPCParameterClient::PutParameter(
899 const std::vector<uint64_t>& keys,
900 const std::vector<std::vector<float>>& values) {
901 #ifdef ENABLE_PERF_REPORT
902 auto start_time = std::chrono::high_resolution_clock::now();
903 #endif
904
905
1/2
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
 20 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
906
907
2/2
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 20 times.
 40 for (int start = 0, index = 0; start < keys.size();
908 20 start += MAX_PARAMETER_BATCH_BRPC, ++index) {
909 int key_size =
910 20 std::min((int)(keys.size() - start), MAX_PARAMETER_BATCH_BRPC);
911
912
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 PutParameterRequest request;
913
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 PutParameterResponse response;
914
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 ParameterCompressor compressor;
915
916
2/2
✓ Branch 0 taken 234 times.
✓ Branch 1 taken 20 times.
 254 for (int i = start; i < start + key_size; i++) {
917 234 auto each_key = keys[i];
918 234 auto& embedding = values[i];
919 234 ParameterPack parameter_pack;
920 234 parameter_pack.key = each_key;
921 234 parameter_pack.dim = embedding.size();
922 234 parameter_pack.emb_data = embedding.data();
923
1/2
✓ Branch 1 taken 234 times.
✗ Branch 2 not taken.
 234 compressor.AddItem(parameter_pack, nullptr);
924 }
925
926
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 brpc::Controller cntl;
927
1/2
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
 20 compressor.AppendToIOBuf(&cntl.request_attachment());
928
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 stub.PutParameter(&cntl, &request, &response, nullptr);
929
930
2/4
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 20 times.
 20 if (cntl.Failed()) {
931 ✗ LOG(ERROR) << "bRPC PutParameter failed: " << cntl.ErrorText();
932 ✗ return false;
933 }
934
4/8
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 20 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 20 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 20 times.
✗ Branch 11 not taken.
 20 }
935
936 #ifdef ENABLE_PERF_REPORT
937 auto end_time = std::chrono::high_resolution_clock::now();
938 auto duration =
939 std::chrono::duration_cast<std::chrono::microseconds>(
940 end_time - start_time)
941 .count();
942 report("ps_client_latency",
943 "PutParameter",
944 "latency_us",
945 static_cast<double>(duration));
946 #endif
947
948 20 return true;
949 20 }
950
951 ✗ int BRPCParameterClient::AsyncGetParameter(
952 const base::ConstArray<uint64_t>& keys, float* values) {
953 ✗ return GetParameter(keys, values);
954 }
955
956 14 int BRPCParameterClient::PutParameter(
957 const base::ConstArray<uint64_t>& keys,
958 const std::vector<std::vector<float>>& values) {
959
1/2
✓ Branch 5 taken 14 times.
✗ Branch 6 not taken.
 14 std::vector<uint64_t> key_vec(keys.Data(), keys.Data() + keys.Size());
960
1/2
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
 14 bool success = PutParameter(key_vec, values);
961
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 28 return success ? 1 : 0;
962 14 }
963
964 ✗ void BRPCParameterClient::Command(recstore::PSCommand command) {
965 ✗ switch (command) {
966 ✗ case recstore::PSCommand::CLEAR_PS:
967 ✗ ClearPS();
968 ✗ break;
969 ✗ case recstore::PSCommand::RELOAD_PS:
970 ✗ LOG(WARNING) << "RELOAD_PS command requires additional parameters";
971 ✗ break;
972 ✗ case recstore::PSCommand::LOAD_FAKE_DATA: {
973 ✗ int64_t fake_data = 1000;
974 ✗ LoadFakeData(fake_data);
975 ✗ } break;
976 ✗ case recstore::PSCommand::DUMP_FAKE_DATA: {
977 ✗ DumpFakeData(4096);
978 ✗ } break;
979 ✗ default:
980 ✗ LOG(ERROR) << "Unknown PS command: " << static_cast<int>(command);
981 ✗ break;
982 }
983 ✗ }
984
985 ✗ int BRPCParameterClient::UpdateParameter(
986 const std::string& table_name,
987 const base::ConstArray<uint64_t>& keys,
988 const std::vector<std::vector<float>>* grads) {
989 #ifdef ENABLE_PERF_REPORT
990 auto start_time = std::chrono::high_resolution_clock::now();
991 const uint64_t trace_id = recstore::g_trace_id;
992 #endif
993 ✗ if (grads == nullptr) {
994 ✗ LOG(ERROR) << "UpdateParameter grads pointer is null";
995 ✗ return -1;
996 }
997 ✗ if (keys.Size() != grads->size()) {
998 ✗ LOG(ERROR) << "UpdateParameter keys/grads size mismatch: " << keys.Size()
999 ✗ << " vs " << grads->size();
1000 ✗ return -1;
1001 }
1002
1003 ✗ ParameterCompressor compressor;
1004 ✗ for (size_t i = 0; i < keys.Size(); ++i) {
1005 ✗ ParameterPack pack;
1006 ✗ pack.key = keys[i];
1007 ✗ pack.dim = grads->at(i).size();
1008 ✗ pack.emb_data = grads->at(i).data();
1009 ✗ compressor.AddItem(pack, nullptr);
1010 }
1011 #ifdef ENABLE_PERF_REPORT
1012 auto serialize_done_time = std::chrono::high_resolution_clock::now();
1013 #endif
1014 ✗ if (keys.Size() == 0) {
1015 ✗ LOG(WARNING) << "UpdateParameter no gradients to send";
1016 ✗ return 0;
1017 }
1018
1019 ✗ UpdateParameterRequest request;
1020 ✗ UpdateParameterResponse response;
1021 request.set_table_name(table_name);
1022
1023 ✗ brpc::Controller cntl;
1024 #ifdef ENABLE_PERF_REPORT
1025 if (trace_id != 0) {
1026 cntl.http_request().SetHeader(
1027 "x-recstore-trace-id", std::to_string(trace_id));
1028 }
1029 auto rpc_start_time = std::chrono::high_resolution_clock::now();
1030 #endif
1031 ✗ compressor.AppendToIOBuf(&cntl.request_attachment());
1032 ✗ recstoreps_brpc::ParameterService_Stub stub(channel_.get());
1033 ✗ stub.UpdateParameter(&cntl, &request, &response, nullptr);
1034 ✗ if (cntl.Failed()) {
1035 ✗ LOG(ERROR) << "UpdateParameter RPC failed: " << cntl.ErrorText();
1036 ✗ return -1;
1037 }
1038
1039 #ifdef ENABLE_PERF_REPORT
1040 auto end_time = std::chrono::high_resolution_clock::now();
1041 auto duration =
1042 std::chrono::duration_cast<std::chrono::microseconds>(
1043 end_time - start_time)
1044 .count();
1045 auto serialize_duration =
1046 std::chrono::duration_cast<std::chrono::microseconds>(
1047 serialize_done_time - start_time)
1048 .count();
1049 auto rpc_duration =
1050 std::chrono::duration_cast<std::chrono::microseconds>(
1051 end_time - rpc_start_time)
1052 .count();
1053 report("ps_client_latency",
1054 "UpdateParameter",
1055 "latency_us",
1056 static_cast<double>(duration));
1057
1058 const uint64_t effective_trace_id =
1059 trace_id == 0
1060 ? static_cast<uint64_t>(
1061 std::chrono::duration_cast<std::chrono::microseconds>(
1062 start_time.time_since_epoch())
1063 .count())
1064 : trace_id;
1065 std::string stage_id =
1066 "brpc_client::EmbUpdate|" + std::to_string(effective_trace_id);
1067 report("embupdate_stages",
1068 stage_id.c_str(),
1069 "client_serialize_us",
1070 static_cast<double>(serialize_duration));
1071 report("embupdate_stages",
1072 stage_id.c_str(),
1073 "client_rpc_us",
1074 static_cast<double>(rpc_duration));
1075 report("embupdate_stages",
1076 stage_id.c_str(),
1077 "client_total_us",
1078 static_cast<double>(duration));
1079 report("embupdate_stages",
1080 stage_id.c_str(),
1081 "client_request_size",
1082 static_cast<double>(keys.Size()));
1083 #endif
1084
1085 ✗ return response.success() ? 0 : -1;
1086 ✗ }
1087
1088 ✗ int BRPCParameterClient::UpdateParameterFlat(
1089 const std::string& table_name,
1090 const base::ConstArray<uint64_t>& keys,
1091 const float* grads,
1092 int64_t num_rows,
1093 int64_t embedding_dim) {
1094 #ifdef ENABLE_PERF_REPORT
1095 auto start_time = std::chrono::high_resolution_clock::now();
1096 const uint64_t trace_id = recstore::g_trace_id;
1097 #endif
1098 ✗ if (keys.Size() == 0) {
1099 ✗ return 0;
1100 }
1101
1102 ✗ ParameterCompressor compressor;
1103 ✗ if (BuildUpdateBlocksFromFlat(
1104 ✗ keys, grads, num_rows, embedding_dim, &compressor) != 0) {
1105 ✗ return -1;
1106 }
1107 #ifdef ENABLE_PERF_REPORT
1108 auto serialize_done_time = std::chrono::high_resolution_clock::now();
1109 #endif
1110
1111 ✗ UpdateParameterRequest request;
1112 ✗ UpdateParameterResponse response;
1113 request.set_table_name(table_name);
1114
1115 ✗ brpc::Controller cntl;
1116 #ifdef ENABLE_PERF_REPORT
1117 if (trace_id != 0) {
1118 cntl.http_request().SetHeader(
1119 "x-recstore-trace-id", std::to_string(trace_id));
1120 }
1121 auto rpc_start_time = std::chrono::high_resolution_clock::now();
1122 #endif
1123 ✗ compressor.AppendToIOBuf(&cntl.request_attachment());
1124 ✗ recstoreps_brpc::ParameterService_Stub stub(channel_.get());
1125 ✗ stub.UpdateParameter(&cntl, &request, &response, nullptr);
1126 ✗ if (cntl.Failed()) {
1127 ✗ LOG(ERROR) << "UpdateParameterFlat RPC failed: " << cntl.ErrorText();
1128 ✗ return -1;
1129 }
1130
1131 #ifdef ENABLE_PERF_REPORT
1132 auto end_time = std::chrono::high_resolution_clock::now();
1133 auto duration =
1134 std::chrono::duration_cast<std::chrono::microseconds>(
1135 end_time - start_time)
1136 .count();
1137 auto serialize_duration =
1138 std::chrono::duration_cast<std::chrono::microseconds>(
1139 serialize_done_time - start_time)
1140 .count();
1141 auto rpc_duration =
1142 std::chrono::duration_cast<std::chrono::microseconds>(
1143 end_time - rpc_start_time)
1144 .count();
1145 report("ps_client_latency",
1146 "UpdateParameterFlat",
1147 "latency_us",
1148 static_cast<double>(duration));
1149
1150 const uint64_t effective_trace_id =
1151 trace_id == 0
1152 ? static_cast<uint64_t>(
1153 std::chrono::duration_cast<std::chrono::microseconds>(
1154 start_time.time_since_epoch())
1155 .count())
1156 : trace_id;
1157 std::string stage_id =
1158 "brpc_client::EmbUpdate|" + std::to_string(effective_trace_id);
1159 report("embupdate_stages",
1160 stage_id.c_str(),
1161 "client_serialize_us",
1162 static_cast<double>(serialize_duration));
1163 report("embupdate_stages",
1164 stage_id.c_str(),
1165 "client_rpc_us",
1166 static_cast<double>(rpc_duration));
1167 report("embupdate_stages",
1168 stage_id.c_str(),
1169 "client_total_us",
1170 static_cast<double>(duration));
1171 report("embupdate_stages",
1172 stage_id.c_str(),
1173 "client_request_size",
1174 static_cast<double>(num_rows));
1175 report("embupdate_stages",
1176 stage_id.c_str(),
1177 "client_embedding_dim",
1178 static_cast<double>(embedding_dim));
1179 #endif
1180
1181 ✗ return response.success() ? 0 : -1;
1182 ✗ }
1183
1184 ✗ int BRPCParameterClient::InitEmbeddingTable(
1185 const std::string& table_name,
1186 const recstore::EmbeddingTableConfig& config) {
1187 #ifdef ENABLE_PERF_REPORT
1188 auto start_time = std::chrono::high_resolution_clock::now();
1189 #endif
1190
1191 ✗ InitEmbeddingTableRequest request;
1192 ✗ InitEmbeddingTableResponse response;
1193 request.set_table_name(table_name);
1194 ✗ request.set_config_payload(config.Serialize());
1195
1196 ✗ brpc::Controller cntl;
1197 ✗ recstoreps_brpc::ParameterService_Stub stub(channel_.get());
1198 ✗ stub.InitEmbeddingTable(&cntl, &request, &response, nullptr);
1199 ✗ if (cntl.Failed()) {
1200 ✗ LOG(ERROR) << "InitEmbeddingTable RPC failed: " << cntl.ErrorText();
1201 ✗ return -1;
1202 }
1203
1204 #ifdef ENABLE_PERF_REPORT
1205 auto end_time = std::chrono::high_resolution_clock::now();
1206 auto duration =
1207 std::chrono::duration_cast<std::chrono::microseconds>(
1208 end_time - start_time)
1209 .count();
1210 report("ps_client_latency",
1211 "InitEmbeddingTable",
1212 "latency_us",
1213 static_cast<double>(duration));
1214 #endif
1215
1216 ✗ return response.success() ? 0 : -1;
1217 ✗ }
1218
1219 8 uint64_t BRPCParameterClient::EmbWriteAsync(const base::RecTensor& keys,
1220 const base::RecTensor& values) {
1221
3/6
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 8 times.
 8 if (keys.dtype() != base::DataType::UINT64 || keys.dim() != 1) {
1222 ✗ LOG(ERROR) << "EmbWriteAsync expects keys as 1D UINT64 tensor, got dtype="
1223 ✗ << base::DataTypeToString(keys.dtype())
1224 ✗ << ", dim=" << keys.dim();
1225 ✗ return 0;
1226 }
1227
3/6
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 8 times.
 8 if (values.dtype() != base::DataType::FLOAT32 || values.dim() != 2) {
1228 ✗ LOG(ERROR)
1229 << "EmbWriteAsync expects values as 2D FLOAT32 tensor, got dtype="
1230 ✗ << base::DataTypeToString(values.dtype()) << ", dim=" << values.dim();
1231 ✗ return 0;
1232 }
1233
3/6
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 8 times.
 8 if (values.shape(0) != keys.shape(0)) {
1234 ✗ LOG(ERROR) << "EmbWriteAsync row mismatch: keys=" << keys.shape(0)
1235 ✗ << ", values=" << values.shape(0);
1236 ✗ return 0;
1237 }
1238
2/4
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
 8 if (values.shape(1) <= 0) {
1239 ✗ LOG(ERROR) << "EmbWriteAsync invalid embedding dim: " << values.shape(1);
1240 ✗ return 0;
1241 }
1242
1243
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 const uint64_t* key_data = keys.data_as<uint64_t>();
1244
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 const float* value_data = values.data_as<float>();
1245
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 int64_t key_count = keys.shape(0);
1246
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 int64_t emb_dim = values.shape(1);
1247
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (key_count == 0) {
1248 ✗ return 0;
1249 }
1250
1251 8 uint64_t prewrite_id = next_prewrite_id_++;
1252 8 int request_num =
1253 8 (static_cast<int>(key_count) + MAX_PARAMETER_BATCH_BRPC - 1) /
1254 MAX_PARAMETER_BATCH_BRPC;
1255
1256
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 auto it = prewrite_batches_.emplace(prewrite_id, request_num).first;
1257 8 struct BrpcPrewriteBatch* pb = &it->second;
1258
1259
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 recstoreps_brpc::ParameterService_Stub stub(channel_.get());
1260
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 for (int start = 0, index = 0; start < key_count;
1261 8 start += MAX_PARAMETER_BATCH_BRPC, ++index) {
1262 int key_size =
1263 8 std::min(static_cast<int>(key_count - start), MAX_PARAMETER_BATCH_BRPC);
1264 8 pb->key_sizes_[index] = key_size;
1265
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 pb->controllers_[index] = std::make_unique<brpc::Controller>();
1266
1267
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 ParameterCompressor compressor;
1268
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 8 times.
 104 for (int i = 0; i < key_size; ++i) {
1269 96 int64_t row = start + i;
1270 96 ParameterPack parameter_pack;
1271 96 parameter_pack.key = key_data[row];
1272 96 parameter_pack.dim = emb_dim;
1273 96 parameter_pack.emb_data = value_data + row * emb_dim;
1274
1/2
✓ Branch 1 taken 96 times.
✗ Branch 2 not taken.
 96 compressor.AddItem(parameter_pack, nullptr);
1275 }
1276
1277
1/2
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
 8 compressor.AppendToIOBuf(&pb->controllers_[index]->request_attachment());
1278
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 google::protobuf::Closure* done = brpc::NewCallback(OnPrewriteDone, pb);
1279
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 stub.PutParameter(
1280 8 pb->controllers_[index].get(),
1281 8 &pb->requests_[index],
1282 8 &pb->responses_[index],
1283 done);
1284 8 }
1285
1286 8 return prewrite_id;
1287 8 }
1288
1289 ✗ bool BRPCParameterClient::IsWriteDone(uint64_t write_id) {
1290 ✗ auto it = prewrite_batches_.find(write_id);
1291 ✗ if (it == prewrite_batches_.end()) {
1292 ✗ LOG(ERROR) << "Invalid prewrite_id: " << write_id;
1293 ✗ return false;
1294 }
1295 ✗ auto& pb = it->second;
1296 ✗ return pb.completed_count_ == pb.batch_size_;
1297 }
1298
1299 8 void BRPCParameterClient::WaitForWrite(uint64_t write_id) {
1300
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 auto it = prewrite_batches_.find(write_id);
1301
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (it == prewrite_batches_.end()) {
1302 ✗ LOG(ERROR) << "Invalid prewrite_id: " << write_id;
1303 ✗ return;
1304 }
1305 8 auto& pb = it->second;
1306
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 for (int i = 0; i < pb.batch_size_; ++i) {
1307
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (!pb.controllers_[i]) {
1308 ✗ continue;
1309 }
1310
2/4
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
✓ Branch 6 taken 8 times.
✗ Branch 7 not taken.
 8 brpc::Join(pb.controllers_[i]->call_id());
1311
2/4
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
 8 if (pb.controllers_[i]->Failed()) {
1312 ✗ LOG(ERROR) << "Async PutParameter failed: "
1313 ✗ << pb.controllers_[i]->ErrorText();
1314 }
1315 }
1316 8 pb.completed_count_ = pb.batch_size_;
1317
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 prewrite_batches_.erase(it);
1318 }
1319
1320 // Register BRPCParameterClient with the factory
1321 using BasePSClient = recstore::BasePSClient;
1322 FACTORY_REGISTER(BasePSClient, brpc, BRPCParameterClient, json);
1323