GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
37 / 48
77.1%
Branches:
32 / 64
50.0%
Functions:
7 / 11
63.6%
ps/base/base_client.h
Line Branch Exec Source
1 #pragma once
2 #include <cstdint>
3 #include <cstring>
4 #include <stdexcept>
5 #include <string>
6 #include <vector>
7
8 #include "base/array.h"
9 #include "base/json.h"
10 #include "base/log.h"
11 #include "base/tensor.h"
12 #include "ps/base/sparse_gradient_commit.h"
13
14 namespace recstore {
15 struct EmbeddingTableConfig {
16 uint64_t num_embeddings = 0;
17 uint64_t embedding_dim = 0;
18 uint64_t table_id = 0;
19
20 20 std::string Serialize() const {
21 20 nlohmann::json payload{{"num_embeddings", num_embeddings},
22 20 {"embedding_dim", embedding_dim},
23
7/14
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 5 taken 20 times.
✗ Branch 6 not taken.
✓ Branch 8 taken 20 times.
✗ Branch 9 not taken.
✓ Branch 12 taken 20 times.
✗ Branch 13 not taken.
✓ Branch 15 taken 20 times.
✗ Branch 16 not taken.
✓ Branch 19 taken 20 times.
✗ Branch 20 not taken.
✓ Branch 22 taken 20 times.
✗ Branch 23 not taken.
 260 {"table_id", table_id}};
24
1/2
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
 40 return payload.dump();
25 20 }
26 };
27
28 enum class PSCommand {
29 CLEAR_PS,
30 RELOAD_PS,
31 LOAD_FAKE_DATA,
32 DUMP_FAKE_DATA,
33 SAVE_CHECKPOINT,
34 LOAD_CHECKPOINT,
35 };
36
37 inline bool
38 570 IsFloatEmbeddingValues(const base::RecTensor& values, int64_t num_keys) {
39
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 570 times.
 570 if (values.dtype() != base::DataType::FLOAT32) {
40 ✗ return false;
41 }
42
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 568 times.
 570 if (num_keys == 0) {
43 2 return true;
44 }
45
2/2
✓ Branch 2 taken 566 times.
✓ Branch 3 taken 2 times.
 1136 return values.dim() == 2 && values.shape(0) == num_keys &&
46
3/6
✓ Branch 0 taken 568 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 566 times.
✗ Branch 4 not taken.
✓ Branch 6 taken 566 times.
✗ Branch 7 not taken.
 1136 values.shape(1) > 0 && values.data() != nullptr;
47 }
48
49 170 inline bool EnsureEmbeddingOutput(base::RecTensor& values, int64_t num_keys) {
50
7/10
✓ Branch 1 taken 170 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 60 times.
✓ Branch 5 taken 110 times.
✓ Branch 7 taken 60 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 60 times.
✗ Branch 10 not taken.
✓ Branch 11 taken 60 times.
✓ Branch 12 taken 110 times.
 170 if (values.dim() == 2 && values.shape(0) == 0 && values.shape(1) > 0 &&
51 num_keys >= 0) {
52 values =
53
2/4
✓ Branch 3 taken 60 times.
✗ Branch 4 not taken.
✓ Branch 6 taken 60 times.
✗ Branch 7 not taken.
 60 base::RecTensor({num_keys, values.shape(1)}, base::DataType::FLOAT32);
54
2/4
✓ Branch 0 taken 60 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 60 times.
✗ Branch 4 not taken.
 60 return num_keys == 0 || values.data() != nullptr;
55 }
56 110 return IsFloatEmbeddingValues(values, num_keys);
57 }
58
59 128 inline void GatherEmbeddingRows(const base::RecTensor& src,
60 const std::vector<size_t>& src_indices,
61 base::RecTensor& dst) {
62 128 const int64_t D = src.shape(1);
63 128 const float* s = src.data_as<float>();
64 128 float* d = dst.data_as<float>();
65
2/2
✓ Branch 1 taken 1454 times.
✓ Branch 2 taken 128 times.
 1582 for (size_t i = 0; i < src_indices.size(); ++i) {
66 2908 std::memcpy(d + i * D,
67 1454 s + src_indices[i] * D,
68 1454 static_cast<size_t>(D) * sizeof(float));
69 }
70 128 }
71
72 108 inline void ScatterEmbeddingRows(const base::RecTensor& src,
73 const std::vector<size_t>& dst_indices,
74 base::RecTensor& dst) {
75 108 const int64_t D = dst.shape(1);
76 108 const float* s = src.data_as<float>();
77 108 float* d = dst.data_as<float>();
78
2/2
✓ Branch 1 taken 586 times.
✓ Branch 2 taken 108 times.
 694 for (size_t i = 0; i < dst_indices.size(); ++i) {
79 586 std::memcpy(d + dst_indices[i] * D,
80 586 s + i * D,
81 586 static_cast<size_t>(D) * sizeof(float));
82 }
83 108 }
84
85 class BasePSClient : public SparseGradientCommitClient {
86 json json_config_;
87
88 public:
89
1/2
✓ Branch 2 taken 86 times.
✗ Branch 3 not taken.
 86 explicit BasePSClient(json config) : json_config_(config) {}
90 86 virtual ~BasePSClient() {}
91
92 virtual int GetParameter(const base::ConstArray<uint64_t>& keys,
93 base::RecTensor& values) = 0;
94
95 virtual int PutParameter(const base::ConstArray<uint64_t>& keys,
96 const base::RecTensor& values) = 0;
97 virtual int UpdateParameter(const std::string& table_name,
98 const base::ConstArray<uint64_t>& keys,
99 const base::RecTensor& grads) = 0;
100
101 virtual int InitEmbeddingTable(const std::string& table_name,
102 const EmbeddingTableConfig& config) = 0;
103
104 virtual void Command(PSCommand command) = 0;
105
106 virtual bool
107 ✗ SaveCheckpoint(const std::string& path, const std::string& metadata) {
108 (void)path;
109 (void)metadata;
110
0/4
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
 ✗ LOG(ERROR) << "Checkpoint save is unsupported by this PS client";
111 ✗ return false;
112 }
113 virtual bool
114 ✗ LoadCheckpoint(const std::string& path, const std::string& metadata) {
115 (void)path;
116 (void)metadata;
117
0/4
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
 ✗ LOG(ERROR) << "Checkpoint load is unsupported by this PS client";
118 ✗ return false;
119 }
120
121 virtual uint64_t
122 PrefetchParameter(const base::ConstArray<uint64_t>& keys) = 0;
123 virtual bool IsPrefetchDone(uint64_t prefetch_id) = 0;
124 virtual void WaitForPrefetch(uint64_t prefetch_id) = 0;
125 virtual bool
126 GetPrefetchResult(uint64_t prefetch_id, base::RecTensor& values) = 0;
127
128 // Asynchronous update. Backends without a native async path fail at
129 // submit time so callers never treat handle 0 as in-flight work.
130 ✗ virtual uint64_t SubmitUpdateParameterAsync(
131 const std::string& table_name,
132 const base::ConstArray<uint64_t>& keys,
133 const base::RecTensor& grads) {
134 (void)table_name;
135 (void)keys;
136 (void)grads;
137
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 ✗ throw std::runtime_error("Async update is not supported by this backend");
138 }
139 ✗ virtual int WaitUpdateParameter(uint64_t update_id) {
140 (void)update_id;
141
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 ✗ throw std::runtime_error("Async update is not supported by this backend");
142 }
143 };
144
145 } // namespace recstore
146