ps/brpc/brpc_ps_client.h - GCC Code Coverage Report

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/shard_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			class BRPCParameterClient : public recstore::ShardClient {
59			public:
60			// New constructor with JSON config
61			explicit BRPCParameterClient(json config);
62
63			// Legacy constructor for backward compatibility
64			explicit BRPCParameterClient(const std::string& host, int port, int shard);
65
66		52	~BRPCParameterClient() {}
67
68			int64_t ExchangeSparseGradientCommit(const std::string& table,
69			const std::string& payload) override;
70
71			int GetParameter(const base::ConstArray<uint64_t>& keys,
72			base::RecTensor& values) override;
73
74			int PutParameter(const base::ConstArray<uint64_t>& keys,
75			const base::RecTensor& values) override;
76
77			void Command(recstore::PSCommand command) override;
78
79			inline int shard() const { return shard_; }
80
81			bool ClearPS();
82
83			bool LoadFakeData(int64_t data);
84
85			bool DumpFakeData(int64_t n);
86
87			bool LoadCkpt(const std::vector<std::string>& model_config_path,
88			const std::vector<std::string>& emb_file_path);
89
90			bool PutParameter(const std::vector<uint64_t>& keys,
91			const base::RecTensor& values);
92
93			int UpdateParameter(const std::string& table_name,
94			const base::ConstArray<uint64_t>& keys,
95			const base::RecTensor& grads) override;
96
97			int InitEmbeddingTable(const std::string& table_name,
98			const recstore::EmbeddingTableConfig& config) override;
99
100			// Prefetch API
101			uint64_t PrefetchParameter(const base::ConstArray<uint64_t>& keys) override;
102			bool IsPrefetchDone(uint64_t prefetch_id) override;
103			void WaitForPrefetch(uint64_t prefetch_id) override;
104			bool
105			GetPrefetchResult(uint64_t prefetch_id, base::RecTensor& values) override;
106
107			protected:
108			bool Initialize();
109
110			std::string host_;
111			int port_;
112			int shard_;
113			int timeout_ms_;
114			int max_retry_;
115
116			// bRPC channel
117			std::shared_ptr<brpc::Channel> channel_;
118
119			std::vector<float> cache_;
120			std::vector<int32_t> offset_;
121
122			private:
123			std::unordered_map<uint64_t, struct BrpcPrefetchBatch> prefetch_batches_;
124			uint64_t next_prefetch_id_ = 1;
125			};
126