GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 60.5% 221 / 0 / 365
Functions: 82.9% 29 / 0 / 35
Branches: 32.4% 268 / 0 / 827
framework/op.cc
Line Branch Exec Source
1 #include "framework/op.h"
2 #include "framework/common/hierkv_local_runtime.h"
3 #include "framework/common/local_shm_op_component.h"
4 #include "framework/common/op_runtime_support.h"
5 #include "framework/common/ps_client_config_adapter.h"
6 #include "ps/client_factory.h"
7 #include "ps/brpc/dist_brpc_ps_client.h"
8 #include "ps/grpc/dist_grpc_ps_client.h"
9 #include "ps/rdma/rdma_ps_client_adapter.h"
10 #include "base/factory.h"
11 #include <algorithm>
12 #include <cctype>
13 #include <cstring>
14 #include <immintrin.h>
15 #include <iostream>
16 #include <stdexcept>
17 #include <vector>
18 #include <unordered_map>
19 #include <filesystem>
20 #include <mutex>
21 #include <memory>
22 #include <numeric>
23 #include <thread>
24 #include <cstdlib>
25 #include <emmintrin.h>
26 #include <string>
27 #include <fstream>
28 #include "base/tensor.h"
29 #include <glog/logging.h>
30 #ifdef ENABLE_PERF_REPORT
31 # include "base/report/report_client.h"
32 #endif
33
34 namespace recstore {
35
36 namespace {
37 24 std::string NormalizeBackendName(std::string backend_name) {
38 24 std::transform(
39 backend_name.begin(),
40 backend_name.end(),
41 backend_name.begin(),
42 136 [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
43 24 return backend_name;
44 }
45
46 58 bool IsReadWriteSuccess(BasePSClient* client, int ret) {
47
3/4
✓ Branch 0 taken 58 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 50 times.
 116 if (dynamic_cast<RDMAPSClientAdapter*>(client) != nullptr ||
48
2/4
✓ Branch 0 taken 58 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 66 dynamic_cast<DistributedGRPCParameterClient*>(client) != nullptr ||
49
5/8
✓ Branch 0 taken 58 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
✓ Branch 6 taken 50 times.
✓ Branch 7 taken 8 times.
 124 dynamic_cast<DistributedBRPCParameterClient*>(client) != nullptr ||
50
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 dynamic_cast<LocalShmPSClient*>(client) != nullptr) {
51 50 return ret == 0;
52 }
53 // Legacy GRPC/BRPC read/write methods return bool-like int values.
54 8 return ret != 0;
55 }
56
57 24 std::string ResolveBackendNameWithHierKV(const json& config) {
58
3/6
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 5 taken 24 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 24 times.
✗ Branch 8 not taken.
 24 if (config.contains("cache_ps") && config["cache_ps"].contains("ps_type")) {
59 const std::string ps_type =
60
4/8
✓ 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.
 24 NormalizeBackendName(config["cache_ps"]["ps_type"].get<std::string>());
61
3/4
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
✓ Branch 4 taken 4 times.
 24 if (IsHierKVBackendName(ps_type)) {
62
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 return ps_type;
63 }
64
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 20 times.
 24 }
65
1/5
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 4 switch (ResolveFrameworkPSClientType(config)) {
66 4 case PSClientType::kGrpc:
67
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 return HasFrameworkDistributedClientConfig(config)
68 ? "distributed_grpc"
69
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 8 : "grpc";
70 ✗ case PSClientType::kBrpc:
71 ✗ return HasFrameworkDistributedClientConfig(config)
72 ? "distributed_brpc"
73 ✗ : "brpc";
74 ✗ case PSClientType::kRdma:
75 ✗ return "rdma";
76 ✗ case PSClientType::kLocalShm:
77 ✗ return "local_shm";
78 }
79
80 ✗ return "unknown";
81 }
82 } // namespace
83
84 70 void validate_keys(const base::RecTensor& keys) {
85
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 68 times.
 70 if (keys.dtype() != base::DataType::UINT64) {
86
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 throw std::invalid_argument("Keys tensor must have dtype UINT64, but got " +
87
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 4 base::DataTypeToString(keys.dtype()));
88 }
89
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 68 times.
 68 if (keys.dim() != 1) {
90 ✗ throw std::invalid_argument("Keys tensor must be 1-dimensional, but has " +
91 ✗ std::to_string(keys.dim()) + " dimensions.");
92 }
93 68 }
94
95 66 void validate_embeddings(const base::RecTensor& embeddings,
96 const std::string& name) {
97
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 66 times.
 66 if (embeddings.dtype() != base::DataType::FLOAT32) {
98 ✗ throw std::invalid_argument(
99 ✗ name + " tensor must have dtype FLOAT32, but got " +
100 ✗ base::DataTypeToString(embeddings.dtype()));
101 }
102
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 66 times.
 66 if (embeddings.dim() != 2) {
103 ✗ throw std::invalid_argument(
104 ✗ name + " tensor must be 2-dimensional, but has " +
105 ✗ std::to_string(embeddings.dim()) + " dimensions.");
106 }
107 // No fixed embedding dimension check for mock.
108 66 }
109
110 2 void KVClientOp::EmbInit(const base::RecTensor& keys,
111 const base::RecTensor& init_values) {
112 2 EmbWrite(keys, init_values);
113 2 }
114
115 2 void KVClientOp::EmbDelete(const base::RecTensor& keys) {
116
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
117 }
118 2 bool KVClientOp::EmbExists(const base::RecTensor& keys) {
119
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
120 }
121
122 2 void KVClientOp::WaitForWrite(uint64_t write_id) {
123
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
124 }
125 2 void KVClientOp::SaveToFile(const std::string& path) {
126
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
127 }
128 2 void KVClientOp::LoadFromFile(const std::string& path) {
129
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
130 }
131
132 ✗ bool KVClientOp::SaveCheckpoint(const std::string& path,
133 const std::string& metadata) {
134 ✗ if (path.empty() || metadata.empty()) {
135 ✗ throw std::invalid_argument(
136 ✗ "checkpoint path and metadata must be non-empty");
137 }
138 #ifdef USE_FAKE_KVCLIENT
139 throw std::runtime_error("Checkpoint save is unsupported by fake KVClient");
140 #else
141 ✗ if (ps_client_ == nullptr) {
142 ✗ throw std::runtime_error("PS client is not initialized");
143 }
144 ✗ return ps_client_->SaveCheckpoint(path, metadata);
145 #endif
146 }
147
148 ✗ bool KVClientOp::LoadCheckpoint(const std::string& path,
149 const std::string& metadata) {
150 ✗ if (path.empty() || metadata.empty()) {
151 ✗ throw std::invalid_argument(
152 ✗ "checkpoint path and metadata must be non-empty");
153 }
154 #ifdef USE_FAKE_KVCLIENT
155 throw std::runtime_error("Checkpoint load is unsupported by fake KVClient");
156 #else
157 ✗ if (ps_client_ == nullptr) {
158 ✗ throw std::runtime_error("PS client is not initialized");
159 }
160 ✗ return ps_client_->LoadCheckpoint(path, metadata);
161 #endif
162 }
163
164 2 uint64_t KVClientOp::EmbWriteAsync(const base::RecTensor& keys,
165 const base::RecTensor& values) {
166
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
167 }
168
169 70 std::shared_ptr<CommonOp> GetKVClientOp() {
170
3/4
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 64 times.
✓ Branch 3 taken 6 times.
✗ Branch 4 not taken.
 70 static std::shared_ptr<CommonOp> instance;
171 static std::once_flag once_flag;
172
1/2
✓ Branch 1 taken 70 times.
✗ Branch 2 not taken.
 70 std::call_once(once_flag, []() {
173 6 instance = std::make_shared<KVClientOp>();
174 6 });
175 70 return instance;
176 }
177
178 } // namespace recstore
179
180 #ifndef USE_FAKE_KVCLIENT
181
182 namespace recstore {
183
184
1/2
✓ Branch 3 taken 34 times.
✗ Branch 4 not taken.
 34 KVClientOp::KVClientOp() {
185
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 10 times.
 34 if (!ps_client_) {
186 try {
187
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 json config = GetGlobalConfig();
188
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 ps_backend_name_ = ResolveBackendNameWithHierKV(config);
189
3/4
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
✓ Branch 4 taken 4 times.
 24 if (IsHierKVBackendName(ps_backend_name_)) {
190
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 20 ConfigureLogging();
191
3/6
✓ 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.
 20 LOG(INFO) << "Initialized local HierKV backend in KVClientOp.";
192 20 return;
193 }
194 4 bool use_rdma = false;
195 try {
196
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 use_rdma = ResolveFrameworkPSClientType(config) == PSClientType::kRdma;
197 ✗ } catch (...) {
198 ✗ use_rdma = false;
199 ✗ }
200
201
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (use_rdma) {
202 ✗ InitializeRdmaProcessRuntime();
203 ✗ ConfigureLogging(false);
204 } else {
205
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 ConfigureLogging();
206 }
207
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 ps_client_holder_ = create_ps_client_from_config(config);
208 4 ps_client_ = ps_client_holder_.get();
209
210
3/6
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 4 times.
✗ Branch 8 not taken.
 4 LOG(INFO) << "PS client initialized successfully.";
211
2/4
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 20 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 24 } catch (const std::exception& e) {
212 ✗ LOG(ERROR) << "Failed to initialize PS client: " << std::string(e.what());
213 ✗ throw;
214 ✗ }
215 }
216 ✗ }
217
218 BasePSClient* KVClientOp::ps_client_ = nullptr;
219 std::unique_ptr<BasePSClient> KVClientOp::ps_client_holder_;
220
221 ✗ void KVClientOp::SetPSConfig(const std::string& host, int port) {
222 ✗ if (IsHierKVBackendName(ps_backend_name_)) {
223 ✗ LOG(INFO) << "HierKV backend ignores set_ps_config host=" << host
224 ✗ << " port=" << port;
225 ✗ return;
226 }
227 ✗ ps_client_holder_.reset();
228 ✗ ps_client_ = nullptr;
229
230 ✗ json file_config = GetGlobalConfig();
231 ✗ int final_port = port;
232 ✗ if (final_port <= 0) {
233 ✗ if (file_config.contains("client") &&
234 ✗ file_config["client"].contains("port")) {
235 ✗ final_port = file_config["client"]["port"].get<int>();
236 ✗ } else if (file_config.contains("cache_ps") &&
237 ✗ file_config["cache_ps"].contains("servers") &&
238 ✗ file_config["cache_ps"]["servers"].is_array() &&
239 ✗ !file_config["cache_ps"]["servers"].empty()) {
240 ✗ final_port = file_config["cache_ps"]["servers"][0]["port"].get<int>();
241 } else {
242 ✗ final_port = 15000;
243 }
244 }
245
246 ✗ std::string final_host = host;
247 ✗ if (final_host.empty()) {
248 ✗ final_host = "127.0.0.1";
249 }
250
251 ✗ json config = file_config;
252 ✗ config.erase("distributed_client");
253 ✗ if (!config.contains("client")) {
254 ✗ config["client"] = json::object();
255 }
256 ✗ config["client"]["host"] = final_host;
257 ✗ config["client"]["port"] = final_port;
258 ✗ config["client"]["shard"] = 0;
259
260 ✗ ps_client_holder_ = create_ps_client_from_config(config);
261 ✗ ps_client_ = ps_client_holder_.get();
262 ✗ ps_backend_name_ = ResolveBackendNameWithHierKV(config);
263 ✗ LOG(INFO) << "Re-initialized PS client with host=" << final_host
264 ✗ << " port=" << final_port;
265 ✗ }
266
267 ✗ void KVClientOp::SetPSBackend(const std::string& backend) {
268 ✗ if (backend.empty()) {
269 ✗ throw std::invalid_argument("backend must be non-empty");
270 }
271
272 ✗ const std::string normalized_backend = NormalizeBackendName(backend);
273 ✗ if (IsHierKVBackendName(normalized_backend)) {
274 ✗ ps_client_holder_.reset();
275 ✗ ps_client_ = nullptr;
276 ✗ ps_backend_name_ = normalized_backend;
277 ✗ LOG(INFO) << "Switched KVClientOp backend to local HierKV runtime.";
278 ✗ return;
279 }
280
281 ✗ json config = GetGlobalConfig();
282 ✗ if (!config.contains("cache_ps")) {
283 ✗ config["cache_ps"] = json::object();
284 }
285 ✗ config["cache_ps"]["ps_type"] = NormalizePSType(backend);
286
287 ✗ ps_client_holder_.reset();
288 ✗ ps_client_ = nullptr;
289 ✗ ps_client_holder_ = create_ps_client_from_config(config);
290 ✗ ps_client_ = ps_client_holder_.get();
291 ✗ ps_backend_name_ = ResolveBackendNameWithHierKV(config);
292 ✗ LOG(INFO) << "Re-initialized PS client with backend=" << ps_backend_name_;
293 ✗ }
294
295 2 std::string KVClientOp::CurrentPSBackend() const { return ps_backend_name_; }
296
297 48 void KVClientOp::EmbRead(const RecTensor& keys, RecTensor& values) {
298
3/4
✓ Branch 1 taken 48 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✓ Branch 4 taken 38 times.
 48 if (IsHierKVBackendName(ps_backend_name_)) {
299
3/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 6 times.
✓ Branch 5 taken 4 times.
 10 GetHierKVLocalRuntime().Read(keys, values);
300 6 return;
301 }
302
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 38 times.
 38 if (ps_client_ == nullptr) {
303 ✗ throw std::runtime_error("PS client is not initialized. Please call "
304 ✗ "KVClientOp::SetPSClient() first.");
305 }
306
307 # ifdef ENABLE_PERF_REPORT
308 auto start_time = std::chrono::high_resolution_clock::now();
309 double start_us =
310 std::chrono::duration_cast<std::chrono::microseconds>(
311 start_time.time_since_epoch())
312 .count();
313 std::string report_id =
314 "op::EmbRead|" + std::to_string(static_cast<uint64_t>(start_us));
315 std::string unique_id =
316 "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us));
317 # endif
318
319
6/12
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 38 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 38 times.
✗ Branch 14 not taken.
✓ Branch 16 taken 38 times.
✗ Branch 17 not taken.
 76 LOG(INFO) << "EmbRead: keys.shape=" << keys.shape(0) << ", values.shape=["
320
6/12
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 38 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 38 times.
✗ Branch 14 not taken.
✓ Branch 16 taken 38 times.
✗ Branch 17 not taken.
 38 << values.shape(0) << ", " << values.shape(1) << "]";
321
5/10
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 38 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 38 times.
✗ Branch 14 not taken.
 76 LOG(INFO) << "EmbRead: keys.data=" << keys.data_as<uint64_t>()
322
3/6
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
 38 << ", values.data=" << values.data_as<float>();
323
2/4
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 38 times.
✗ Branch 4 not taken.
 38 if (keys.shape(0) > 0) {
324
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 std::ostringstream oss;
325
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 oss << "EmbRead: keys start with: ";
326
3/4
✓ Branch 1 taken 142 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 104 times.
✓ Branch 5 taken 38 times.
 142 for (int i = 0; i < std::min((int64_t)10, keys.shape(0)); ++i)
327
3/6
✓ Branch 1 taken 104 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 104 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 104 times.
✗ Branch 8 not taken.
 104 oss << keys.data_as<uint64_t>()[i] << ", ";
328
4/8
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 38 times.
✗ Branch 11 not taken.
 38 LOG(INFO) << oss.str();
329 38 }
330
2/4
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 38 times.
✗ Branch 4 not taken.
 38 if (values.shape(0) > 0) {
331
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 std::ostringstream oss;
332
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 oss << "EmbRead: values start with: ";
333
3/4
✓ Branch 1 taken 140 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 102 times.
✓ Branch 5 taken 38 times.
 140 for (int i = 0; i < std::min((int64_t)10, values.shape(0)); ++i) {
334
1/2
✓ Branch 1 taken 102 times.
✗ Branch 2 not taken.
 102 oss << "[";
335
3/4
✓ Branch 1 taken 1052 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 950 times.
✓ Branch 5 taken 102 times.
 1052 for (int j = 0; j < std::min((int64_t)10, values.shape(1)); ++j) {
336
4/8
✓ Branch 1 taken 950 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 950 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 950 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 950 times.
✗ Branch 11 not taken.
 950 oss << values.data_as<float>()[i * values.shape(1) + j] << ", ";
337 }
338
1/2
✓ Branch 1 taken 102 times.
✗ Branch 2 not taken.
 102 oss << "] ";
339 }
340
4/8
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 38 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 38 times.
✗ Branch 11 not taken.
 38 LOG(INFO) << oss.str();
341 38 }
342
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 validate_keys(keys);
343
2/4
✓ Branch 2 taken 38 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 38 times.
✗ Branch 6 not taken.
 38 validate_embeddings(values, "Values");
344
345
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 const int64_t L = keys.shape(0);
346
3/4
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 36 times.
 38 if (values.shape(0) != L) {
347
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 throw std::invalid_argument(
348
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
 4 "Dimension mismatch: Keys has length " + std::to_string(L) +
349
3/6
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
 8 " but values has length " + std::to_string(values.shape(0)));
350 }
351
1/2
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
 36 const uint64_t* keys_data = keys.data_as<uint64_t>();
352 36 base::ConstArray<uint64_t> keys_array(keys_data, L);
353
1/2
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
 36 float* values_data = values.data_as<float>();
354
355
1/2
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
 36 const int64_t D = values.shape(1);
356 36 const size_t total = static_cast<size_t>(L) * static_cast<size_t>(D);
357
1/2
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
 36 std::fill_n(values_data, total, 0.0f);
358
359
1/2
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
 36 int ret = ps_client_->GetParameter(keys_array, values_data);
360
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 34 times.
 36 if (!IsReadWriteSuccess(ps_client_, ret)) {
361
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Failed to read embeddings from PS client.");
362 }
363
364 # ifdef ENABLE_PERF_REPORT
365 auto end_time = std::chrono::high_resolution_clock::now();
366 auto duration =
367 std::chrono::duration_cast<std::chrono::microseconds>(
368 end_time - start_time)
369 .count();
370 std::string op_latency_key =
371 "EmbRead|" + std::to_string(static_cast<uint64_t>(start_us));
372 report("op_latency",
373 op_latency_key.c_str(),
374 "recstore_us",
375 static_cast<double>(duration));
376
377 report("embread_stages",
378 report_id.c_str(),
379 "duration_us",
380 static_cast<double>(duration));
381
382 report("embread_stages",
383 report_id.c_str(),
384 "request_size",
385 static_cast<double>(keys.shape(0)));
386
387 FlameGraphData op_data = {
388 "op::EmbRead",
389 start_us,
390 0, // level
391 static_cast<double>(duration),
392 static_cast<double>(duration)};
393 report_flame_graph("emb_read_flame_map", unique_id.c_str(), op_data);
394 # endif
395 }
396
397 4 void KVClientOp::EmbUpdate(const base::RecTensor& keys,
398 const base::RecTensor& grads) {
399
2/4
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
 8 EmbUpdate("default", keys, grads);
400 ✗ }
401
402 10 void KVClientOp::EmbUpdate(const std::string& table_name,
403 const base::RecTensor& keys,
404 const base::RecTensor& grads) {
405
3/4
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 6 times.
 10 if (IsHierKVBackendName(ps_backend_name_)) {
406
3/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 2 times.
 4 GetHierKVLocalRuntime().Update(table_name, keys, grads);
407 2 return;
408 }
409
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (ps_client_ == nullptr) {
410 ✗ throw std::runtime_error("PS client is not initialized. Please call "
411 ✗ "KVClientOp::SetPSClient() first.");
412 }
413
414 # ifdef ENABLE_PERF_REPORT
415 auto start_time = std::chrono::high_resolution_clock::now();
416 const uint64_t trace_id = static_cast<uint64_t>(
417 std::chrono::duration_cast<std::chrono::microseconds>(
418 start_time.time_since_epoch())
419 .count());
420 struct TraceGuard {
421 explicit TraceGuard(uint64_t new_trace_id)
422 : previous_trace_id_(recstore::g_trace_id) {
423 recstore::g_trace_id = new_trace_id;
424 }
425 ~TraceGuard() { recstore::g_trace_id = previous_trace_id_; }
426 uint64_t previous_trace_id_;
427 } trace_guard(trace_id);
428 # endif
429
430 6 int64_t validate_done_us = 0;
431
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 validate_keys(keys);
432
2/4
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 6 times.
✗ Branch 6 not taken.
 6 validate_embeddings(grads, "Grads");
433
434
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 const int64_t L = keys.shape(0);
435
2/4
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
 6 if (grads.shape(0) != L) {
436 ✗ throw std::invalid_argument(
437 ✗ "Dimension mismatch: Keys has length " + std::to_string(L) +
438 ✗ " but grads has length " + std::to_string(grads.shape(0)));
439 }
440
441
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 const int64_t D = grads.shape(1);
442
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (D <= 0) {
443 ✗ throw std::invalid_argument(
444 ✗ "Invalid grad dimension D: " + std::to_string(D));
445 }
446
447 # ifdef ENABLE_PERF_REPORT
448 auto validate_done_time = std::chrono::high_resolution_clock::now();
449 validate_done_us =
450 std::chrono::duration_cast<std::chrono::microseconds>(
451 validate_done_time - start_time)
452 .count();
453 # endif
454
455
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 const uint64_t* keys_data = keys.data_as<uint64_t>();
456 6 base::ConstArray<uint64_t> keys_array(keys_data, L);
457
458
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 const float* grads_data = grads.data_as<float>();
459 int ret =
460
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 ps_client_->UpdateParameterFlat(table_name, keys_array, grads_data, L, D);
461
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 if (ret != 0) {
462
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Failed to update embeddings via PS client.");
463 }
464
465 # ifdef ENABLE_PERF_REPORT
466 auto end_time = std::chrono::high_resolution_clock::now();
467 auto duration =
468 std::chrono::duration_cast<std::chrono::microseconds>(
469 end_time - start_time)
470 .count();
471 double start_us =
472 std::chrono::duration_cast<std::chrono::microseconds>(
473 start_time.time_since_epoch())
474 .count();
475 std::string op_latency_key =
476 "EmbUpdate|" + std::to_string(static_cast<uint64_t>(start_us));
477 report("op_latency",
478 op_latency_key.c_str(),
479 "recstore_us",
480 static_cast<double>(duration));
481
482 std::string update_stage_id =
483 "op_client::EmbUpdate|" + std::to_string(trace_id);
484 report("embupdate_stages",
485 update_stage_id.c_str(),
486 "op_total_us",
487 static_cast<double>(duration));
488 report("embupdate_stages",
489 update_stage_id.c_str(),
490 "op_validate_us",
491 static_cast<double>(validate_done_us));
492 report("embupdate_stages",
493 update_stage_id.c_str(),
494 "request_size",
495 static_cast<double>(L));
496 report("embupdate_stages",
497 update_stage_id.c_str(),
498 "embedding_dim",
499 static_cast<double>(D));
500 # endif
501 }
502
503 ✗ uint64_t KVClientOp::EmbUpdateAsync(const std::string& table_name,
504 const base::RecTensor& keys,
505 const base::RecTensor& grads) {
506 ✗ if (ps_client_ == nullptr) {
507 ✗ throw std::runtime_error("PS client is not initialized");
508 }
509 ✗ auto* rdma_client = dynamic_cast<RDMAPSClientAdapter*>(ps_client_);
510 ✗ if (rdma_client == nullptr) {
511 ✗ throw std::runtime_error(
512 ✗ "Asynchronous embedding updates require the RDMA backend");
513 }
514 ✗ validate_keys(keys);
515 ✗ validate_embeddings(grads, "Grads");
516 ✗ if (keys.shape(0) != grads.shape(0) || grads.shape(1) <= 0) {
517 ✗ throw std::invalid_argument("Invalid asynchronous embedding update shape");
518 }
519 ✗ const int64_t rows = keys.shape(0);
520 ✗ const int64_t dim = grads.shape(1);
521 ✗ return rdma_client->SubmitUpdateParameterFlatAsync(
522 table_name,
523 ✗ base::ConstArray<uint64_t>(keys.data_as<uint64_t>(), rows),
524 ✗ grads.data_as<float>(),
525 rows,
526 ✗ dim);
527 }
528
529 ✗ void KVClientOp::WaitForEmbUpdate(uint64_t update_id) {
530 ✗ auto* rdma_client = dynamic_cast<RDMAPSClientAdapter*>(ps_client_);
531 ✗ if (rdma_client == nullptr) {
532 ✗ throw std::runtime_error(
533 ✗ "Asynchronous embedding updates require the RDMA backend");
534 }
535 ✗ if (rdma_client->WaitUpdateParameterFlat(update_id) != 0) {
536 ✗ throw std::runtime_error("Failed to complete asynchronous RDMA update");
537 }
538 ✗ }
539
540 18 bool KVClientOp::InitEmbeddingTable(const std::string& table_name,
541 const EmbeddingTableConfig& config) {
542
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 14 times.
 18 if (IsHierKVBackendName(ps_backend_name_)) {
543 4 return GetHierKVLocalRuntime().InitEmbeddingTable(table_name, config);
544 }
545
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (ps_client_ == nullptr) {
546 ✗ throw std::runtime_error("PS client is not initialized. Please call "
547 ✗ "KVClientOp::SetPSClient() first.");
548 }
549
550 # ifdef ENABLE_PERF_REPORT
551 auto start_time = std::chrono::high_resolution_clock::now();
552 # endif
553 14 int ret = ps_client_->InitEmbeddingTable(table_name, config);
554 # ifdef ENABLE_PERF_REPORT
555 auto end_time = std::chrono::high_resolution_clock::now();
556 auto duration =
557 std::chrono::duration_cast<std::chrono::microseconds>(
558 end_time - start_time)
559 .count();
560 double start_us =
561 std::chrono::duration_cast<std::chrono::microseconds>(
562 start_time.time_since_epoch())
563 .count();
564 std::string op_latency_key =
565 "InitEmbeddingTable|" + std::to_string(static_cast<uint64_t>(start_us));
566 // report(table_name, key, metric_name, value)
567 report("op_latency",
568 op_latency_key.c_str(),
569 "recstore_us",
570 static_cast<double>(duration));
571 # endif
572 14 return ret == 0;
573 }
574
575 36 void KVClientOp::EmbWrite(const RecTensor& keys, const RecTensor& values) {
576
3/4
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 14 times.
✓ Branch 4 taken 22 times.
 36 if (IsHierKVBackendName(ps_backend_name_)) {
577
3/4
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 6 times.
✓ Branch 5 taken 8 times.
 14 GetHierKVLocalRuntime().Write(keys, values);
578 6 return;
579 }
580
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 22 times.
 22 if (ps_client_ == nullptr) {
581 ✗ throw std::runtime_error("PS client is not initialized. Please call "
582 ✗ "KVClientOp::SetPSClient() first.");
583 }
584
585 # ifdef ENABLE_PERF_REPORT
586 auto start_time = std::chrono::high_resolution_clock::now();
587 # endif
588
589
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 validate_keys(keys);
590
2/4
✓ Branch 2 taken 22 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 22 times.
✗ Branch 6 not taken.
 22 validate_embeddings(values, "Values");
591
592
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 const int64_t L = keys.shape(0);
593 22 const auto& values_shape = values.shape();
594
2/4
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 22 times.
 22 if (values.shape(0) != L) {
595 ✗ throw std::invalid_argument(
596 ✗ "Dimension mismatch: Keys has length " + std::to_string(L) +
597 ✗ " but values has length " + std::to_string(values.shape(0)));
598 }
599
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 const int64_t D = values.shape(1);
600
601
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 const uint64_t* keys_data = keys.data_as<uint64_t>();
602 22 base::ConstArray<uint64_t> keys_array(keys_data, L);
603
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 const float* values_data = values.data_as<float>();
604
605 22 const int64_t total_values = L * D;
606
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 22 times.
 22 if (values_shape[0] * values_shape[1] != total_values) {
607 ✗ throw std::invalid_argument(
608 ✗ "Values total elements mismatch: expected " +
609 ✗ std::to_string(total_values) + ", but got " +
610 ✗ std::to_string(values_shape[0] * values_shape[1]));
611 }
612
613
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 22 times.
 22 if (D <= 0) {
614 ✗ throw std::invalid_argument(
615 ✗ "Invalid embedding dimension D: " + std::to_string(D));
616 }
617
618 22 std::vector<std::vector<float>> values_vector;
619
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 values_vector.reserve(L);
620
2/2
✓ Branch 0 taken 600 times.
✓ Branch 1 taken 22 times.
 622 for (int64_t i = 0; i < L; ++i) {
621
1/2
✓ Branch 2 taken 600 times.
✗ Branch 3 not taken.
 600 std::vector<float> row(D);
622 600 std::memcpy(row.data(), values_data + i * D, D * sizeof(float));
623 600 asm volatile("" ::: "memory");
624 _mm_mfence();
625
1/2
✓ Branch 2 taken 600 times.
✗ Branch 3 not taken.
 600 values_vector.push_back(std::move(row));
626 600 }
627
628
3/6
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
 22 LOG(INFO) << "=== Keys Array Info ===";
629
4/8
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 22 times.
✗ Branch 11 not taken.
 22 LOG(INFO) << "Keys size: " << L;
630
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 22 if (L > 0) {
631
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 std::ostringstream keys_stream;
632
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 keys_stream << "First 3 keys: ";
633
2/2
✓ Branch 1 taken 62 times.
✓ Branch 2 taken 22 times.
 84 for (int64_t i = 0; i < std::min(L, static_cast<int64_t>(3)); ++i) {
634
2/4
✓ Branch 2 taken 62 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 62 times.
✗ Branch 6 not taken.
 62 keys_stream << keys_array[i] << " ";
635 }
636
4/8
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 22 times.
✗ Branch 11 not taken.
 22 LOG(INFO) << keys_stream.str();
637 22 }
638
639
3/6
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
 22 LOG(INFO) << "=== Values Vector Info ===";
640
4/8
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 22 times.
✗ Branch 11 not taken.
 22 LOG(INFO) << "Values total elements: " << total_values;
641
4/8
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 22 times.
✗ Branch 11 not taken.
 22 LOG(INFO) << "Embedding dimension D: " << D;
642
2/4
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 22 times.
✗ Branch 3 not taken.
 22 if (L > 0 && D > 0) {
643
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 std::ostringstream values_stream;
644
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 values_stream << "First 3 embeddings (each first 3 items): ";
645
2/2
✓ Branch 1 taken 62 times.
✓ Branch 2 taken 22 times.
 84 for (int64_t i = 0; i < std::min(L, static_cast<int64_t>(3)); ++i) {
646
1/2
✓ Branch 1 taken 62 times.
✗ Branch 2 not taken.
 62 values_stream << "[";
647
2/2
✓ Branch 1 taken 186 times.
✓ Branch 2 taken 62 times.
 248 for (int64_t j = 0; j < std::min(D, static_cast<int64_t>(3)); ++j) {
648
2/4
✓ Branch 3 taken 186 times.
✗ Branch 4 not taken.
✓ Branch 6 taken 186 times.
✗ Branch 7 not taken.
 186 values_stream << values_vector[i][j] << " ";
649 }
650
1/2
✓ Branch 1 taken 62 times.
✗ Branch 2 not taken.
 62 values_stream << "] ";
651 }
652
4/8
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 22 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 22 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 22 times.
✗ Branch 11 not taken.
 22 LOG(INFO) << values_stream.str();
653 22 }
654
655
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 int ret = ps_client_->PutParameter(keys_array, values_vector);
656
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 20 times.
 22 if (!IsReadWriteSuccess(ps_client_, ret)) {
657
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Failed to write embeddings to PS client.");
658 }
659
660 # ifdef ENABLE_PERF_REPORT
661 auto end_time = std::chrono::high_resolution_clock::now();
662 auto duration =
663 std::chrono::duration_cast<std::chrono::microseconds>(
664 end_time - start_time)
665 .count();
666 double start_us =
667 std::chrono::duration_cast<std::chrono::microseconds>(
668 start_time.time_since_epoch())
669 .count();
670 std::string op_latency_key =
671 "EmbWrite|" + std::to_string(static_cast<uint64_t>(start_us));
672 report("op_latency",
673 op_latency_key.c_str(),
674 "recstore_us",
675 static_cast<double>(duration));
676 # endif
677 22 }
678
679 4 void KVClientOp::EmbInit(const base::RecTensor& keys,
680 const InitStrategy& strategy) {
681 4 validate_keys(keys);
682 2 }
683
684 uint64_t
685 8 KVClientOp::EmbPrefetch(const base::RecTensor& keys, const RecTensor& values) {
686
3/4
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 4 times.
 8 if (IsHierKVBackendName(ps_backend_name_)) {
687
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 4 int64_t embedding_dim = values.dim() == 2 ? values.shape(1) : -1;
688
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (embedding_dim <= 0) {
689 ✗ embedding_dim = GetHierKVLocalRuntime().DefaultEmbeddingDim();
690 }
691
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 4 return GetHierKVLocalRuntime().Prefetch(keys, embedding_dim);
692 }
693
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 const uint64_t* keys_data = keys.data_as<uint64_t>();
694
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 int64_t L = keys.shape(0);
695 4 base::ConstArray<uint64_t> keys_array(keys_data, L);
696
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 return ps_client_->PrefetchParameter(keys_array);
697 }
698
699 6 bool KVClientOp::IsPrefetchDone(uint64_t prefetch_id) {
700
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
 6 if (IsHierKVBackendName(ps_backend_name_)) {
701 4 return GetHierKVLocalRuntime().IsPrefetchDone(prefetch_id);
702 }
703 2 return ps_client_->IsPrefetchDone(prefetch_id);
704 }
705
706 6 void KVClientOp::WaitForPrefetch(uint64_t prefetch_id) {
707
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
 6 if (IsHierKVBackendName(ps_backend_name_)) {
708 2 GetHierKVLocalRuntime().WaitForPrefetch(prefetch_id);
709 2 return;
710 }
711 4 ps_client_->WaitForPrefetch(prefetch_id);
712 }
713
714 4 void KVClientOp::GetPretchResult(uint64_t prefetch_id,
715 std::vector<std::vector<float>>* values) {
716
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
 4 if (IsHierKVBackendName(ps_backend_name_)) {
717 2 GetHierKVLocalRuntime().ConsumePrefetch(prefetch_id, values);
718 2 return;
719 }
720 2 ps_client_->GetPrefetchResult(prefetch_id, values);
721 }
722
723 6 void KVClientOp::GetPretchResultFlat(
724 uint64_t prefetch_id,
725 std::vector<float>* values,
726 int64_t* num_rows,
727 int64_t embedding_dim) {
728
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
 6 if (IsHierKVBackendName(ps_backend_name_)) {
729 2 GetHierKVLocalRuntime().ConsumePrefetchFlat(
730 prefetch_id, values, num_rows, embedding_dim);
731 2 return;
732 }
733 4 ps_client_->GetPrefetchResultFlat(
734 prefetch_id, values, num_rows, embedding_dim);
735 }
736
737 2 bool KVClientOp::IsWriteDone(uint64_t write_id) {
738 // return ps_client_->IsWriteDone(write_id);
739
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw std::runtime_error("Not impl");
740 }
741
742 namespace testing {} // namespace testing
743
744 } // namespace recstore
745
746 #else
747
748 # include "common/op_mock.cc"
749
750 #endif // USE_FAKE_KVCLIENT
751