GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 69.7% 131 / 0 / 188
Functions: 66.7% 10 / 0 / 15
Branches: 36.0% 100 / 0 / 278
ps/rdma/allshards_ps_client.cc
Line Branch Exec Source
1 #include "allshards_ps_client.h"
2
3 #include <algorithm>
4 #include <boost/coroutine2/all.hpp>
5 #include <cstring>
6 #include <limits>
7 #include <memory>
8 #include <stdexcept>
9 #include <thread>
10 #include <vector>
11
12 #include "ps/rdma/rdma_common.h"
13
14 DECLARE_int32(value_size);
15 DECLARE_int32(max_kv_num_per_request);
16
17 10 AllShardsParameterClientWrapper::AllShardsParameterClientWrapper(
18 const std::vector<BaseParameterClient*>& clients,
19 int num_shards,
20 const std::string& hash_method,
21 10 const std::vector<int>& shard_ids)
22 : BaseParameterClient("", 0, 0),
23 10 clients_(clients),
24 10 num_shards_(num_shards),
25
4/8
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 10 times.
✗ Branch 6 not taken.
✓ Branch 10 taken 10 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 10 times.
✗ Branch 14 not taken.
 10 hash_method_(hash_method) {
26
2/8
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 10 times.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 10 CHECK_EQ(static_cast<int>(clients_.size()), num_shards_);
27
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 8 times.
 10 if (!shard_ids.empty()) {
28
2/8
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 2 CHECK_EQ(static_cast<int>(shard_ids.size()), num_shards_);
29
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 for (int i = 0; i < num_shards_; ++i) {
30
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 shard_to_client_index_[shard_ids[static_cast<std::size_t>(i)]] = i;
31 }
32 } else {
33
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 for (int i = 0; i < num_shards_; ++i) {
34
1/2
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
 16 shard_to_client_index_[i] = i;
35 }
36 }
37 10 }
38
39 std::vector<AllShardsParameterClientWrapper::ShardChunk>
40 6 AllShardsParameterClientWrapper::BuildChunks(
41 base::ConstArray<uint64_t> keys) const {
42 return recstore::shard_routing::BuildChunks(
43 keys,
44 6 num_shards_,
45 6 hash_method_,
46 6 shard_to_client_index_,
47 6 static_cast<std::size_t>(FLAGS_max_kv_num_per_request));
48 }
49
50 4 void AllShardsParameterClientWrapper::WaitShardRpcsCooperatively(
51 const std::vector<PendingShardRpc>& shard_rpcs) const {
52 using Coroutine = boost::coroutines2::coroutine<void>;
53 4 std::vector<std::unique_ptr<Coroutine::pull_type>> waiters;
54
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 waiters.reserve(shard_rpcs.size());
55
2/2
✓ Branch 4 taken 12 times.
✓ Branch 5 taken 4 times.
 16 for (const auto& pending : shard_rpcs) {
56
3/6
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 12 times.
✗ Branch 8 not taken.
 12 waiters.emplace_back(std::make_unique<Coroutine::pull_type>(
57 12 [this, pending](Coroutine::push_type& sink) {
58 auto* client =
59 12 clients_[static_cast<std::size_t>(pending.client_index)];
60
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
 12 while (!client->QueryRPCFinished(pending.rpc_id)) {
61 ✗ sink();
62 }
63 12 client->WaitRPCFinish(pending.rpc_id);
64 12 }));
65 }
66
67
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
 8 while (!waiters.empty()) {
68
2/2
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 4 times.
 16 for (auto it = waiters.begin(); it != waiters.end();) {
69 12 auto& waiter = *it;
70
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 12 if (*waiter) {
71 ✗ (*waiter)();
72 }
73
1/2
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 12 if (!*waiter) {
74
1/2
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 12 it = waiters.erase(it);
75 } else {
76 ✗ ++it;
77 }
78 }
79
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
 4 if (!waiters.empty()) {
80 ✗ std::this_thread::yield();
81 }
82 }
83 4 }
84
85 2 int AllShardsParameterClientWrapper::GetParameter(
86 base::ConstArray<uint64_t> keys, std::vector<std::vector<float>>* values) {
87 2 values->clear();
88
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (keys.Size() == 0) {
89 ✗ return 0;
90 }
91
92 2 const int embedding_dim = FLAGS_value_size / sizeof(float);
93
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 2 std::vector<float> flat(keys.Size() * embedding_dim + 1, 0.0f);
94
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 int rpc_id = GetParameter(keys, flat.data(), false, 0);
95
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 WaitRPCFinish(rpc_id);
96 const auto* status_word =
97 2 petps::FixedSlotStatusWord(flat.data(), keys.Size(), FLAGS_value_size);
98
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (*status_word != static_cast<std::int32_t>(petps::RpcStatus::kOk)) {
99
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 RevokeRPCResource(rpc_id);
100 2 return -1;
101 }
102
103 ✗ petps::CopyFlatRowsToVectors(
104 ✗ flat.data(),
105 ✗ keys.Size(),
106 static_cast<std::size_t>(embedding_dim),
107 values);
108 ✗ RevokeRPCResource(rpc_id);
109 ✗ return 0;
110 2 }
111
112 6 int AllShardsParameterClientWrapper::GetParameter(
113 base::ConstArray<uint64_t> keys,
114 float* values,
115 bool isAsync,
116 int async_req_id) {
117 6 BatchRequest batch;
118 6 batch.user_buffer = values;
119 6 batch.total_key_count = keys.Size();
120 auto* batch_status_word =
121 6 petps::FixedSlotStatusWord(values, keys.Size(), FLAGS_value_size);
122 6 *batch_status_word = static_cast<std::int32_t>(petps::RpcStatus::kPending);
123
124
3/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 7 taken 16 times.
✓ Branch 8 taken 6 times.
 22 for (const auto& chunk : BuildChunks(keys)) {
125 16 void* recv = clients_[chunk.client_index]->GetReceiveBuffer(
126
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 chunk.keys.size() * static_cast<std::size_t>(FLAGS_value_size) +
127 sizeof(std::int32_t));
128
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 32 int rpc_id = clients_[chunk.client_index]->GetParameter(
129 16 base::ConstArray<uint64_t>(chunk.keys),
130 static_cast<float*>(recv),
131 isAsync,
132 async_req_id);
133
1/2
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
 16 batch.shard_rpcs.push_back(PendingShardRpc{
134 16 chunk.shard_id,
135 16 chunk.client_index,
136 rpc_id,
137
1/2
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
 16 chunk.positions,
138 recv,
139 16 chunk.keys.size(),
140 });
141 6 }
142
143 6 std::uint64_t batch_id = 0;
144 {
145
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 std::lock_guard<std::mutex> guard(batches_mu_);
146 6 batch_id = batch_rpc_id_acc_++;
147 6 if (batch_id >
148
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
 6 static_cast<std::uint64_t>(std::numeric_limits<int>::max())) {
149
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 throw std::runtime_error("allshards batch rpc id overflow int range: " +
150
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 4 std::to_string(batch_id));
151 }
152
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 batches_[batch_id] = std::move(batch);
153 6 }
154
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!isAsync) {
155
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 WaitRPCFinish(static_cast<int>(batch_id));
156 }
157 4 return static_cast<int>(batch_id);
158 6 }
159
160 2 void AllShardsParameterClientWrapper::InitThread() {
161
2/2
✓ Branch 5 taken 4 times.
✓ Branch 6 taken 2 times.
 6 for (auto* client : clients_) {
162
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 client->InitThread();
163 }
164 2 }
165
166 ✗ void AllShardsParameterClientWrapper::Barrier(const std::string& ss, int k) {
167 ✗ CHECK(!clients_.empty());
168 ✗ clients_.front()->Barrier(ss, k);
169 ✗ }
170
171 ✗ void* AllShardsParameterClientWrapper::GetReceiveBuffer(size_t size) {
172 ✗ return new char[size];
173 }
174
175 ✗ bool AllShardsParameterClientWrapper::QueryRPCFinished(int rpc_id) {
176 ✗ std::lock_guard<std::mutex> guard(batches_mu_);
177 ✗ auto it = batches_.find(rpc_id);
178 ✗ CHECK(it != batches_.end());
179
180 ✗ for (const auto& pending : it->second.shard_rpcs) {
181 ✗ if (!clients_[pending.client_index]->QueryRPCFinished(pending.rpc_id)) {
182 ✗ return false;
183 }
184 }
185
186 ✗ return recstore::shard_routing::FinalizeBatchIfNeeded(&it->second,
187 ✗ FLAGS_value_size);
188 ✗ }
189
190 8 void AllShardsParameterClientWrapper::WaitRPCFinish(int rpc_id) {
191 8 std::vector<PendingShardRpc> shard_rpcs;
192 {
193
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 std::lock_guard<std::mutex> guard(batches_mu_);
194
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 auto it = batches_.find(rpc_id);
195
2/12
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 15 not taken.
✓ Branch 16 taken 8 times.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
 8 CHECK(it != batches_.end());
196
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
 8 if (it->second.assembled) {
197 4 return;
198 }
199
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 shard_rpcs = it->second.shard_rpcs;
200
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
 8 }
201
202
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 WaitShardRpcsCooperatively(shard_rpcs);
203
204 {
205
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 std::lock_guard<std::mutex> guard(batches_mu_);
206
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 auto it = batches_.find(rpc_id);
207
2/12
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 15 not taken.
✓ Branch 16 taken 4 times.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
 4 CHECK(it != batches_.end());
208 4 recstore::shard_routing::FinalizeBatchIfNeeded(&it->second,
209 FLAGS_value_size);
210 4 }
211
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
 8 }
212
213 4 void AllShardsParameterClientWrapper::RevokeRPCResource(int rpc_id) {
214
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 std::lock_guard<std::mutex> guard(batches_mu_);
215
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 auto it = batches_.find(rpc_id);
216
2/12
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 15 not taken.
✓ Branch 16 taken 4 times.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
 4 CHECK(it != batches_.end());
217
218
2/2
✓ Branch 6 taken 12 times.
✓ Branch 7 taken 4 times.
 16 for (const auto& pending : it->second.shard_rpcs) {
219
1/2
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 12 clients_[pending.client_index]->RevokeRPCResource(pending.rpc_id);
220 }
221
222
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 batches_.erase(it);
223 4 }
224
225 4 int AllShardsParameterClientWrapper::PutParameter(
226 const std::vector<uint64_t>& keys,
227 const std::vector<std::vector<float>>& values) {
228
2/8
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
 4 CHECK_EQ(keys.size(), values.size());
229
230
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 std::vector<std::vector<uint64_t>> shard_keys(num_shards_);
231
1/2
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 std::vector<std::vector<std::vector<float>>> shard_values(num_shards_);
232
233
2/2
✓ Branch 1 taken 16 times.
✓ Branch 2 taken 4 times.
 20 for (std::size_t i = 0; i < keys.size(); ++i) {
234 const int shard =
235
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 recstore::shard_routing::PartitionKey(keys[i], num_shards_, hash_method_);
236
1/2
✓ Branch 3 taken 16 times.
✗ Branch 4 not taken.
 16 shard_keys[static_cast<std::size_t>(shard)].push_back(keys[i]);
237
1/2
✓ Branch 3 taken 16 times.
✗ Branch 4 not taken.
 16 shard_values[static_cast<std::size_t>(shard)].push_back(values[i]);
238 }
239
240
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 4 times.
 12 for (int shard = 0; shard < num_shards_; ++shard) {
241
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 const int client_index = shard_to_client_index_.at(shard);
242 8 for (std::size_t offset = 0;
243
2/2
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 8 times.
 16 offset < shard_keys[static_cast<std::size_t>(shard)].size();
244 8 offset += static_cast<std::size_t>(FLAGS_max_kv_num_per_request)) {
245 8 const std::size_t end = std::min(
246 16 offset + static_cast<std::size_t>(FLAGS_max_kv_num_per_request),
247 8 shard_keys[static_cast<std::size_t>(shard)].size());
248 std::vector<uint64_t> key_slice(
249 8 shard_keys[static_cast<std::size_t>(shard)].begin() + offset,
250
1/2
✓ Branch 7 taken 8 times.
✗ Branch 8 not taken.
 16 shard_keys[static_cast<std::size_t>(shard)].begin() + end);
251 std::vector<std::vector<float>> value_slice(
252 8 shard_values[static_cast<std::size_t>(shard)].begin() + offset,
253
1/2
✓ Branch 7 taken 8 times.
✗ Branch 8 not taken.
 16 shard_values[static_cast<std::size_t>(shard)].begin() + end);
254
1/2
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 int rc = clients_[client_index]->PutParameter(key_slice, value_slice);
255
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (rc != 0) {
256 ✗ return rc;
257 }
258
2/4
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
 8 }
259 }
260
261 4 return 0;
262 4 }
263
264 ✗ int AllShardsParameterClientWrapper::InitEmbeddingTable(
265 const std::string& table_name,
266 std::uint64_t num_embeddings,
267 std::uint64_t embedding_dim) {
268 ✗ for (auto* client : clients_) {
269 const int rc =
270 ✗ client->InitEmbeddingTable(table_name, num_embeddings, embedding_dim);
271 ✗ if (rc != 0) {
272 ✗ return rc;
273 }
274 }
275 ✗ return 0;
276 }
277
278 ✗ int AllShardsParameterClientWrapper::UpdateParameter(
279 const std::string& table_name,
280 base::ConstArray<uint64_t> keys,
281 const std::vector<std::vector<float>>* grads) {
282 ✗ if (grads == nullptr) {
283 ✗ return -1;
284 }
285 ✗ if (keys.Size() != grads->size()) {
286 ✗ return -1;
287 }
288 ✗ if (keys.Size() == 0) {
289 ✗ return 0;
290 }
291
292 ✗ std::vector<std::vector<uint64_t>> shard_keys(num_shards_);
293 ✗ std::vector<std::vector<std::vector<float>>> shard_grads(num_shards_);
294
295 ✗ for (std::size_t i = 0; i < keys.Size(); ++i) {
296 const int shard =
297 ✗ recstore::shard_routing::PartitionKey(keys[i], num_shards_, hash_method_);
298 ✗ shard_keys[static_cast<std::size_t>(shard)].push_back(keys[i]);
299 ✗ shard_grads[static_cast<std::size_t>(shard)].push_back((*grads)[i]);
300 }
301
302 ✗ for (int shard = 0; shard < num_shards_; ++shard) {
303 ✗ if (shard_keys[static_cast<std::size_t>(shard)].empty()) {
304 ✗ continue;
305 }
306 ✗ const int client_index = shard_to_client_index_.at(shard);
307 ✗ const int rc = clients_[client_index]->UpdateParameter(
308 table_name,
309 ✗ base::ConstArray<uint64_t>(shard_keys[static_cast<std::size_t>(shard)]),
310 ✗ &shard_grads[static_cast<std::size_t>(shard)]);
311 ✗ if (rc != 0) {
312 ✗ return rc;
313 }
314 }
315
316 ✗ return 0;
317 ✗ }
318