GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
25 / 25
100.0%
Branches:
7 / 14
50.0%
Functions:
4 / 4
100.0%
ps/grpc/grpc_ps_client.h
Line Branch Exec Source
1 #pragma once
2
3 #include <cstdint>
4 #include <future>
5 #include <mutex>
6 #include <string>
7 #include <unordered_map>
8 #include <vector>
9
10 #include "base/array.h"
11 #include "base/flatc.h"
12 #include "base/init.h"
13 #include "base/json.h"
14 #include "ps/base/shard_client.h"
15 #include "ps/base/parameters.h"
16 #include "ps.grpc.pb.h"
17 #include "ps.pb.h"
18 #include "base/tensor.h"
19
20 using grpc::Channel;
21 using grpc::ClientContext;
22 using grpc::Status;
23 using recstoreps::CommandRequest;
24 using recstoreps::CommandResponse;
25 using recstoreps::GetParameterRequest;
26 using recstoreps::GetParameterResponse;
27 using recstoreps::PSCommand;
28 using recstoreps::PutParameterRequest;
29 using recstoreps::PutParameterResponse;
30
31 using base::ConstArray;
32 using json = nlohmann::json;
33
34 static const int MAX_PARAMETER_BATCH = 2000;
35
36 struct PrefetchBatch {
37 112 PrefetchBatch(int request_num) {
38 112 batch_size_ = request_num;
39
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 key_sizes_.resize(request_num);
40
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 status_.resize(request_num);
41
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 contexts_.resize(request_num);
42
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 requests_.resize(request_num);
43
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 responses_.resize(request_num);
44
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 response_readers_.resize(request_num);
45
1/2
✓ Branch 1 taken 112 times.
✗ Branch 2 not taken.
 112 cqs_ = std::make_unique<grpc::CompletionQueue>();
46 112 completed_count_ = 0;
47 112 }
48
49 112 PrefetchBatch(PrefetchBatch&& other) noexcept
50 112 : key_sizes_(std::move(other.key_sizes_)),
51 112 status_(std::move(other.status_)),
52 112 contexts_(std::move(other.contexts_)),
53 112 requests_(std::move(other.requests_)),
54 112 responses_(std::move(other.responses_)),
55 112 response_readers_(std::move(other.response_readers_)),
56 112 batch_size_(other.batch_size_),
57 112 cqs_(std::move(other.cqs_)),
58 112 completed_count_(other.completed_count_) {
59 112 other.batch_size_ = 0;
60 112 }
61 PrefetchBatch(const PrefetchBatch&) = delete;
62 PrefetchBatch& operator=(const PrefetchBatch&) = delete;
63
64 std::vector<int> key_sizes_;
65 std::vector<Status> status_;
66 std::vector<std::unique_ptr<ClientContext>> contexts_;
67 std::vector<GetParameterRequest> requests_;
68 std::vector<GetParameterResponse> responses_;
69 std::vector<
70 std::unique_ptr<grpc::ClientAsyncResponseReader<GetParameterResponse>>>
71 response_readers_;
72
73 int batch_size_;
74 int completed_count_;
75 std::unique_ptr<grpc::CompletionQueue> cqs_;
76 };
77
78 class GRPCParameterClient : public recstore::ShardClient {
79 public:
80 // New constructor with JSON config
81 explicit GRPCParameterClient(json config);
82
83 // Legacy constructor for backward compatibility
84 explicit GRPCParameterClient(const std::string& host, int port, int shard);
85
86 68 ~GRPCParameterClient() {}
87
88 int64_t ExchangeSparseGradientCommit(const std::string& table,
89 const std::string& payload) override;
90
91 int GetParameter(const base::ConstArray<uint64_t>& keys,
92 base::RecTensor& values) override;
93
94 int PutParameter(const base::ConstArray<uint64_t>& keys,
95 const base::RecTensor& values) override;
96
97 void Command(recstore::PSCommand command) override;
98
99 inline int shard() const { return shard_; }
100
101 bool ClearPS();
102
103 bool LoadFakeData(int64_t data);
104
105 // Write n bytes of random floats to storage at key 0. n must be a positive
106 // multiple of sizeof(float).
107 bool DumpFakeData(int64_t n);
108
109 bool LoadCkpt(const std::vector<std::string>& model_config_path,
110 const std::vector<std::string>& emb_file_path);
111
112 bool
113 SaveCheckpoint(const std::string& path, const std::string& metadata) override;
114 bool
115 LoadCheckpoint(const std::string& path, const std::string& metadata) override;
116
117 bool PutParameter(const std::vector<uint64_t>& keys,
118 const base::RecTensor& values);
119
120 int UpdateParameter(const std::string& table_name,
121 const base::ConstArray<uint64_t>& keys,
122 const base::RecTensor& grads) override;
123
124 int InitEmbeddingTable(const std::string& table_name,
125 const recstore::EmbeddingTableConfig& config) override;
126
127 uint64_t PrefetchParameter(const base::ConstArray<uint64_t>& keys) override;
128 bool IsPrefetchDone(uint64_t prefetch_id) override;
129 void WaitForPrefetch(uint64_t prefetch_id) override;
130 bool
131 GetPrefetchResult(uint64_t prefetch_id, base::RecTensor& values) override;
132
133 protected:
134 36 bool Initialize() { return true; }
135 std::string host_;
136 int port_;
137 int shard_;
138 int nr_clients_;
139 std::vector<float> cache_;
140 std::vector<int32_t> offset_;
141 std::vector<int> get_param_key_sizes_;
142 std::vector<Status> get_param_status_;
143 std::vector<GetParameterRequest> get_param_requests_;
144 std::vector<GetParameterResponse> get_param_responses_;
145 std::vector<std::unique_ptr<grpc::ClientContext>> get_param_contexts_;
146 std::vector<
147 std::unique_ptr<grpc::ClientAsyncResponseReader<GetParameterResponse>>>
148 get_param_resonse_readers_;
149 std::shared_ptr<Channel> channel_;
150 std::vector<std::unique_ptr<recstoreps::ParameterService::Stub>> stubs_;
151 std::unique_ptr<grpc::CompletionQueue> cq_;
152
153 private:
154 std::mutex prefetch_mu_;
155 std::unordered_map<uint64_t, struct PrefetchBatch> prefetch_batches_;
156 // start from 1
157 uint64_t next_prefetch_id_ = 1;
158 };
159