GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 51.6% 313 / 0 / 607
Functions: 63.0% 17 / 0 / 27
Branches: 22.0% 225 / 0 / 1024
ps/grpc/grpc_ps_client.cpp
Line Branch Exec Source
1 #include "grpc_ps_client.h"
2
3 #include <fmt/format.h>
4 #include <grpcpp/grpcpp.h>
5
6 #include <cstdint>
7 #include <cstring>
8 #include <future>
9 #include <string>
10 #include <vector>
11
12 #include "base/array.h"
13 #include "base/factory.h"
14 #include "base/flatc.h"
15 #include "base/log.h"
16 #include "base/timer.h"
17 #include "ps/base/parameters.h"
18 #include "ps.grpc.pb.h"
19 #include "ps.pb.h"
20
21 #ifdef ENABLE_PERF_REPORT
22 # include "base/report/report_client.h"
23 # include <chrono>
24 #endif
25
26 using grpc::Channel;
27 using grpc::ClientAsyncResponseReader;
28 using grpc::ClientContext;
29 using grpc::Status;
30 using recstoreps::CommandRequest;
31 using recstoreps::CommandResponse;
32 using recstoreps::GetParameterRequest;
33 using recstoreps::GetParameterResponse;
34 using recstoreps::InitEmbeddingTableRequest;
35 using recstoreps::InitEmbeddingTableResponse;
36 using recstoreps::PSCommand;
37 using recstoreps::PutParameterRequest;
38 using recstoreps::PutParameterResponse;
39 using recstoreps::UpdateParameterRequest;
40 using recstoreps::UpdateParameterResponse;
41
42 namespace {
43
44 178 void SetRpcDeadline(grpc::ClientContext* context, int timeout_ms = 15000) {
45
1/2
✓ Branch 1 taken 178 times.
✗ Branch 2 not taken.
 178 context->set_deadline(
46
1/2
✓ Branch 3 taken 178 times.
✗ Branch 4 not taken.
 178 std::chrono::system_clock::now() + std::chrono::milliseconds(timeout_ms));
47 178 }
48
49 ✗ int BuildUpdateBlocksFromFlat(
50 const base::ConstArray<uint64_t>& keys,
51 const float* grads,
52 int64_t num_rows,
53 int64_t embedding_dim,
54 ParameterCompressor* compressor) {
55 ✗ if (grads == nullptr) {
56 ✗ LOG(ERROR) << "UpdateParameterFlat grads pointer is null";
57 ✗ return -1;
58 }
59 ✗ if (num_rows < 0 || embedding_dim <= 0) {
60 ✗ LOG(ERROR) << "UpdateParameterFlat invalid shape: rows=" << num_rows
61 ✗ << " dim=" << embedding_dim;
62 ✗ return -1;
63 }
64 ✗ if (keys.Size() != static_cast<size_t>(num_rows)) {
65 ✗ LOG(ERROR) << "UpdateParameterFlat keys/grads size mismatch: "
66 ✗ << keys.Size() << " vs " << num_rows;
67 ✗ return -1;
68 }
69
70 ✗ for (int64_t i = 0; i < num_rows; ++i) {
71 ✗ ParameterPack pack;
72 ✗ pack.key = keys[static_cast<size_t>(i)];
73 ✗ pack.dim = embedding_dim;
74 ✗ pack.emb_data = grads + i * embedding_dim;
75 ✗ compressor->AddItem(pack, nullptr);
76 }
77 ✗ return 0;
78 }
79
80 } // namespace
81
82 DEFINE_int32(get_parameter_threads, 4, "get clients per shard");
83 DEFINE_bool(parameter_client_random_init, false, "");
84
85 // New constructor that takes JSON config.
86 /*
87 Example: load config from file
88 std::ifstream config_file(FLAGS_config_path);
89 nlohmann::json ex;
90 config_file >> ex;
91 json client_config = ex["client"];
92
93 */
94 34 GRPCParameterClient::GRPCParameterClient(json config)
95
1/2
✓ Branch 2 taken 34 times.
✗ Branch 3 not taken.
 34 : recstore::BasePSClient(config) {
96 // Extract fields from JSON config
97
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 host_ = config.value("host", "localhost");
98
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 port_ = config.value("port", 15000);
99
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 shard_ = config.value("shard", 0);
100 34 nr_clients_ = FLAGS_get_parameter_threads;
101 34 Initialize();
102
103
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 grpc::ChannelArguments args;
104
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 args.SetMaxReceiveMessageSize(-1);
105
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 34 args.SetMaxSendMessageSize(-1);
106
107 34 channel_ = grpc::CreateCustomChannel(
108
2/4
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 34 times.
✗ Branch 5 not taken.
 102 fmt::format("{}:{}", host_, port_),
109
1/2
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
 68 grpc::InsecureChannelCredentials(),
110 34 args);
111
2/4
✓ Branch 1 taken 34 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 34 times.
✗ Branch 5 not taken.
 34 auto* raw_cq = new grpc::CompletionQueue();
112 34 cq_.reset(raw_cq);
113
114
2/2
✓ Branch 0 taken 136 times.
✓ Branch 1 taken 34 times.
 170 for (int i = 0; i < nr_clients_; i++) {
115
1/2
✓ Branch 2 taken 136 times.
✗ Branch 3 not taken.
 136 stubs_.push_back(nullptr);
116
1/2
✓ Branch 3 taken 136 times.
✗ Branch 4 not taken.
 136 stubs_[i] = recstoreps::ParameterService::NewStub(channel_);
117
4/8
✓ Branch 1 taken 136 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 136 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 136 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 136 times.
✗ Branch 11 not taken.
 136 LOG(INFO) << "Init PS Client Shard " << i;
118 }
119 34 }
120
121 // Legacy constructor for backward compatibility
122 2 GRPCParameterClient::GRPCParameterClient(
123 2 const std::string& host, int port, int shard)
124 : recstore::BasePSClient(
125
14/28
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 9 taken 2 times.
✗ Branch 10 not taken.
✓ Branch 12 taken 2 times.
✗ Branch 13 not taken.
✓ Branch 16 taken 2 times.
✗ Branch 17 not taken.
✓ Branch 19 taken 2 times.
✗ Branch 20 not taken.
✓ Branch 22 taken 6 times.
✓ Branch 23 taken 2 times.
✓ Branch 25 taken 4 times.
✓ Branch 26 taken 2 times.
✓ Branch 28 taken 4 times.
✓ Branch 29 taken 2 times.
✓ Branch 31 taken 4 times.
✓ Branch 32 taken 2 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.
 22 json{{"host", host}, {"port", port}, {"shard", shard}}),
126 2 host_(host),
127 2 port_(port),
128 2 shard_(shard),
129
2/4
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
 6 nr_clients_(FLAGS_get_parameter_threads) {
130 2 Initialize();
131
132
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 grpc::ChannelArguments args;
133
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 args.SetMaxReceiveMessageSize(-1);
134
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 args.SetMaxSendMessageSize(-1);
135
136 2 channel_ = grpc::CreateCustomChannel(
137
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 4 fmt::format("{}:{}", host, port),
138
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 4 grpc::InsecureChannelCredentials(),
139 2 args);
140
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 auto* raw_cq = new grpc::CompletionQueue();
141 2 cq_.reset(raw_cq);
142
143
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 2 times.
 10 for (int i = 0; i < nr_clients_; i++) {
144
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 stubs_.push_back(nullptr);
145
1/2
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
 8 stubs_[i] = recstoreps::ParameterService::NewStub(channel_);
146
4/8
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 8 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 8 times.
✗ Branch 11 not taken.
 8 LOG(INFO) << "Init PS Client Shard " << i;
147 }
148 2 }
149
150 ✗ int GRPCParameterClient::GetParameter(const base::ConstArray<uint64_t>& keys,
151 float* values) {
152 #ifdef ENABLE_PERF_REPORT
153 auto start_time = std::chrono::high_resolution_clock::now();
154 #endif
155
156 ✗ if (FLAGS_parameter_client_random_init) {
157 ✗ CHECK(0) << "todo implement";
158 return true;
159 }
160
161 ✗ get_param_key_sizes_.clear();
162 ✗ get_param_status_.clear();
163 ✗ get_param_requests_.clear();
164 ✗ get_param_responses_.clear();
165 ✗ get_param_resonse_readers_.clear();
166 ✗ get_param_contexts_.clear();
167
168 int request_num =
169 ✗ (keys.Size() + MAX_PARAMETER_BATCH - 1) / MAX_PARAMETER_BATCH;
170 ✗ get_param_status_.resize(request_num);
171 ✗ get_param_requests_.resize(request_num);
172 ✗ get_param_responses_.resize(request_num);
173 ✗ get_param_contexts_.resize(request_num);
174
175 ✗ for (int start = 0, index = 0; start < keys.Size();
176 ✗ start += MAX_PARAMETER_BATCH, ++index) {
177 ✗ int key_size = std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH);
178 ✗ get_param_key_sizes_.emplace_back(key_size);
179 ✗ auto& status = get_param_status_[index];
180 ✗ auto& request = get_param_requests_[index];
181 ✗ auto& response = get_param_responses_[index];
182 ✗ request.set_keys(reinterpret_cast<const char*>(&keys[start]),
183 ✗ sizeof(uint64_t) * key_size);
184 // rpc
185 // grpc::ClientContext context;
186 ✗ if (!get_param_contexts_[index]) {
187 ✗ get_param_contexts_[index] = std::make_unique<grpc::ClientContext>();
188 }
189 std::unique_ptr<ClientAsyncResponseReader<GetParameterResponse>> rpc =
190 ✗ stubs_[0]->AsyncGetParameter(
191 ✗ get_param_contexts_[index].get(), request, cq_.get());
192 ✗ rpc->Finish(&response, &status, reinterpret_cast<void*>(index));
193 ✗ }
194 ✗ int get = 0;
195 ✗ while (get != request_num) {
196 void* got_tag;
197 ✗ bool ok = false;
198 ✗ cq_->Next(&got_tag, &ok);
199 ✗ if (!ok) {
200 ✗ LOG(ERROR) << "error";
201 }
202 ✗ get++;
203 }
204 #ifdef ENABLE_PERF_REPORT
205 auto after_rpc_time = std::chrono::high_resolution_clock::now();
206 auto rpc_duration =
207 std::chrono::duration_cast<std::chrono::microseconds>(
208 after_rpc_time - start_time)
209 .count();
210 double start_us_for_rpc =
211 std::chrono::duration_cast<std::chrono::microseconds>(
212 start_time.time_since_epoch())
213 .count();
214 std::string report_id_for_rpc =
215 "grpc_client::GetParameter|" +
216 std::to_string(static_cast<uint64_t>(start_us_for_rpc));
217 report("embread_stages",
218 report_id_for_rpc.c_str(),
219 "rpc_duration_us",
220 static_cast<double>(rpc_duration));
221 #endif
222 ✗ size_t get_embedding_acc = 0;
223 ✗ int old_dimension = -1;
224
225 ✗ for (int i = 0; i < get_param_responses_.size(); ++i) {
226 ✗ auto& response = get_param_responses_[i];
227 ✗ int key_size = get_param_key_sizes_[i];
228 auto parameters = reinterpret_cast<const ParameterCompressReader*>(
229 ✗ response.parameter_value().data());
230
231 ✗ if (parameters->size != key_size) {
232 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
233 ✗ << key_size;
234 ✗ return false;
235 }
236
237 ✗ for (int index = 0; index < parameters->item_size(); ++index) {
238 ✗ auto item = parameters->item(index);
239 ✗ if (item->dim != 0) {
240 ✗ if (old_dimension == -1)
241 ✗ old_dimension = item->dim;
242 ✗ CHECK_EQ(item->dim, old_dimension);
243 ✗ std::copy_n(
244 ✗ item->embedding, item->dim, values + item->dim * get_embedding_acc);
245 } else {
246 ✗ RECSTORE_LOG_EVERY_MS(ERROR, 2000)
247 ✗ << "error; not find key " << keys[get_embedding_acc] << " in ps";
248 }
249 ✗ get_embedding_acc++;
250 }
251 }
252
253 #ifdef ENABLE_PERF_REPORT
254 auto end_time = std::chrono::high_resolution_clock::now();
255 auto duration =
256 std::chrono::duration_cast<std::chrono::microseconds>(
257 end_time - start_time)
258 .count();
259 double start_us =
260 std::chrono::duration_cast<std::chrono::microseconds>(
261 start_time.time_since_epoch())
262 .count();
263
264 auto deserialize_duration =
265 std::chrono::duration_cast<std::chrono::microseconds>(
266 end_time - after_rpc_time)
267 .count();
268
269 report("embread_stages",
270 "grpc_client::GetParameter",
271 "deserialize_duration_us",
272 static_cast<double>(deserialize_duration));
273
274 report("embread_stages",
275 "grpc_client::GetParameter",
276 "duration_us",
277 static_cast<double>(duration));
278
279 report("embread_stages",
280 "grpc_client::GetParameter",
281 "request_size",
282 static_cast<double>(keys.Size()));
283
284 FlameGraphData grpc_client_data = {
285 "grpc_ps_client::GetParameter",
286 start_us,
287 1, // level
288 static_cast<double>(duration),
289 static_cast<double>(duration)};
290
291 std::string unique_id = "embread_debug";
292 report_flame_graph("emb_read_flame_map", unique_id.c_str(), grpc_client_data);
293 #endif
294
295 ✗ return true;
296 }
297
298 80 int GRPCParameterClient::GetParameter(const base::ConstArray<uint64_t>& keys,
299 std::vector<std::vector<float>>* values) {
300 #ifdef ENABLE_PERF_REPORT
301 auto start_time = std::chrono::high_resolution_clock::now();
302 #endif
303
304
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
 80 if (FLAGS_parameter_client_random_init) {
305 ✗ values->clear();
306 ✗ values->reserve(keys.Size());
307 ✗ for (size_t i = 0; i < keys.Size(); i++)
308 ✗ values->emplace_back(std::vector<float>(128, 0.1));
309
310 ✗ return true;
311 }
312
313 80 values->clear();
314 80 get_param_key_sizes_.clear();
315 80 get_param_status_.clear();
316 80 get_param_requests_.clear();
317 80 get_param_responses_.clear();
318 80 get_param_resonse_readers_.clear();
319 80 get_param_contexts_.clear();
320
321 80 values->reserve(keys.Size());
322
323 int request_num =
324 80 (keys.Size() + MAX_PARAMETER_BATCH - 1) / MAX_PARAMETER_BATCH;
325
326 80 get_param_status_.resize(request_num);
327 80 get_param_requests_.resize(request_num);
328 80 get_param_responses_.resize(request_num);
329 80 get_param_contexts_.resize(request_num);
330
331
2/2
✓ Branch 1 taken 80 times.
✓ Branch 2 taken 80 times.
 160 for (int start = 0, index = 0; start < keys.Size();
332 80 start += MAX_PARAMETER_BATCH, ++index) {
333 80 int key_size = std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH);
334
1/2
✓ Branch 1 taken 80 times.
✗ Branch 2 not taken.
 80 get_param_key_sizes_.emplace_back(key_size);
335 80 auto& status = get_param_status_[index];
336 80 auto& request = get_param_requests_[index];
337 80 auto& response = get_param_responses_[index];
338 80 request.set_keys(reinterpret_cast<const char*>(&keys[start]),
339 80 sizeof(uint64_t) * key_size);
340 // rpc
341 // grpc::ClientContext context;
342
1/2
✓ Branch 2 taken 80 times.
✗ Branch 3 not taken.
 80 if (!get_param_contexts_[index]) {
343
1/2
✓ Branch 1 taken 80 times.
✗ Branch 2 not taken.
 80 get_param_contexts_[index] = std::make_unique<grpc::ClientContext>();
344 }
345
2/4
✓ Branch 5 taken 80 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 80 times.
✗ Branch 9 not taken.
 160 get_param_resonse_readers_.emplace_back(stubs_[0]->AsyncGetParameter(
346 80 get_param_contexts_[index].get(), request, cq_.get()));
347 80 auto& rpc = get_param_resonse_readers_.back();
348 // GetParameter(&context, request, &response);
349
1/2
✓ Branch 2 taken 80 times.
✗ Branch 3 not taken.
 80 rpc->Finish(&response, &status, reinterpret_cast<void*>(index));
350 }
351
352 80 int get = 0;
353
2/2
✓ Branch 0 taken 80 times.
✓ Branch 1 taken 80 times.
 160 while (get != request_num) {
354 void* got_tag;
355 80 bool ok = false;
356
1/2
✓ Branch 2 taken 80 times.
✗ Branch 3 not taken.
 80 cq_->Next(&got_tag, &ok);
357
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
 80 if (unlikely(!ok)) {
358 ✗ LOG(ERROR) << "error";
359 }
360 80 get++;
361 }
362
363 #ifdef ENABLE_PERF_REPORT
364 auto after_rpc_time = std::chrono::high_resolution_clock::now();
365 auto rpc_duration =
366 std::chrono::duration_cast<std::chrono::microseconds>(
367 after_rpc_time - start_time)
368 .count();
369 double start_us_for_rpc =
370 std::chrono::duration_cast<std::chrono::microseconds>(
371 start_time.time_since_epoch())
372 .count();
373 std::string report_id_for_rpc =
374 "grpc_client::GetParameter|" +
375 std::to_string(static_cast<uint64_t>(start_us_for_rpc));
376 report("embread_stages",
377 report_id_for_rpc.c_str(),
378 "rpc_duration_us",
379 static_cast<double>(rpc_duration));
380 #endif
381
382
2/2
✓ Branch 1 taken 80 times.
✓ Branch 2 taken 80 times.
 160 for (int i = 0; i < get_param_responses_.size(); ++i) {
383 80 auto& response = get_param_responses_[i];
384 80 int key_size = get_param_key_sizes_[i];
385 auto parameters = reinterpret_cast<const ParameterCompressReader*>(
386 80 response.parameter_value().data());
387
388
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 80 times.
 80 if (unlikely(parameters->size != key_size)) {
389 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
390 ✗ << key_size;
391 ✗ return false;
392 }
393
394
2/2
✓ Branch 1 taken 362 times.
✓ Branch 2 taken 80 times.
 442 for (int index = 0; index < parameters->item_size(); ++index) {
395 362 auto item = parameters->item(index);
396
2/2
✓ Branch 0 taken 356 times.
✓ Branch 1 taken 6 times.
 362 if (item->dim != 0) {
397
1/2
✓ Branch 1 taken 356 times.
✗ Branch 2 not taken.
 356 values->emplace_back(
398
1/2
✓ Branch 2 taken 356 times.
✗ Branch 3 not taken.
 712 std::vector<float>(item->embedding, item->embedding + item->dim));
399 } else {
400
2/4
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 6 times.
✗ Branch 6 not taken.
 6 values->emplace_back(std::vector<float>(0));
401 }
402 }
403 }
404
405 #ifdef ENABLE_PERF_REPORT
406 auto end_time = std::chrono::high_resolution_clock::now();
407 auto duration =
408 std::chrono::duration_cast<std::chrono::microseconds>(
409 end_time - start_time)
410 .count();
411 double start_us =
412 std::chrono::duration_cast<std::chrono::microseconds>(
413 start_time.time_since_epoch())
414 .count();
415
416 auto deserialize_duration =
417 std::chrono::duration_cast<std::chrono::microseconds>(
418 end_time - after_rpc_time)
419 .count();
420
421 report("embread_stages",
422 "grpc_client::GetParameter",
423 "deserialize_duration_us",
424 static_cast<double>(deserialize_duration));
425
426 report("embread_stages",
427 "grpc_client::GetParameter",
428 "duration_us",
429 static_cast<double>(duration));
430
431 report("embread_stages",
432 "grpc_client::GetParameter",
433 "request_size",
434 static_cast<double>(keys.Size()));
435
436 FlameGraphData grpc_client_data = {
437 "grpc_ps_client::GetParameter",
438 start_us,
439 1, // level
440 static_cast<double>(duration),
441 static_cast<double>(duration)};
442
443 std::string unique_id = "embread_debug";
444 report_flame_graph("emb_read_flame_map", unique_id.c_str(), grpc_client_data);
445 #endif
446
447 80 return true;
448 }
449
450 // return prefetch id
451 uint64_t
452 112 GRPCParameterClient::PrefetchParameter(const base::ConstArray<uint64_t>& keys) {
453 int request_num =
454 112 (keys.Size() + MAX_PARAMETER_BATCH - 1) / MAX_PARAMETER_BATCH;
455
456
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 struct PrefetchBatch pb(request_num);
457
458
2/2
✓ Branch 1 taken 112 times.
✓ Branch 2 taken 112 times.
 224 for (int start = 0, index = 0; start < keys.Size();
459 112 start += MAX_PARAMETER_BATCH, ++index) {
460 112 int key_size = std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH);
461 112 pb.key_sizes_[index] = key_size;
462 112 auto& status = pb.status_[index];
463
1/2
✓ Branch 2 taken 112 times.
✗ Branch 3 not taken.
 112 if (!pb.contexts_[index]) {
464
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 pb.contexts_[index] = std::make_unique<grpc::ClientContext>();
465 }
466 112 auto& request = pb.requests_[index];
467 112 auto& response = pb.responses_[index];
468 112 request.set_keys(reinterpret_cast<const char*>(&keys[start]),
469 112 sizeof(uint64_t) * key_size);
470 // rpc
471 // grpc::ClientContext context;
472
2/4
✓ Branch 5 taken 112 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 112 times.
✗ Branch 9 not taken.
 224 pb.response_readers_.emplace_back(stubs_[0]->AsyncGetParameter(
473 112 pb.contexts_[index].get(), request, pb.cqs_.get()));
474 112 auto& rpc = pb.response_readers_.back();
475 // GetParameter(&context, request, &response);
476
1/2
✓ Branch 2 taken 112 times.
✗ Branch 3 not taken.
 112 rpc->Finish(&response, &status, reinterpret_cast<void*>(index));
477 }
478 112 uint64_t prefetch_id = 0;
479 {
480
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 std::lock_guard<std::mutex> lk(prefetch_mu_);
481 112 prefetch_id = next_prefetch_id_++;
482
1/2
✓ Branch 2 taken 112 times.
✗ Branch 3 not taken.
 112 prefetch_batches_.emplace(prefetch_id, std::move(pb));
483 112 }
484
485 112 return prefetch_id;
486 112 }
487
488 4 bool GRPCParameterClient::IsPrefetchDone(uint64_t prefetch_id) {
489
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 std::lock_guard<std::mutex> lk(prefetch_mu_);
490
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 auto it = prefetch_batches_.find(prefetch_id);
491
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (it == prefetch_batches_.end()) {
492 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
493 ✗ return false;
494 }
495 4 auto& pb = it->second;
496 4 int request_num = pb.batch_size_;
497 4 int get = 0;
498
499
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (pb.completed_count_ == pb.batch_size_) {
500 4 return true;
501 }
502
503 ✗ void* got_tag = nullptr;
504 ✗ bool ok = false;
505 auto deadline =
506 ✗ std::chrono::system_clock::now() + std::chrono::milliseconds(0);
507 for (;;) {
508 ✗ auto status = pb.cqs_->AsyncNext(&got_tag, &ok, deadline);
509 ✗ if (status == grpc::CompletionQueue::NextStatus::GOT_EVENT) {
510 ✗ if (unlikely(!ok)) {
511 ✗ LOG(ERROR) << "CompletionQueue returned not ok for prefetch";
512 }
513 ✗ pb.completed_count_++;
514 ✗ if (pb.completed_count_ == pb.batch_size_)
515 ✗ break;
516 deadline =
517 ✗ std::chrono::system_clock::now() + std::chrono::milliseconds(0);
518 ✗ continue;
519 ✗ } else if (status == grpc::CompletionQueue::NextStatus::TIMEOUT) {
520 ✗ break;
521 } else {
522 ✗ LOG(ERROR) << "CompletionQueue shutdown during prefetch";
523 ✗ break;
524 }
525 ✗ }
526 ✗ return (pb.completed_count_ == pb.batch_size_);
527 4 }
528
529 212 void GRPCParameterClient::WaitForPrefetch(uint64_t prefetch_id) {
530
1/2
✓ Branch 1 taken 212 times.
✗ Branch 2 not taken.
 212 std::lock_guard<std::mutex> lk(prefetch_mu_);
531
1/2
✓ Branch 1 taken 212 times.
✗ Branch 2 not taken.
 212 auto it = prefetch_batches_.find(prefetch_id);
532
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 212 times.
 212 if (it == prefetch_batches_.end()) {
533 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
534 ✗ return;
535 }
536 212 auto& pb = it->second;
537 212 void* got_tag = nullptr;
538 212 bool ok = false;
539 212 int idle_rounds = 0;
540 212 constexpr auto kPollInterval = std::chrono::milliseconds(200);
541 212 constexpr int kMaxIdleRounds = 150; // 30s
542
2/2
✓ Branch 0 taken 112 times.
✓ Branch 1 taken 212 times.
 324 while (pb.completed_count_ < pb.batch_size_) {
543
1/2
✓ Branch 2 taken 112 times.
✗ Branch 3 not taken.
 112 auto deadline = std::chrono::system_clock::now() + kPollInterval;
544
1/2
✓ Branch 2 taken 112 times.
✗ Branch 3 not taken.
 112 auto status = pb.cqs_->AsyncNext(&got_tag, &ok, deadline);
545
1/2
✓ Branch 0 taken 112 times.
✗ Branch 1 not taken.
 112 if (status == grpc::CompletionQueue::NextStatus::GOT_EVENT) {
546 112 idle_rounds = 0;
547
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 112 times.
 112 if (unlikely(!ok)) {
548 ✗ LOG(ERROR) << "CompletionQueue returned not ok for prefetch";
549 }
550 112 pb.completed_count_++;
551 112 continue;
552 }
553 ✗ if (status == grpc::CompletionQueue::NextStatus::TIMEOUT) {
554 ✗ idle_rounds++;
555 ✗ if (idle_rounds >= kMaxIdleRounds) {
556 ✗ LOG(ERROR) << "WaitForPrefetch timed out for prefetch_id "
557 ✗ << prefetch_id << ", completed " << pb.completed_count_
558 ✗ << "/" << pb.batch_size_;
559 ✗ break;
560 }
561 ✗ continue;
562 }
563 ✗ if (status == grpc::CompletionQueue::NextStatus::SHUTDOWN) {
564 ✗ LOG(ERROR) << "CompletionQueue shutdown while waiting prefetch";
565 ✗ break;
566 }
567 }
568
1/2
✓ Branch 1 taken 212 times.
✗ Branch 2 not taken.
 212 }
569
570 112 bool GRPCParameterClient::GetPrefetchResult(
571 uint64_t prefetch_id, std::vector<std::vector<float>>* values) {
572
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 std::lock_guard<std::mutex> lk(prefetch_mu_);
573
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 auto it = prefetch_batches_.find(prefetch_id);
574
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 112 times.
 112 if (it == prefetch_batches_.end()) {
575 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
576 ✗ return false;
577 }
578 112 auto& pb = it->second;
579 112 int request_num = pb.batch_size_;
580
581 112 values->clear();
582 112 int keys_size = 0;
583
2/2
✓ Branch 4 taken 112 times.
✓ Branch 5 taken 112 times.
 224 for (const auto& size : pb.key_sizes_) {
584 112 keys_size += size;
585 }
586
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 values->reserve(keys_size);
587
588
2/2
✓ Branch 0 taken 112 times.
✓ Branch 1 taken 112 times.
 224 for (int i = 0; i < request_num; ++i) {
589 112 auto& response = pb.responses_[i];
590 112 int key_size = pb.key_sizes_[i];
591 auto parameters = reinterpret_cast<const ParameterCompressReader*>(
592
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 response.parameter_value().data());
593
594
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 112 times.
 112 if (unlikely(parameters->size != key_size)) {
595 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
596 ✗ << key_size;
597 ✗ return false;
598 }
599
600
2/2
✓ Branch 1 taken 520 times.
✓ Branch 2 taken 112 times.
 632 for (int index = 0; index < parameters->item_size(); ++index) {
601
1/2
✓ Branch 1 taken 520 times.
✗ Branch 2 not taken.
 520 auto item = parameters->item(index);
602
1/2
✓ Branch 0 taken 520 times.
✗ Branch 1 not taken.
 520 if (item->dim != 0) {
603
1/2
✓ Branch 1 taken 520 times.
✗ Branch 2 not taken.
 520 values->emplace_back(
604
1/2
✓ Branch 2 taken 520 times.
✗ Branch 3 not taken.
 1040 std::vector<float>(item->embedding, item->embedding + item->dim));
605 } else {
606 ✗ values->emplace_back(std::vector<float>(0));
607 }
608 }
609 }
610
611 112 return true;
612 112 }
613
614 ✗ bool GRPCParameterClient::GetPrefetchResultFlat(
615 uint64_t prefetch_id,
616 std::vector<float>* values,
617 int64_t* num_rows,
618 int64_t embedding_dim) {
619 ✗ std::lock_guard<std::mutex> lk(prefetch_mu_);
620 ✗ auto it = prefetch_batches_.find(prefetch_id);
621 ✗ if (it == prefetch_batches_.end()) {
622 ✗ LOG(ERROR) << "Invalid prefetch_id: " << prefetch_id;
623 ✗ return false;
624 }
625 ✗ if (values == nullptr || num_rows == nullptr) {
626 ✗ LOG(ERROR) << "GetPrefetchResultFlat output pointer is null";
627 ✗ return false;
628 }
629
630 ✗ auto& pb = it->second;
631 ✗ int request_num = pb.batch_size_;
632 ✗ int total_keys = 0;
633 ✗ for (const auto& size : pb.key_sizes_) {
634 ✗ total_keys += size;
635 }
636
637 ✗ *num_rows = static_cast<int64_t>(total_keys);
638 ✗ values->assign(
639 ✗ static_cast<size_t>(*num_rows) * static_cast<size_t>(embedding_dim),
640 ✗ 0.0f);
641
642 ✗ size_t row_offset = 0;
643 ✗ for (int i = 0; i < request_num; ++i) {
644 ✗ auto& response = pb.responses_[i];
645 ✗ int key_size = pb.key_sizes_[i];
646 auto parameters = reinterpret_cast<const ParameterCompressReader*>(
647 ✗ response.parameter_value().data());
648
649 ✗ if (unlikely(parameters->size != key_size)) {
650 ✗ LOG(ERROR) << "GetParameter error: " << parameters->size << " vs "
651 ✗ << key_size;
652 ✗ return false;
653 }
654
655 ✗ for (int index = 0; index < parameters->item_size();
656 ✗ ++index, ++row_offset) {
657 ✗ auto item = parameters->item(index);
658 ✗ if (item->dim != 0) {
659 const int64_t copy_d =
660 ✗ std::min<int64_t>(embedding_dim, static_cast<int64_t>(item->dim));
661 ✗ std::memcpy(values->data() + row_offset * embedding_dim,
662 ✗ item->embedding,
663 ✗ static_cast<size_t>(copy_d) * sizeof(float));
664 }
665 }
666 }
667
668 ✗ return true;
669 ✗ }
670
671 26 bool GRPCParameterClient::ClearPS() {
672
1/2
✓ Branch 1 taken 26 times.
✗ Branch 2 not taken.
 26 CommandRequest request;
673
1/2
✓ Branch 1 taken 26 times.
✗ Branch 2 not taken.
 26 CommandResponse response;
674
1/2
✓ Branch 1 taken 26 times.
✗ Branch 2 not taken.
 26 request.set_command(PSCommand::CLEAR_PS);
675
1/2
✓ Branch 1 taken 26 times.
✗ Branch 2 not taken.
 26 grpc::ClientContext context;
676
1/2
✓ Branch 1 taken 26 times.
✗ Branch 2 not taken.
 26 SetRpcDeadline(&context);
677
1/2
✓ Branch 3 taken 26 times.
✗ Branch 4 not taken.
 26 grpc::Status status = stubs_[0]->Command(&context, request, &response);
678
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 26 times.
 26 if (!status.ok()) {
679 ✗ LOG(ERROR) << "gRPC ClearPS failed: " << status.error_code() << " "
680 ✗ << status.error_message();
681 }
682 52 return status.ok();
683 26 }
684
685 // Read n bytes from the server. The server does not access storage;
686 // it generates data randomly instead.
687 2 bool GRPCParameterClient::LoadFakeData(int64_t n) {
688
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 CommandRequest request;
689
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 CommandResponse response;
690
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 request.set_command(PSCommand::LOAD_FAKE_DATA);
691
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 request.add_arg1(&n, sizeof(int64_t));
692
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 grpc::ClientContext context;
693
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 SetRpcDeadline(&context);
694
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 2 grpc::Status status = stubs_[0]->Command(&context, request, &response);
695
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (!status.ok()) {
696 ✗ LOG(ERROR) << "gRPC LoadFakeData failed: " << status.error_code() << " "
697 ✗ << status.error_message();
698 ✗ return false;
699 }
700
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 2 if (response.reply().size() != static_cast<size_t>(n)) {
701 ✗ LOG(ERROR) << "gRPC LoadFakeData reply size mismatch: expected " << n
702 ✗ << ", got " << response.reply().size();
703 ✗ return false;
704 }
705 2 return true;
706 2 }
707
708 // Write n bytes(random generated) into the server
709 2 bool GRPCParameterClient::DumpFakeData(int64_t n) {
710
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 CommandRequest request;
711
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 CommandResponse response;
712
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 request.set_command(PSCommand::DUMP_FAKE_DATA);
713
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 request.add_arg1(&n, sizeof(int64_t));
714
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 grpc::ClientContext context;
715
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 SetRpcDeadline(&context);
716
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 2 grpc::Status status = stubs_[0]->Command(&context, request, &response);
717
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (!status.ok()) {
718 ✗ LOG(ERROR) << "gRPC DumpFakeData failed: " << status.error_code() << " "
719 ✗ << status.error_message();
720 ✗ return false;
721 }
722
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
 2 if (response.reply() != "ok") {
723 ✗ LOG(ERROR) << "gRPC DumpFakeData unexpected reply: " << response.reply();
724 ✗ return false;
725 }
726 2 return true;
727 2 }
728
729 ✗ bool GRPCParameterClient::LoadCkpt(
730 const std::vector<std::string>& model_config_path,
731 const std::vector<std::string>& emb_file_path) {
732 ✗ CommandRequest request;
733 ✗ CommandResponse response;
734 ✗ request.set_command(PSCommand::RELOAD_PS);
735
736 ✗ for (auto& each : model_config_path) {
737 ✗ request.add_arg1(each);
738 }
739 ✗ for (auto& each : emb_file_path) {
740 ✗ request.add_arg2(each);
741 }
742 ✗ grpc::ClientContext context;
743 ✗ SetRpcDeadline(&context, 30000);
744 ✗ grpc::Status status = stubs_[0]->Command(&context, request, &response);
745 ✗ return status.ok();
746 ✗ }
747
748 ✗ bool GRPCParameterClient::SaveCheckpoint(const std::string& path,
749 const std::string& metadata) {
750 ✗ CommandRequest request;
751 ✗ CommandResponse response;
752 ✗ request.set_command(PSCommand::SAVE_CHECKPOINT);
753 ✗ request.add_arg1(path);
754 ✗ request.add_arg2(metadata);
755 ✗ grpc::ClientContext context;
756 ✗ SetRpcDeadline(&context, 600000);
757 ✗ const grpc::Status status = stubs_[0]->Command(&context, request, &response);
758 ✗ if (!status.ok()) {
759 ✗ LOG(ERROR) << "gRPC checkpoint save failed: " << status.error_code() << " "
760 ✗ << status.error_message();
761 }
762 ✗ return status.ok();
763 ✗ }
764
765 ✗ bool GRPCParameterClient::LoadCheckpoint(const std::string& path,
766 const std::string& metadata) {
767 ✗ CommandRequest request;
768 ✗ CommandResponse response;
769 ✗ request.set_command(PSCommand::LOAD_CHECKPOINT);
770 ✗ request.add_arg1(path);
771 ✗ request.add_arg2(metadata);
772 ✗ grpc::ClientContext context;
773 ✗ SetRpcDeadline(&context, 600000);
774 ✗ const grpc::Status status = stubs_[0]->Command(&context, request, &response);
775 ✗ if (!status.ok()) {
776 ✗ LOG(ERROR) << "gRPC checkpoint load failed: " << status.error_code() << " "
777 ✗ << status.error_message();
778 }
779 ✗ return status.ok();
780 ✗ }
781
782 146 bool GRPCParameterClient::PutParameter(
783 const std::vector<uint64_t>& keys,
784 const std::vector<std::vector<float>>& values) {
785
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 146 times.
 146 if (keys.size() != values.size()) {
786 ✗ LOG(ERROR) << "PutParameter keys/values size mismatch: " << keys.size()
787 ✗ << " vs " << values.size();
788 ✗ return false;
789 }
790
2/2
✓ Branch 1 taken 146 times.
✓ Branch 2 taken 146 times.
 292 for (int start = 0, index = 0; start < keys.size();
791 146 start += MAX_PARAMETER_BATCH, ++index) {
792 146 int key_size = std::min((int)(keys.size() - start), MAX_PARAMETER_BATCH);
793
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 PutParameterRequest request;
794
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 PutParameterResponse response;
795
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 ParameterCompressor compressor;
796 146 std::vector<std::string> blocks;
797
2/2
✓ Branch 0 taken 1244 times.
✓ Branch 1 taken 146 times.
 1390 for (int i = start; i < start + key_size; i++) {
798 1244 auto each_key = keys[i];
799 1244 auto& embedding = values[i];
800 1244 ParameterPack parameter_pack;
801 1244 parameter_pack.key = each_key;
802 1244 parameter_pack.dim = embedding.size();
803 1244 parameter_pack.emb_data = embedding.data();
804
1/2
✓ Branch 1 taken 1244 times.
✗ Branch 2 not taken.
 1244 compressor.AddItem(parameter_pack, &blocks);
805 }
806
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 compressor.ToBlock(&blocks);
807
2/8
✓ Branch 4 taken 146 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 146 times.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 146 CHECK_EQ(blocks.size(), 1);
808
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 request.mutable_parameter_value()->swap(blocks[0]);
809
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 grpc::ClientContext context;
810
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 SetRpcDeadline(&context);
811
1/2
✓ Branch 3 taken 146 times.
✗ Branch 4 not taken.
 146 grpc::Status status = stubs_[0]->PutParameter(&context, request, &response);
812
1/2
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
 146 if (status.ok()) {
813 146 continue;
814 } else {
815 ✗ std::cout << status.error_code() << ": " << status.error_message()
816 ✗ << std::endl;
817 ✗ return false;
818 }
819
6/12
✓ Branch 1 taken 146 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 146 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 146 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 146 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 146 times.
✗ Branch 14 not taken.
✓ Branch 16 taken 146 times.
✗ Branch 17 not taken.
 876 }
820 146 return true;
821 }
822
823 2 int GRPCParameterClient::UpdateParameter(
824 const std::string& table_name,
825 const base::ConstArray<uint64_t>& keys,
826 const std::vector<std::vector<float>>* grads) {
827 #ifdef ENABLE_PERF_REPORT
828 auto start_time = std::chrono::high_resolution_clock::now();
829 const uint64_t trace_id = recstore::g_trace_id;
830 #endif
831
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (grads == nullptr) {
832 ✗ LOG(ERROR) << "UpdateParameter grads pointer is null";
833 ✗ return -1;
834 }
835
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (keys.Size() != grads->size()) {
836 ✗ LOG(ERROR) << "UpdateParameter keys/grads size mismatch: " << keys.Size()
837 ✗ << " vs " << grads->size();
838 ✗ return -1;
839 }
840
841
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 ParameterCompressor compressor;
842
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
 6 for (size_t i = 0; i < keys.Size(); ++i) {
843 4 ParameterPack pack;
844 4 pack.key = keys[i];
845
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 pack.dim = grads->at(i).size();
846
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 pack.emb_data = grads->at(i).data();
847
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 compressor.AddItem(pack, nullptr);
848 }
849 #ifdef ENABLE_PERF_REPORT
850 auto serialize_done_time = std::chrono::high_resolution_clock::now();
851 #endif
852
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (keys.Size() == 0) {
853 ✗ LOG(WARNING) << "UpdateParameter no gradients to send";
854 ✗ return 0;
855 }
856
857
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 UpdateParameterRequest request;
858
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 UpdateParameterResponse response;
859 request.set_table_name(table_name);
860
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 compressor.ToBlock(request.mutable_gradients());
861
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 2 if (request.gradients().empty()) {
862 ✗ LOG(WARNING) << "UpdateParameter no serialized gradients payload";
863 ✗ return 0;
864 }
865
866
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 grpc::ClientContext context;
867
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 SetRpcDeadline(&context);
868 #ifdef ENABLE_PERF_REPORT
869 if (trace_id != 0) {
870 context.AddMetadata("x-recstore-trace-id", std::to_string(trace_id));
871 }
872 auto rpc_start_time = std::chrono::high_resolution_clock::now();
873 #endif
874 grpc::Status status =
875
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 2 stubs_[0]->UpdateParameter(&context, request, &response);
876 #ifdef ENABLE_PERF_REPORT
877 auto end_time = std::chrono::high_resolution_clock::now();
878 auto serialize_duration =
879 std::chrono::duration_cast<std::chrono::microseconds>(
880 serialize_done_time - start_time)
881 .count();
882 auto rpc_duration =
883 std::chrono::duration_cast<std::chrono::microseconds>(
884 end_time - rpc_start_time)
885 .count();
886 auto total_duration =
887 std::chrono::duration_cast<std::chrono::microseconds>(
888 end_time - start_time)
889 .count();
890 std::string stage_id =
891 "grpc_client::EmbUpdate|" +
892 std::to_string(
893 trace_id == 0
894 ? static_cast<uint64_t>(
895 std::chrono::duration_cast< std::chrono::microseconds>(
896 start_time.time_since_epoch())
897 .count())
898 : trace_id);
899 report("embupdate_stages",
900 stage_id.c_str(),
901 "client_serialize_us",
902 static_cast<double>(serialize_duration));
903 report("embupdate_stages",
904 stage_id.c_str(),
905 "client_rpc_us",
906 static_cast<double>(rpc_duration));
907 report("embupdate_stages",
908 stage_id.c_str(),
909 "client_total_us",
910 static_cast<double>(total_duration));
911 report("embupdate_stages",
912 stage_id.c_str(),
913 "client_request_size",
914 static_cast<double>(keys.Size()));
915 #endif
916
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (!status.ok()) {
917 ✗ LOG(ERROR) << "UpdateParameter RPC failed: " << status.error_message();
918 ✗ return -1;
919 }
920
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 2 return response.success() ? 0 : -1;
921 2 }
922
923 ✗ int GRPCParameterClient::UpdateParameterFlat(
924 const std::string& table_name,
925 const base::ConstArray<uint64_t>& keys,
926 const float* grads,
927 int64_t num_rows,
928 int64_t embedding_dim) {
929 #ifdef ENABLE_PERF_REPORT
930 auto start_time = std::chrono::high_resolution_clock::now();
931 const uint64_t trace_id = recstore::g_trace_id;
932 #endif
933 ✗ if (keys.Size() == 0) {
934 ✗ return 0;
935 }
936
937 ✗ ParameterCompressor compressor;
938 ✗ if (BuildUpdateBlocksFromFlat(
939 ✗ keys, grads, num_rows, embedding_dim, &compressor) != 0) {
940 ✗ return -1;
941 }
942 #ifdef ENABLE_PERF_REPORT
943 auto serialize_done_time = std::chrono::high_resolution_clock::now();
944 #endif
945 ✗ UpdateParameterRequest request;
946 ✗ UpdateParameterResponse response;
947 request.set_table_name(table_name);
948 ✗ compressor.ToBlock(request.mutable_gradients());
949 ✗ if (request.gradients().empty()) {
950 ✗ return 0;
951 }
952
953 ✗ grpc::ClientContext context;
954 ✗ SetRpcDeadline(&context);
955 #ifdef ENABLE_PERF_REPORT
956 if (trace_id != 0) {
957 context.AddMetadata("x-recstore-trace-id", std::to_string(trace_id));
958 }
959 auto rpc_start_time = std::chrono::high_resolution_clock::now();
960 #endif
961 grpc::Status status =
962 ✗ stubs_[0]->UpdateParameter(&context, request, &response);
963 #ifdef ENABLE_PERF_REPORT
964 auto end_time = std::chrono::high_resolution_clock::now();
965 auto serialize_duration =
966 std::chrono::duration_cast<std::chrono::microseconds>(
967 serialize_done_time - start_time)
968 .count();
969 auto rpc_duration =
970 std::chrono::duration_cast<std::chrono::microseconds>(
971 end_time - rpc_start_time)
972 .count();
973 auto total_duration =
974 std::chrono::duration_cast<std::chrono::microseconds>(
975 end_time - start_time)
976 .count();
977 std::string stage_id =
978 "grpc_client::EmbUpdate|" +
979 std::to_string(
980 trace_id == 0
981 ? static_cast<uint64_t>(
982 std::chrono::duration_cast< std::chrono::microseconds>(
983 start_time.time_since_epoch())
984 .count())
985 : trace_id);
986 report("embupdate_stages",
987 stage_id.c_str(),
988 "client_serialize_us",
989 static_cast<double>(serialize_duration));
990 report("embupdate_stages",
991 stage_id.c_str(),
992 "client_rpc_us",
993 static_cast<double>(rpc_duration));
994 report("embupdate_stages",
995 stage_id.c_str(),
996 "client_total_us",
997 static_cast<double>(total_duration));
998 report("embupdate_stages",
999 stage_id.c_str(),
1000 "client_request_size",
1001 static_cast<double>(num_rows));
1002 report("embupdate_stages",
1003 stage_id.c_str(),
1004 "client_embedding_dim",
1005 static_cast<double>(embedding_dim));
1006 #endif
1007 ✗ if (!status.ok()) {
1008 ✗ LOG(ERROR) << "UpdateParameterFlat RPC failed: " << status.error_message();
1009 ✗ return -1;
1010 }
1011 ✗ return response.success() ? 0 : -1;
1012 ✗ }
1013
1014 20 int GRPCParameterClient::InitEmbeddingTable(
1015 const std::string& table_name,
1016 const recstore::EmbeddingTableConfig& config) {
1017
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 InitEmbeddingTableRequest request;
1018
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 InitEmbeddingTableResponse response;
1019 request.set_table_name(table_name);
1020
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 40 request.set_config_payload(config.Serialize());
1021
1022
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 grpc::ClientContext context;
1023 grpc::Status status =
1024
1/2
✓ Branch 3 taken 20 times.
✗ Branch 4 not taken.
 20 stubs_[0]->InitEmbeddingTable(&context, request, &response);
1025
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 20 times.
 20 if (!status.ok()) {
1026 ✗ LOG(ERROR) << "InitEmbeddingTable RPC failed: " << status.error_message();
1027 ✗ return -1;
1028 }
1029
2/4
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
✗ Branch 4 not taken.
 20 return response.success() ? 0 : -1;
1030 20 }
1031
1032 // BasePSClient pure virtual implementations
1033 // int GRPCParameterClient::GetParameter(const base::ConstArray<uint64_t>& keys,
1034 // float* values) {
1035 // return GetParameter(ConstArray<uint64_t>(keys.Data(), keys.Size()), values)
1036 // ? 0 : -1;
1037 // }
1038
1039 ✗ int GRPCParameterClient::AsyncGetParameter(
1040 const base::ConstArray<uint64_t>& keys, float* values) {
1041 ✗ return GetParameter(keys, values);
1042 }
1043
1044 144 int GRPCParameterClient::PutParameter(
1045 const base::ConstArray<uint64_t>& keys,
1046 const std::vector<std::vector<float>>& values) {
1047
1/2
✓ Branch 5 taken 144 times.
✗ Branch 6 not taken.
 144 std::vector<uint64_t> key_vec(keys.Data(), keys.Data() + keys.Size());
1048
1/2
✓ Branch 1 taken 144 times.
✗ Branch 2 not taken.
 144 bool success = PutParameter(key_vec, values);
1049
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 144 times.
 144 if (!success) {
1050 ✗ LOG(ERROR) << "PutParameter batch failed";
1051 }
1052
1/2
✓ Branch 0 taken 144 times.
✗ Branch 1 not taken.
 288 return success ? 1 : 0;
1053 144 }
1054
1055 ✗ void GRPCParameterClient::Command(recstore::PSCommand command) {
1056 ✗ switch (command) {
1057 ✗ case recstore::PSCommand::CLEAR_PS:
1058 ✗ ClearPS();
1059 ✗ break;
1060 ✗ case recstore::PSCommand::RELOAD_PS:
1061
1062 ✗ LOG(WARNING) << "RELOAD_PS command requires additional parameters";
1063 ✗ break;
1064 ✗ case recstore::PSCommand::LOAD_FAKE_DATA: {
1065 ✗ int64_t fake_data = 1000;
1066 ✗ LoadFakeData(fake_data);
1067 ✗ } break;
1068 ✗ case recstore::PSCommand::DUMP_FAKE_DATA: {
1069 ✗ DumpFakeData(4096);
1070 ✗ } break;
1071 ✗ case recstore::PSCommand::SAVE_CHECKPOINT:
1072 case recstore::PSCommand::LOAD_CHECKPOINT:
1073 ✗ LOG(WARNING) << "Checkpoint command requires path and metadata";
1074 ✗ break;
1075 ✗ default:
1076 ✗ LOG(ERROR) << "Unknown PS command: " << static_cast<int>(command);
1077 ✗ break;
1078 }
1079 ✗ }
1080
1081 8 uint64_t GRPCParameterClient::EmbWriteAsync(
1082 const base::ConstArray<uint64_t>& keys,
1083 const std::vector<std::vector<float>>& values) {
1084 8 uint64_t prewrite_id = next_prewrite_id_++;
1085 int request_num =
1086 8 (keys.Size() + MAX_PARAMETER_BATCH - 1) / MAX_PARAMETER_BATCH;
1087
1088
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 struct PrewriteBatch pb(request_num);
1089
1090
2/2
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 8 times.
 16 for (int start = 0, index = 0; start < keys.Size();
1091 8 start += MAX_PARAMETER_BATCH, ++index) {
1092 8 int key_size = std::min((int)(keys.Size() - start), MAX_PARAMETER_BATCH);
1093 8 pb.key_sizes_[index] = key_size;
1094 8 auto& status = pb.status_[index];
1095
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 if (!pb.contexts_[index]) {
1096
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 pb.contexts_[index] = std::make_unique<grpc::ClientContext>();
1097 }
1098 8 auto& request = pb.requests_[index];
1099 8 auto& response = pb.responses_[index];
1100
1101 // Pack key/embedding pairs
1102
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 ParameterCompressor compressor;
1103 8 std::vector<std::string> blocks;
1104
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 8 times.
 104 for (int i = start; i < start + key_size; i++) {
1105 96 auto each_key = keys[i];
1106 96 auto& embedding = values[i];
1107 96 ParameterPack parameter_pack;
1108 96 parameter_pack.key = each_key;
1109 96 parameter_pack.dim = embedding.size();
1110 96 parameter_pack.emb_data = embedding.data();
1111
1/2
✓ Branch 1 taken 96 times.
✗ Branch 2 not taken.
 96 compressor.AddItem(parameter_pack, &blocks);
1112 }
1113
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 compressor.ToBlock(&blocks);
1114
2/8
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 8 times.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 8 CHECK_EQ(blocks.size(), 1);
1115
1116
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 request.mutable_parameter_value()->swap(blocks[0]);
1117
1118 // Issue async RPC
1119
2/4
✓ Branch 5 taken 8 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
 16 pb.response_readers_.emplace_back(stubs_[0]->AsyncPutParameter(
1120 8 pb.contexts_[index].get(), request, pb.cqs_.get()));
1121
1122 // Async call; completion via CQ tag
1123 8 auto& rpc = pb.response_readers_.back();
1124
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 rpc->Finish(&response, &status, reinterpret_cast<void*>(index));
1125 8 }
1126
1127 // Store batch state in prewrite_batches_
1128
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 prewrite_batches_.emplace(prewrite_id, std::move(pb));
1129 8 return prewrite_id;
1130 8 }
1131
1132 ✗ bool GRPCParameterClient::IsWriteDone(uint64_t write_id) {
1133 ✗ LOG(ERROR) << "IsWriteDone not implemented!";
1134 ✗ auto it = prewrite_batches_.find(write_id);
1135 ✗ if (it == prewrite_batches_.end()) {
1136 ✗ LOG(ERROR) << "Invalid prewrite_id: " << write_id;
1137 ✗ return false;
1138 }
1139 ✗ auto& pb = it->second;
1140 ✗ return (pb.completed_count_ == pb.batch_size_);
1141 }
1142
1143 8 void GRPCParameterClient::WaitForWrite(uint64_t write_id) {
1144
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 auto it = prewrite_batches_.find(write_id);
1145
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (it == prewrite_batches_.end()) {
1146 ✗ LOG(ERROR) << "Invalid prewrite_id: " << write_id;
1147 ✗ return;
1148 }
1149 8 auto& pb = it->second;
1150
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 while (pb.completed_count_ < pb.batch_size_) {
1151 8 void* got_tag = nullptr;
1152 8 bool ok = false;
1153
2/4
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (!pb.cqs_->Next(&got_tag, &ok)) {
1154 ✗ break;
1155 }
1156
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (unlikely(!ok)) {
1157 ✗ LOG(ERROR) << "Completion queue returned not ok for write";
1158 ✗ continue;
1159 }
1160 8 pb.completed_count_++;
1161 }
1162 }
1163
1164 // Register GRPCParameterClient with the factory
1165 using BasePSClient = recstore::BasePSClient;
1166 FACTORY_REGISTER(BasePSClient, grpc, GRPCParameterClient, json);
1167