GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 100.0% 16 / 0 / 16
Functions: 100.0% 3 / 0 / 3
Branches: 50.0% 7 / 0 / 14
ps/brpc/brpc_ps_client.h
Line Branch Exec Source
1 #pragma once
2
3 #include <brpc/channel.h>
4 #include <brpc/controller.h>
5 #include <butil/logging.h>
6
7 #include <cstdint>
8 #include <memory>
9 #include <string>
10 #include <unordered_map>
11 #include <vector>
12 #include <atomic>
13
14 #include "base/array.h"
15 #include "base/flatc.h"
16 #include "base/json.h"
17 #include "base/tensor.h"
18 #include "ps/base/base_client.h"
19 #include "ps/base/parameters.h"
20 #include "ps_brpc.pb.h"
21
22 using json = nlohmann::json;
23
24 // Increased from 2000 to 65536 to reduce per-call protobuf message and
25 // brpc::Controller allocation overhead. The server max_batch_keys_size is
26 // also 65536, so this stays within the single-request limit.
27 static const int MAX_PARAMETER_BATCH_BRPC = 65536;
28
29 // Prefetch batch structure for bRPC
30 struct BrpcPrefetchBatch {
31 10 BrpcPrefetchBatch(int request_num) {
32 10 batch_size_ = request_num;
33
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 key_sizes_.resize(request_num);
34
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 responses_.resize(request_num);
35
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 controllers_.resize(request_num);
36 10 completed_count_ = 0;
37 10 }
38
39 BrpcPrefetchBatch(BrpcPrefetchBatch&& other) noexcept
40 : key_sizes_(std::move(other.key_sizes_)),
41 responses_(std::move(other.responses_)),
42 controllers_(std::move(other.controllers_)),
43 batch_size_(other.batch_size_),
44 completed_count_(other.completed_count_.load()) {
45 other.batch_size_ = 0;
46 }
47
48 BrpcPrefetchBatch(const BrpcPrefetchBatch&) = delete;
49 BrpcPrefetchBatch& operator=(const BrpcPrefetchBatch&) = delete;
50
51 std::vector<int> key_sizes_;
52 std::vector<recstoreps_brpc::GetParameterResponse> responses_;
53 std::vector<std::unique_ptr<brpc::Controller>> controllers_;
54 int batch_size_;
55 std::atomic<int> completed_count_;
56 };
57
58 struct BrpcPrewriteBatch {
59 8 BrpcPrewriteBatch(int request_num) {
60 8 batch_size_ = request_num;
61
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 key_sizes_.resize(request_num);
62
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 requests_.resize(request_num);
63
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 responses_.resize(request_num);
64
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 controllers_.resize(request_num);
65 8 completed_count_ = 0;
66 8 }
67
68 BrpcPrewriteBatch(BrpcPrewriteBatch&& other) noexcept
69 : key_sizes_(std::move(other.key_sizes_)),
70 requests_(std::move(other.requests_)),
71 responses_(std::move(other.responses_)),
72 controllers_(std::move(other.controllers_)),
73 batch_size_(other.batch_size_),
74 completed_count_(other.completed_count_.load()) {
75 other.batch_size_ = 0;
76 }
77
78 BrpcPrewriteBatch(const BrpcPrewriteBatch&) = delete;
79 BrpcPrewriteBatch& operator=(const BrpcPrewriteBatch&) = delete;
80
81 std::vector<int> key_sizes_;
82 std::vector<recstoreps_brpc::PutParameterRequest> requests_;
83 std::vector<recstoreps_brpc::PutParameterResponse> responses_;
84 std::vector<std::unique_ptr<brpc::Controller>> controllers_;
85 int batch_size_;
86 std::atomic<int> completed_count_;
87 };
88
89 class BRPCParameterClient : public recstore::BasePSClient {
90 public:
91 // New constructor with JSON config
92 explicit BRPCParameterClient(json config);
93
94 // Legacy constructor for backward compatibility
95 explicit BRPCParameterClient(const std::string& host, int port, int shard);
96
97 54 ~BRPCParameterClient() {}
98
99 // BasePSClient pure virtual implementations
100 virtual int
101 GetParameter(const base::ConstArray<uint64_t>& keys, float* values) override;
102
103 int AsyncGetParameter(const base::ConstArray<uint64_t>& keys,
104 float* values) override;
105
106 int PutParameter(const base::ConstArray<uint64_t>& keys,
107 const std::vector<std::vector<float>>& values) override;
108
109 void Command(recstore::PSCommand command) override;
110
111 // Legacy API methods
112 int GetParameter(const base::ConstArray<uint64_t>& keys,
113 std::vector<std::vector<float>>* values);
114
115 inline int shard() const { return shard_; }
116
117 bool ClearPS();
118
119 bool LoadFakeData(int64_t data);
120
121 bool DumpFakeData(int64_t n);
122
123 bool LoadCkpt(const std::vector<std::string>& model_config_path,
124 const std::vector<std::string>& emb_file_path);
125
126 bool PutParameter(const std::vector<uint64_t>& keys,
127 const std::vector<std::vector<float>>& values);
128
129 int UpdateParameter(const std::string& table_name,
130 const base::ConstArray<uint64_t>& keys,
131 const std::vector<std::vector<float>>* grads);
132 int UpdateParameterFlat(const std::string& table_name,
133 const base::ConstArray<uint64_t>& keys,
134 const float* grads,
135 int64_t num_rows,
136 int64_t embedding_dim) override;
137
138 int InitEmbeddingTable(const std::string& table_name,
139 const recstore::EmbeddingTableConfig& config);
140
141 // Prefetch API
142 uint64_t PrefetchParameter(const base::ConstArray<uint64_t>& keys);
143 bool IsPrefetchDone(uint64_t prefetch_id);
144 void WaitForPrefetch(uint64_t prefetch_id);
145 bool GetPrefetchResult(uint64_t prefetch_id,
146 std::vector<std::vector<float>>* values);
147 bool GetPrefetchResultFlat(uint64_t prefetch_id,
148 std::vector<float>* values,
149 int64_t* num_rows,
150 int64_t embedding_dim) override;
151
152 virtual uint64_t
153 EmbWriteAsync(const base::RecTensor& keys, const base::RecTensor& values);
154 virtual bool IsWriteDone(uint64_t write_id);
155 virtual void WaitForWrite(uint64_t write_id);
156
157 protected:
158 bool Initialize();
159
160 std::string host_;
161 int port_;
162 int shard_;
163 int timeout_ms_;
164 int max_retry_;
165
166 // bRPC channel
167 std::shared_ptr<brpc::Channel> channel_;
168
169 std::vector<float> cache_;
170 std::vector<int32_t> offset_;
171
172 private:
173 std::unordered_map<uint64_t, struct BrpcPrefetchBatch> prefetch_batches_;
174 std::unordered_map<uint64_t, struct BrpcPrewriteBatch> prewrite_batches_;
175 uint64_t next_prefetch_id_ = 1;
176 uint64_t next_prewrite_id_ = 1;
177 };
178