GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
1 / 11
9.1%
Branches:
0 / 4
0.0%
Functions:
1 / 6
16.7%
ps/rdma/rdma_ps_client_adapter.h
Line Branch Exec Source
1 #pragma once
2
3 #include <memory>
4 #include <mutex>
5 #include <string>
6 #include <thread>
7 #include <unordered_map>
8 #include <unordered_set>
9 #include <vector>
10
11 #include "base/json.h"
12 #include "ps/base/base_client.h"
13 #include "ps/rdma/rdma_deployment.h"
14 #include "ps/rdma/rdma_distributed_executor.h"
15 #include "ps/rdma/petps_client.h"
16 #include "ps/rdma/shard_routing.h"
17
18 namespace recstore {
19
20 struct EmbeddedRdmaClientIdentity {
21 int client_index = 0;
22 int num_client_processes = 1;
23 int global_id = 0;
24 };
25
26 // Derives per-process RDMA mesh identity for embedded PyTorch clients.
27 EmbeddedRdmaClientIdentity
28 ResolveEmbeddedRdmaClientIdentity(int num_shards, int configured_num_clients);
29
30 void InitializeRdmaProcessRuntime();
31
32 class RDMAPSClientAdapter : public BasePSClient {
33 // RdmaRawAccess exposes the raw async RPC lifecycle (submit → poll → wait →
34 // revoke) to the low-level RDMA benchmarks without widening the production
35 // BasePSClient surface.
36 friend class RdmaRawAccess;
37
38 public:
39 explicit RDMAPSClientAdapter(json config);
40 10 ~RDMAPSClientAdapter() override = default;
41
42 int64_t SparseGradientCommit(const std::string& table,
43 SparseGradientCommitHeader h,
44 const base::ConstArray<uint64_t>& keys,
45 const base::RecTensor& grads) override;
46
47 int GetParameter(const base::ConstArray<uint64_t>& keys,
48 base::RecTensor& values) override;
49 int PutParameter(const base::ConstArray<uint64_t>& keys,
50 const base::RecTensor& values) override;
51 int UpdateParameter(const std::string& table_name,
52 const base::ConstArray<uint64_t>& keys,
53 const base::RecTensor& grads) override;
54 uint64_t SubmitUpdateParameterAsync(const std::string& table_name,
55 const base::ConstArray<uint64_t>& keys,
56 const base::RecTensor& grads) override;
57 int WaitUpdateParameter(uint64_t update_id) override;
58 int InitEmbeddingTable(const std::string& table_name,
59 const EmbeddingTableConfig& config) override;
60 void Command(PSCommand command) override;
61 uint64_t PrefetchParameter(const base::ConstArray<uint64_t>& keys) override;
62 bool IsPrefetchDone(uint64_t prefetch_id) override;
63 void WaitForPrefetch(uint64_t prefetch_id) override;
64 bool
65 GetPrefetchResult(uint64_t prefetch_id, base::RecTensor& values) override;
66
67 private:
68 struct TableState {
69 EmbeddingTableConfig config;
70 int tag = 0;
71 };
72
73 using ShardChunk = shard_routing::ShardChunk;
74
75 void EnsureClientInitialized();
76 void EnsureThreadInitialized();
77 void EnsureTableReady(const std::string& table_name, int64_t embedding_dim);
78 int64_t DefaultEmbeddingDimOrThrow() const;
79 int64_t EmbeddingDimForKeys(base::ConstArray<uint64_t> keys) const;
80 std::size_t MaxGetKeysPerRpc(int64_t embedding_dim) const;
81 std::size_t MaxPutKeysPerRpc(int64_t embedding_dim) const;
82 std::size_t MaxInFlightGetRpcs() const;
83 std::vector<ShardChunk> BuildChunks(base::ConstArray<uint64_t> keys,
84 std::size_t max_keys_per_rpc) const;
85 int SubmitGetParameter(
86 base::ConstArray<uint64_t> keys,
87 float* values,
88 bool isAsync,
89 int async_req_id,
90 int64_t embedding_dim,
91 bool write_status_word = true);
92 bool QueryRPCFinished(int rpc_id);
93 void WaitRPCFinish(int rpc_id);
94 void RevokeRPCResource(int rpc_id);
95
96 json config_;
97 std::mutex init_mu_;
98 // GET windows and BagPipe submissions share finite RC slots. Serialize
99 // window ownership, while GPU work and PS rank aggregation remain async.
100 std::mutex io_mu_;
101 std::mutex thread_init_mu_;
102 mutable std::mutex state_mu_;
103 bool initialized_ = false;
104 std::unordered_set<std::thread::id> initialized_threads_;
105 ResolvedRdmaDeployment deployment_;
106 std::vector<std::unique_ptr<petps::PetPSClient>> shard_clients_;
107 petps::PetPSClient* client_ = nullptr;
108 std::unordered_map<std::string, TableState> tables_;
109 std::unordered_map<int, int64_t> tag_to_dim_;
110 std::unique_ptr<RdmaDistributedExecutor> executor_;
111 };
112
113 // Friend accessor for the low-level RDMA benchmarks. Forwards the adapter's
114 // private raw-async RPC lifecycle so a benchmark can drive multi-shard async
115 // GETs (submit → poll → wait → revoke) directly, without the benchmark
116 // reimplementing routing/batch bookkeeping.
117 class RdmaRawAccess {
118 public:
119 ✗ explicit RdmaRawAccess(RDMAPSClientAdapter* adapter) : adapter_(adapter) {}
120
121 ✗ int SubmitGetParameter(base::ConstArray<uint64_t> keys,
122 float* values,
123 bool isAsync,
124 int async_req_id,
125 int64_t embedding_dim) {
126
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ std::lock_guard<std::mutex> io_guard(adapter_->io_mu_);
127
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ return adapter_->SubmitGetParameter(
128 ✗ keys, values, isAsync, async_req_id, embedding_dim);
129 ✗ }
130 ✗ bool QueryRPCFinished(int rpc_id) {
131 ✗ return adapter_->QueryRPCFinished(rpc_id);
132 }
133 ✗ void WaitRPCFinish(int rpc_id) { adapter_->WaitRPCFinish(rpc_id); }
134 ✗ void RevokeRPCResource(int rpc_id) { adapter_->RevokeRPCResource(rpc_id); }
135
136 private:
137 RDMAPSClientAdapter* adapter_;
138 };
139
140 } // namespace recstore
141