GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 28.5% 80 / 0 / 281
Functions: 36.8% 7 / 0 / 19
Branches: 17.0% 75 / 0 / 442
ps/base/cache_ps_impl.h
Line Branch Exec Source
1 #pragma once
2
3 #include <algorithm>
4 #include <atomic>
5 #include <boost/coroutine2/all.hpp>
6 #include <cstring>
7 #include <cstdint>
8 #include <experimental/filesystem>
9 #include <mutex>
10 #include <random>
11 #include <stdexcept>
12 #include <unordered_map>
13 #include <vector>
14
15 #include "base/array.h"
16 #include "base/factory.h"
17 #include "base/log.h" // NOLINT
18 #include "base/timer.h"
19 #include "parameters.h"
20 #include "storage/kv_engine/base_kv.h"
21 #include "storage/kv_engine/engine_factory.h"
22 #include "storage/kv_engine/engine_selector.h"
23 #include "optimizer/optimizer.h"
24 #include "ps/local_shm/local_shm_stage_report.h"
25
26 #ifdef ENABLE_PERF_REPORT
27 # include <chrono>
28 # include "base/report/report_client.h"
29 #endif
30
31 using boost::coroutines2::coroutine;
32
33 static const int KEY_CNT = 12543670;
34
35 template <typename key_t>
36 struct TaskElement {
37 TaskElement(const base::ConstArray<key_t>& keys,
38 const base::MutableArray<ParameterPack>& packs,
39 std::atomic_bool* promise)
40 : keys(keys), packs(packs), promise(promise) {}
41
42 TaskElement() {}
43
44 base::ConstArray<key_t> keys;
45 base::MutableArray<ParameterPack> packs;
46 std::atomic_bool* promise;
47 };
48
49 class CachePS {
50 public:
51 using key_t = uint64_t;
52
53 struct FlatGetProfile {
54 std::uint64_t batch_get_ns = 0;
55 std::uint64_t index_lookup_ns = 0;
56 std::uint64_t zero_fill_ns = 0;
57 std::uint64_t row_copy_ns = 0;
58 std::uint64_t rows = 0;
59 std::uint64_t value_bytes = 0;
60 std::uint64_t missing_rows = 0;
61 };
62
63 using DirectFixedRow = BaseKV::DirectFixedRow;
64 using RDMABackingRegion = BaseKV::RDMABackingRegion;
65
66 30 CachePS(json config) {
67
5/10
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 30 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 30 times.
✗ Branch 11 not taken.
✓ Branch 13 taken 30 times.
✗ Branch 14 not taken.
 30 LOG(INFO) << "cache ps config: " << config.dump(2);
68 30 BaseKVConfig kv_config;
69
2/4
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
 30 kv_config.num_threads_ = config["num_threads"].get<int>();
70
2/4
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
 30 kv_config.json_config_ = config["base_kv_config"];
71
2/4
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
 30 auto r = base::ResolveEngine(kv_config);
72
1/2
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
 30 base_kv_.reset(base::Factory<BaseKV, const BaseKVConfig&>::NewInstance(
73 r.engine, r.cfg));
74
3/6
✓ Branch 1 taken 30 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 30 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 30 times.
✗ Branch 8 not taken.
 30 optimizer_ = CreateOptimizer(config.value("optimizer", json::object()));
75 30 }
76
77 30 ~CachePS() {}
78
79 ✗ bool Initialize(const std::vector<std::string>& model_config_path,
80 const std::vector<std::string>& emb_file_path) {
81 ✗ LOG(INFO) << "Before Load CKPT";
82 ✗ LoadCkpt(model_config_path, emb_file_path);
83 ✗ LOG(INFO) << "After Load CKPT";
84 ✗ return true;
85 }
86
87 ✗ void Clear() {
88 ✗ std::lock_guard<std::mutex> lock(checkpoint_mu_);
89 ✗ base_kv_->clear();
90 ✗ active_checkpoint_identity_.clear();
91 ✗ checkpoint_dirty_ = base_kv_->CheckpointRecordCount() != 0;
92 ✗ }
93
94 void LoadFakeData(int64_t key_capacity, int value_size) {
95 std::lock_guard<std::mutex> lock(checkpoint_mu_);
96 checkpoint_dirty_ = true;
97 base_kv_->LoadFakeData(key_capacity, value_size);
98 }
99
100 ✗ bool LoadCkpt(const std::vector<std::string>& model_config_path,
101 const std::vector<std::string>& emb_file_path) {
102 // base_kv_->loadCkpt();
103 // LoadFakeData(KEY_CNT);
104 ✗ return true;
105 }
106
107 ✗ bool SaveCheckpoint(const std::string& path, const std::string& metadata) {
108 ✗ const std::string identity = CheckpointIdentity(metadata);
109 ✗ std::lock_guard<std::mutex> lock(checkpoint_mu_);
110 ✗ if (!base_kv_->SaveCheckpoint(path, metadata)) {
111 ✗ LOG(ERROR) << "Failed to save checkpoint: " << path;
112 ✗ return false;
113 }
114 ✗ active_checkpoint_identity_ = identity;
115 ✗ checkpoint_dirty_ = false;
116 ✗ return true;
117 ✗ }
118
119 ✗ bool LoadCheckpoint(const std::string& path,
120 const std::string& expected_metadata) {
121 ✗ const std::string identity = CheckpointIdentity(expected_metadata);
122 ✗ std::lock_guard<std::mutex> lock(checkpoint_mu_);
123 ✗ if (!checkpoint_dirty_ && active_checkpoint_identity_ == identity) {
124 ✗ return true;
125 }
126 ✗ if (checkpoint_dirty_) {
127 ✗ throw std::runtime_error(
128 ✗ "checkpoint load rejected: parameter server has unsaved updates");
129 }
130 ✗ if (!active_checkpoint_identity_.empty()) {
131 ✗ throw std::runtime_error(
132 ✗ "checkpoint load rejected: active checkpoint identity mismatch");
133 }
134 ✗ if (base_kv_->CheckpointRecordCount() != 0) {
135 ✗ throw std::runtime_error(
136 ✗ "checkpoint load rejected: parameter server is not fresh");
137 }
138 try {
139 ✗ if (!base_kv_->LoadCheckpoint(path, expected_metadata)) {
140 ✗ checkpoint_dirty_ = base_kv_->CheckpointRecordCount() != 0;
141 ✗ LOG(ERROR) << "Failed to load checkpoint: " << path;
142 ✗ return false;
143 }
144 ✗ } catch (...) {
145 ✗ checkpoint_dirty_ = base_kv_->CheckpointRecordCount() != 0;
146 ✗ throw;
147 ✗ }
148 ✗ active_checkpoint_identity_ = identity;
149 ✗ checkpoint_dirty_ = false;
150 ✗ return true;
151 ✗ }
152
153 void PutSingleParameter(
154 const uint64_t key, const void* data, const int dim, const int tid) {
155 std::lock_guard<std::mutex> lock(checkpoint_mu_);
156 checkpoint_dirty_ = true;
157 base_kv_->Put(key, std::string_view((char*)data, dim * sizeof(float)), tid);
158 }
159
160 22 void PutDenseParameterBatch(
161 const uint64_t* keys,
162 const float* values,
163 int key_count,
164 int embedding_dim,
165 int tid) {
166
2/4
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 22 times.
 22 if (key_count <= 0 || embedding_dim <= 0) {
167 ✗ return;
168 }
169
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 std::lock_guard<std::mutex> lock(checkpoint_mu_);
170 22 checkpoint_dirty_ = true;
171 22 base::ConstArray<uint64_t> key_array(keys, key_count);
172 22 std::vector<base::ConstArray<float>> value_slices;
173
1/2
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
 22 value_slices.reserve(static_cast<std::size_t>(key_count));
174
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 22 times.
 60 for (int i = 0; i < key_count; ++i) {
175
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 value_slices.emplace_back(values + i * embedding_dim, embedding_dim);
176 }
177
1/2
✓ Branch 2 taken 22 times.
✗ Branch 3 not taken.
 22 base_kv_->BatchPut(key_array, &value_slices, tid);
178 22 }
179
180 ✗ void PutSingleParameter(const ParameterCompressItem* item, int tid) {
181 ✗ std::lock_guard<std::mutex> lock(checkpoint_mu_);
182 ✗ checkpoint_dirty_ = true;
183 ✗ auto key = item->key;
184 ✗ auto dim = item->dim;
185 ✗ base_kv_->Put(
186 ✗ key, std::string_view((char*)item->data(), dim * sizeof(float)), tid);
187 ✗ }
188
189 void PutParameter(coroutine<void>::push_type& sink,
190 const ParameterCompressReader* reader,
191 int tid) {
192 std::lock_guard<std::mutex> lock(checkpoint_mu_);
193 checkpoint_dirty_ = true;
194 std::vector<uint64_t> keys_vec;
195 std::vector<base::ConstArray<float>> values;
196 for (int i = 0; i < reader->item_size(); i++) {
197 keys_vec.emplace_back(reader->item(i)->key);
198 values.emplace_back(
199 (float*)reader->item(i)->data(), reader->item(i)->dim);
200 }
201 base::ConstArray<uint64_t> keys(keys_vec);
202
203 base_kv_->BatchPut(sink, keys, &values, tid);
204 }
205
206 ✗ void PutParameter(const ParameterCompressReader* reader, int tid) {
207 ✗ std::lock_guard<std::mutex> lock(checkpoint_mu_);
208 ✗ checkpoint_dirty_ = true;
209 ✗ std::vector<uint64_t> keys_vec;
210 ✗ std::vector<base::ConstArray<float>> values;
211 ✗ for (int i = 0; i < reader->item_size(); i++) {
212 ✗ keys_vec.emplace_back(reader->item(i)->key);
213 ✗ values.emplace_back(
214 ✗ (float*)reader->item(i)->data(), reader->item(i)->dim);
215 }
216 ✗ base::ConstArray<uint64_t> keys(keys_vec);
217
218 ✗ base_kv_->BatchPut(keys, &values, tid);
219 ✗ }
220
221 bool GetParameterRun2Completion(key_t key, ParameterPack& pack, int tid) {
222 std::vector<uint64_t> keys = {key};
223 base::ConstArray<uint64_t> keys_array(keys);
224 std::vector<base::ConstArray<float>> values;
225
226 base_kv_->BatchGet(keys_array, &values, tid);
227 base::ConstArray<float> value = values[0];
228
229 if (value.Size() == 0) {
230 pack.key = key;
231 pack.dim = 0;
232 pack.emb_data = nullptr;
233 RECSTORE_LOG_EVERY_MS(ERROR, 1000) << "key " << key << " not existing";
234 return false;
235 }
236 pack.key = key;
237 pack.dim = value.Size();
238 pack.emb_data = value.Data();
239 // LOG(ERROR) << "Get key " << key << " dim " << pack.dim;
240 return true;
241 }
242
243 10 bool GetParameterRun2Completion(base::ConstArray<uint64_t> keys,
244 std::vector<ParameterPack>& packs,
245 int tid) {
246 #ifdef ENABLE_PERF_REPORT
247 auto start_time = std::chrono::high_resolution_clock::now();
248 #endif
249 10 const auto batch_get_start = std::chrono::steady_clock::now();
250 10 std::vector<base::ConstArray<float>> values;
251
1/2
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 base_kv_->BatchGet(keys, &values, tid);
252
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 recstore::ReportLocalShmStageMetric(
253 "cache_ps_get_batch_get_us",
254 recstore::LocalShmElapsedUs(batch_get_start));
255
256 10 const auto pack_build_start = std::chrono::steady_clock::now();
257
2/2
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 10 times.
 30 for (int i = 0; i < keys.Size(); i++) {
258
1/2
✓ Branch 6 taken 20 times.
✗ Branch 7 not taken.
 20 packs.emplace_back(keys[i], values[i].Size(), values[i].Data());
259 }
260
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 recstore::ReportLocalShmStageMetric(
261 "cache_ps_get_pack_us", recstore::LocalShmElapsedUs(pack_build_start));
262
263 #ifdef ENABLE_PERF_REPORT
264 auto end_time = std::chrono::high_resolution_clock::now();
265 double start_us =
266 std::chrono::duration_cast<std::chrono::microseconds>(
267 start_time.time_since_epoch())
268 .count();
269 auto duration =
270 std::chrono::duration_cast<std::chrono::microseconds>(
271 end_time - start_time)
272 .count();
273
274 std::string report_id = "cache_ps::GetParameterRun2Completion|" +
275 std::to_string(static_cast<uint64_t>(start_us));
276
277 report("embread_stages",
278 report_id.c_str(),
279 "duration_us",
280 static_cast<double>(duration));
281
282 report("embread_stages",
283 report_id.c_str(),
284 "request_size",
285 static_cast<double>(keys.Size()));
286
287 std::string unique_id =
288 "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us));
289 FlameGraphData fg_data = {
290 "cache_ps::GetParameterRun2Completion",
291 start_us,
292 3, // level
293 static_cast<double>(duration),
294 static_cast<double>(duration)};
295 report_flame_graph("emb_read_flame_map", unique_id.c_str(), fg_data);
296 #endif
297 10 return true;
298 10 }
299
300 14 bool GetParameterFlat(
301 base::ConstArray<uint64_t> keys,
302 float* values,
303 int64_t num_rows,
304 int64_t embedding_dim,
305 int tid,
306 FlatGetProfile* profile = nullptr) {
307
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (values == nullptr) {
308 ✗ LOG(ERROR) << "GetParameterFlat values pointer is null";
309 ✗ return false;
310 }
311
2/4
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 14 times.
 14 if (num_rows < 0 || embedding_dim <= 0) {
312 ✗ LOG(ERROR) << "GetParameterFlat invalid shape rows=" << num_rows
313 ✗ << " dim=" << embedding_dim;
314 ✗ return false;
315 }
316
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 14 times.
 14 if (keys.Size() != static_cast<size_t>(num_rows)) {
317 ✗ LOG(ERROR) << "GetParameterFlat keys size mismatch " << keys.Size()
318 ✗ << " vs " << num_rows;
319 ✗ return false;
320 }
321
322 14 const auto batch_get_start = std::chrono::steady_clock::now();
323 14 BaseKV::BatchGetFlatStats flat_stats;
324 14 BaseKV::BatchGetFlatStats* flat_stats_ptr =
325
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 profile != nullptr ? &flat_stats : nullptr;
326
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 const bool flat_ok = base_kv_->BatchGetFlat(
327 keys, values, num_rows, embedding_dim, tid, flat_stats_ptr);
328
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 2 times.
 14 if (flat_ok) {
329
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 if (profile != nullptr) {
330 ✗ profile->batch_get_ns = static_cast<std::uint64_t>(
331 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
332 ✗ std::chrono::steady_clock::now() - batch_get_start)
333 ✗ .count());
334 ✗ profile->rows = static_cast<std::uint64_t>(num_rows);
335 ✗ profile->value_bytes =
336 ✗ static_cast<std::uint64_t>(num_rows) *
337 ✗ static_cast<std::uint64_t>(embedding_dim) * sizeof(float);
338 ✗ profile->zero_fill_ns = flat_stats.zero_fill_ns;
339 ✗ profile->index_lookup_ns = flat_stats.index_lookup_ns;
340 ✗ profile->row_copy_ns = flat_stats.row_copy_ns;
341 ✗ profile->missing_rows = flat_stats.missing_rows;
342 }
343
1/2
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 12 recstore::ReportLocalShmStageMetric(
344 "cache_ps_get_batch_get_us",
345 recstore::LocalShmElapsedUs(batch_get_start));
346 12 recstore::ReportLocalShmStageMetric("cache_ps_get_copy_us", 0);
347 12 return true;
348 }
349
350 2 std::vector<base::ConstArray<float>> value_slices;
351
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 value_slices.reserve(static_cast<std::size_t>(num_rows));
352
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 base_kv_->BatchGet(keys, &value_slices, tid);
353
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (profile != nullptr) {
354 ✗ profile->batch_get_ns = static_cast<std::uint64_t>(
355 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
356 ✗ std::chrono::steady_clock::now() - batch_get_start)
357 ✗ .count());
358 ✗ profile->rows = static_cast<std::uint64_t>(num_rows);
359 ✗ profile->value_bytes =
360 ✗ static_cast<std::uint64_t>(num_rows) *
361 ✗ static_cast<std::uint64_t>(embedding_dim) * sizeof(float);
362 }
363
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 recstore::ReportLocalShmStageMetric(
364 "cache_ps_get_batch_get_us",
365 recstore::LocalShmElapsedUs(batch_get_start));
366
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (value_slices.size() != static_cast<size_t>(num_rows)) {
367 ✗ LOG(ERROR) << "GetParameterFlat BatchGet returned " << value_slices.size()
368 ✗ << " rows, expected " << num_rows;
369 ✗ return false;
370 }
371
372
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 for (int64_t row = 0; row < num_rows; ++row) {
373 2 const auto& slice = value_slices[static_cast<size_t>(row)];
374
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 4 if (slice.Size() != 0 &&
375
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 static_cast<int64_t>(slice.Size()) != embedding_dim) {
376
4/8
✓ 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.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 4 LOG(ERROR) << "GetParameterFlat embedding_dim mismatch at row=" << row
377
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
 2 << " key=" << keys[static_cast<size_t>(row)] << " expected="
378
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
 2 << embedding_dim << " actual=" << slice.Size();
379 2 return false;
380 }
381 }
382
383 ✗ const auto row_copy_start = std::chrono::steady_clock::now();
384 ✗ std::uint64_t missing_zero_fill_ns = 0;
385 ✗ for (int64_t row = 0; row < num_rows; ++row) {
386 ✗ const auto& slice = value_slices[static_cast<size_t>(row)];
387 ✗ if (slice.Size() > 0) {
388 ✗ std::memcpy(values + row * embedding_dim,
389 ✗ slice.Data(),
390 ✗ static_cast<std::size_t>(embedding_dim) * sizeof(float));
391 } else {
392 ✗ const auto missing_zero_start = std::chrono::steady_clock::now();
393 ✗ std::memset(values + row * embedding_dim,
394 0,
395 ✗ static_cast<std::size_t>(embedding_dim) * sizeof(float));
396 ✗ if (profile != nullptr) {
397 ✗ missing_zero_fill_ns += static_cast<std::uint64_t>(
398 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
399 ✗ std::chrono::steady_clock::now() - missing_zero_start)
400 ✗ .count());
401 ✗ ++profile->missing_rows;
402 }
403 }
404 }
405 ✗ if (profile != nullptr) {
406 ✗ profile->zero_fill_ns = missing_zero_fill_ns;
407 ✗ profile->row_copy_ns = static_cast<std::uint64_t>(
408 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
409 ✗ std::chrono::steady_clock::now() - row_copy_start)
410 ✗ .count());
411 }
412 ✗ recstore::ReportLocalShmStageMetric(
413 "cache_ps_get_copy_us", recstore::LocalShmElapsedUs(row_copy_start));
414 ✗ return true;
415 2 }
416
417 ✗ bool ProbeParameterIndex(base::ConstArray<uint64_t> keys,
418 int tid,
419 FlatGetProfile* profile = nullptr) {
420 ✗ const auto batch_get_start = std::chrono::steady_clock::now();
421 ✗ BaseKV::BatchGetFlatStats flat_stats;
422 ✗ BaseKV::BatchGetFlatStats* flat_stats_ptr =
423 ✗ profile != nullptr ? &flat_stats : nullptr;
424 ✗ const bool ok = base_kv_->BatchGetIndexOnly(keys, tid, flat_stats_ptr);
425 ✗ if (profile != nullptr) {
426 ✗ profile->batch_get_ns = static_cast<std::uint64_t>(
427 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
428 ✗ std::chrono::steady_clock::now() - batch_get_start)
429 ✗ .count());
430 ✗ profile->rows = static_cast<std::uint64_t>(keys.Size());
431 ✗ profile->value_bytes = 0;
432 ✗ profile->missing_rows = flat_stats.missing_rows;
433 }
434 ✗ return ok;
435 }
436
437 ✗ bool GetParameterDirectFixedRows(
438 base::ConstArray<uint64_t> keys,
439 int64_t num_rows,
440 int64_t embedding_dim,
441 int tid,
442 std::vector<DirectFixedRow>* rows,
443 FlatGetProfile* profile = nullptr) {
444 ✗ const auto batch_get_start = std::chrono::steady_clock::now();
445 ✗ BaseKV::BatchGetFlatStats flat_stats;
446 ✗ BaseKV::BatchGetFlatStats* flat_stats_ptr =
447 ✗ profile != nullptr ? &flat_stats : nullptr;
448 ✗ const bool ok = base_kv_->BatchGetDirectFixedRows(
449 keys, num_rows, embedding_dim, tid, rows, flat_stats_ptr);
450 ✗ if (profile != nullptr) {
451 ✗ profile->batch_get_ns = static_cast<std::uint64_t>(
452 ✗ std::chrono::duration_cast< std::chrono::nanoseconds>(
453 ✗ std::chrono::steady_clock::now() - batch_get_start)
454 ✗ .count());
455 ✗ profile->rows = static_cast<std::uint64_t>(num_rows);
456 ✗ profile->value_bytes =
457 ✗ static_cast<std::uint64_t>(num_rows) *
458 ✗ static_cast<std::uint64_t>(embedding_dim) * sizeof(float);
459 ✗ profile->zero_fill_ns = flat_stats.zero_fill_ns;
460 ✗ profile->index_lookup_ns = flat_stats.index_lookup_ns;
461 ✗ profile->row_copy_ns = 0;
462 ✗ profile->missing_rows = flat_stats.missing_rows;
463 }
464 ✗ return ok;
465 }
466
467 ✗ RDMABackingRegion GetRDMABackingRegion() const {
468 ✗ return base_kv_->GetRDMABackingRegion();
469 }
470
471 bool GetParameterRun2Completion(
472 coroutine<void>::push_type& sink,
473 base::ConstArray<uint64_t> keys,
474 std::vector<ParameterPack>& pack,
475 int tid) {
476 std::vector<base::ConstArray<float>> values;
477
478 base_kv_->BatchGet(sink, keys, &values, tid);
479
480 for (int i = 0; i < keys.Size(); i++) {
481 pack.emplace_back(keys[i], values[i].Size(), values[i].Data());
482 }
483 return true;
484 }
485
486 /// optimizer interface
487
488 24 bool InitTable(const std::string& table_name,
489 uint64_t num_embeddings,
490 uint64_t embedding_dim) {
491
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 std::lock_guard<std::mutex> lock(checkpoint_mu_);
492 24 EmbeddingTableConfig config{num_embeddings, embedding_dim};
493
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 const auto existing = table_configs_.find(table_name);
494
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 24 times.
 24 if (existing != table_configs_.end()) {
495 ✗ const auto& old = existing->second;
496 ✗ if (old.num_embeddings == num_embeddings &&
497 ✗ old.embedding_dim == embedding_dim) {
498 ✗ return true;
499 }
500 ✗ LOG(ERROR) << "Embedding table config mismatch for '" << table_name
501 ✗ << "': existing=[" << old.num_embeddings << ", "
502 ✗ << old.embedding_dim << "] requested=[" << num_embeddings
503 ✗ << ", " << embedding_dim << "]";
504 ✗ return false;
505 }
506
5/10
✓ Branch 3 taken 24 times.
✗ Branch 4 not taken.
✓ Branch 7 taken 24 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 24 times.
✗ Branch 11 not taken.
✓ Branch 14 taken 24 times.
✓ Branch 15 taken 24 times.
✗ Branch 19 not taken.
✗ Branch 20 not taken.
 48 optimizer_->Init({table_name}, config, base_kv_.get());
507
1/2
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
 24 table_configs_.emplace(table_name, config);
508 24 return true;
509 24 }
510
511 ✗ bool UpdateParameter(const std::string& table_name,
512 const ParameterCompressReader* reader,
513 unsigned tid) {
514 ✗ std::lock_guard<std::mutex> lock(checkpoint_mu_);
515 ✗ if (!optimizer_) {
516 ✗ LOG(ERROR) << "Optimizer not initialized. Please call InitTable first.";
517 ✗ return false;
518 }
519
520 ✗ checkpoint_dirty_ = true;
521 ✗ optimizer_->Update(table_name, reader, tid);
522 ✗ return true;
523 ✗ }
524
525 6 bool UpdateParameterFlat(
526 const std::string& table_name,
527 const base::ConstArray<uint64_t>& keys,
528 const float* grads,
529 int64_t num_rows,
530 int64_t embedding_dim,
531 unsigned tid) {
532
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 std::lock_guard<std::mutex> lock(checkpoint_mu_);
533
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (grads == nullptr) {
534 ✗ LOG(ERROR) << "UpdateParameterFlat grads pointer is null";
535 ✗ return false;
536 }
537
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 if (num_rows < 0 || embedding_dim <= 0) {
538 ✗ LOG(ERROR) << "UpdateParameterFlat invalid shape rows=" << num_rows
539 ✗ << " dim=" << embedding_dim;
540 ✗ return false;
541 }
542
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
 6 if (keys.Size() != static_cast<size_t>(num_rows)) {
543 ✗ LOG(ERROR) << "UpdateParameterFlat keys size mismatch " << keys.Size()
544 ✗ << " vs " << num_rows;
545 ✗ return false;
546 }
547
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
 6 if (!optimizer_) {
548 ✗ LOG(ERROR) << "Optimizer not initialized. Please call InitTable first.";
549 ✗ return false;
550 }
551 6 checkpoint_dirty_ = true;
552
2/4
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 6 times.
✗ Branch 6 not taken.
 6 optimizer_->UpdateFlat(
553 table_name, keys, grads, num_rows, embedding_dim, tid);
554 6 return true;
555 6 }
556
557 private:
558 ✗ static std::string CheckpointIdentity(const std::string& metadata) {
559 ✗ const json parsed = json::parse(metadata);
560 ✗ if (!parsed.is_object() || !parsed.contains("identity") ||
561 ✗ !parsed["identity"].is_object() || !parsed.contains("checkpoint_id") ||
562 ✗ !parsed["checkpoint_id"].is_string() ||
563 ✗ parsed["checkpoint_id"].get<std::string>().empty()) {
564 ✗ throw std::invalid_argument(
565 ✗ "checkpoint metadata requires identity and non-empty checkpoint_id");
566 }
567 ✗ return parsed.dump();
568 ✗ }
569
570 std::unique_ptr<BaseKV> base_kv_;
571 std::unique_ptr<Optimizer> optimizer_;
572 std::mutex checkpoint_mu_;
573 std::unordered_map<std::string, EmbeddingTableConfig> table_configs_;
574 std::string active_checkpoint_identity_;
575 bool checkpoint_dirty_ = false;
576 std::atomic<bool> stopFlag_{false};
577 };
578