ps/rdma/rdma_ps_client_adapter.h - GCC Code Coverage Report

Directory:	src/
Coverage:	low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%

	Coverage	Exec / Excl / Total
Lines:	100.0%	1 / 0 / 1
Functions:	100.0%	1 / 0 / 1
Branches:	-%	0 / 0 / 0

    ps/rdma/rdma_ps_client_adapter.h
    
        Line
        Branch
        Exec
        Source
      
        #pragma once
      
        #include <memory>
      
        #include <mutex>
      
        #include <string>
      
        #include <thread>
      
        #include <unordered_map>
      
        #include <unordered_set>
      
        #include <vector>
      
        #include "base/json.h"
      
        #include "ps/base/base_client.h"
      
        #include "ps/rdma/petps_client.h"
      
        #include "ps/rdma/shard_routing.h"
      
        namespace recstore {
      
        struct EmbeddedRdmaClientIdentity {
      
          int client_index         = 0;
      
          int num_client_processes = 1;
      
          int global_id            = 0;
      
        };
      
        // Derives per-process RDMA mesh identity for embedded PyTorch clients.
      
        EmbeddedRdmaClientIdentity ResolveEmbeddedRdmaClientIdentity(int num_shards);
      
        void InitializeRdmaProcessRuntime();
      
        class RDMAPSClientAdapter : public BasePSClient {
      
        public:
      
          explicit RDMAPSClientAdapter(json config);
      
        8
          ~RDMAPSClientAdapter() override = default;
      
          int GetParameter(const base::ConstArray<uint64_t>& keys,
      
                           float* values) override;
      
          int PutParameter(const base::ConstArray<uint64_t>& keys,
      
                           const std::vector<std::vector<float>>& values) override;
      
          int UpdateParameter(const std::string& table_name,
      
                              const base::ConstArray<uint64_t>& keys,
      
                              const std::vector<std::vector<float>>* grads) override;
      
          int UpdateParameterFlat(const std::string& table_name,
      
                                  const base::ConstArray<uint64_t>& keys,
      
                                  const float* grads,
      
                                  int64_t num_rows,
      
                                  int64_t embedding_dim) override;
      
          uint64_t SubmitUpdateParameterFlatAsync(
      
              const std::string& table_name,
      
              const base::ConstArray<uint64_t>& keys,
      
              const float* grads,
      
              int64_t num_rows,
      
              int64_t embedding_dim);
      
          int WaitUpdateParameterFlat(uint64_t update_id);
      
          int InitEmbeddingTable(const std::string& table_name,
      
                                 const EmbeddingTableConfig& config) override;
      
          int AsyncGetParameter(const base::ConstArray<uint64_t>& keys,
      
                                float* values) override;
      
          void Command(PSCommand command) override;
      
          uint64_t PrefetchParameter(const base::ConstArray<uint64_t>& keys) override;
      
          bool IsPrefetchDone(uint64_t prefetch_id) override;
      
          void WaitForPrefetch(uint64_t prefetch_id) override;
      
          bool GetPrefetchResult(uint64_t prefetch_id,
      
                                 std::vector<std::vector<float>>* values) override;
      
          bool GetPrefetchResultFlat(uint64_t prefetch_id,
      
                                     std::vector<float>* values,
      
                                     int64_t* num_rows,
      
                                     int64_t embedding_dim) override;
      
        private:
      
          struct TableState {
      
            EmbeddingTableConfig config;
      
          };
      
          using PendingShardRpc = shard_routing::PendingShardRpc;
      
          using BatchRequest    = shard_routing::BatchRequest;
      
          using ShardChunk      = shard_routing::ShardChunk;
      
          struct PrefetchState {
      
            std::shared_ptr<std::vector<float>> buffer;
      
            int rpc_id             = -1;
      
            int64_t key_count      = 0;
      
            int64_t embedding_dim  = 0;
      
            bool borrowed_response = false;
      
            bool batch_response    = false;
      
          };
      
          struct PendingUpdate {
      
            std::vector<std::pair<int, int>> shard_rpcs;
      
            std::thread::id owner;
      
          };
      
          void EnsureClientInitialized();
      
          void EnsureThreadInitialized();
      
          void EnsureTableReady(const std::string& table_name, int64_t embedding_dim);
      
          int64_t DefaultEmbeddingDimOrThrow() const;
      
          std::size_t MaxGetKeysPerRpc() const;
      
          std::size_t MaxPutKeysPerRpc() const;
      
          std::size_t MaxInFlightGetRpcs() const;
      
          std::vector<ShardChunk> BuildChunks(base::ConstArray<uint64_t> keys) const;
      
          void
      
          WaitShardRpcsCooperatively(const std::vector<PendingShardRpc>& shard_rpcs);
      
          int SubmitGetParameter(base::ConstArray<uint64_t> keys,
      
                                 float* values,
      
                                 bool isAsync,
      
                                 int async_req_id);
      
          bool QueryRPCFinished(int rpc_id);
      
          void WaitRPCFinish(int rpc_id);
      
          void RevokeRPCResource(int rpc_id);
      
          const float* BorrowPrefetchResult(const PrefetchState& state,
      
                                            std::int32_t* status_code,
      
                                            std::size_t* response_bytes);
      
          PrefetchState GetPrefetchState(uint64_t prefetch_id);
      
          void MarkPrefetchConsumed(uint64_t prefetch_id);
      
          json config_;
      
          std::mutex init_mu_;
      
          std::mutex thread_init_mu_;
      
          std::mutex state_mu_;
      
          bool initialized_ = false;
      
          std::unordered_set<std::thread::id> initialized_threads_;
      
          std::vector<std::unique_ptr<petps::PetPSClient>> shard_clients_;
      
          BaseParameterClient* client_ = nullptr;
      
          int num_shards_              = 1;
      
          std::string hash_method_     = "city_hash";
      
          std::unordered_map<int, int> shard_to_client_index_;
      
          int batch_rpc_id_acc_ = -1;
      
          mutable std::mutex batches_mu_;
      
          std::unordered_map<int, BatchRequest> batches_;
      
          std::unordered_map<std::string, TableState> tables_;
      
          std::unordered_map<uint64_t, PrefetchState> prefetches_;
      
          uint64_t next_prefetch_id_ = 1;
      
          std::unordered_map<uint64_t, PendingUpdate> pending_updates_;
      
          uint64_t next_update_id_ = 1;
      
        };
      
        } // namespace recstore

Line	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/petps_client.h"
14		#include "ps/rdma/shard_routing.h"
15
16		namespace recstore {
17
18		struct EmbeddedRdmaClientIdentity {
19		int client_index = 0;
20		int num_client_processes = 1;
21		int global_id = 0;
22		};
23
24		// Derives per-process RDMA mesh identity for embedded PyTorch clients.
25		EmbeddedRdmaClientIdentity ResolveEmbeddedRdmaClientIdentity(int num_shards);
26
27		void InitializeRdmaProcessRuntime();
28
29		class RDMAPSClientAdapter : public BasePSClient {
30		public:
31		explicit RDMAPSClientAdapter(json config);
32	8	~RDMAPSClientAdapter() override = default;
33
34		int GetParameter(const base::ConstArray<uint64_t>& keys,
35		float* values) override;
36		int PutParameter(const base::ConstArray<uint64_t>& keys,
37		const std::vector<std::vector<float>>& values) override;
38		int UpdateParameter(const std::string& table_name,
39		const base::ConstArray<uint64_t>& keys,
40		const std::vector<std::vector<float>>* grads) override;
41		int UpdateParameterFlat(const std::string& table_name,
42		const base::ConstArray<uint64_t>& keys,
43		const float* grads,
44		int64_t num_rows,
45		int64_t embedding_dim) override;
46		uint64_t SubmitUpdateParameterFlatAsync(
47		const std::string& table_name,
48		const base::ConstArray<uint64_t>& keys,
49		const float* grads,
50		int64_t num_rows,
51		int64_t embedding_dim);
52		int WaitUpdateParameterFlat(uint64_t update_id);
53		int InitEmbeddingTable(const std::string& table_name,
54		const EmbeddingTableConfig& config) override;
55		int AsyncGetParameter(const base::ConstArray<uint64_t>& keys,
56		float* values) override;
57		void Command(PSCommand command) override;
58		uint64_t PrefetchParameter(const base::ConstArray<uint64_t>& keys) override;
59		bool IsPrefetchDone(uint64_t prefetch_id) override;
60		void WaitForPrefetch(uint64_t prefetch_id) override;
61		bool GetPrefetchResult(uint64_t prefetch_id,
62		std::vector<std::vector<float>>* values) override;
63		bool GetPrefetchResultFlat(uint64_t prefetch_id,
64		std::vector<float>* values,
65		int64_t* num_rows,
66		int64_t embedding_dim) override;
67
68		private:
69		struct TableState {
70		EmbeddingTableConfig config;
71		};
72
73		using PendingShardRpc = shard_routing::PendingShardRpc;
74		using BatchRequest = shard_routing::BatchRequest;
75		using ShardChunk = shard_routing::ShardChunk;
76
77		struct PrefetchState {
78		std::shared_ptr<std::vector<float>> buffer;
79		int rpc_id = -1;
80		int64_t key_count = 0;
81		int64_t embedding_dim = 0;
82		bool borrowed_response = false;
83		bool batch_response = false;
84		};
85
86		struct PendingUpdate {
87		std::vector<std::pair<int, int>> shard_rpcs;
88		std::thread::id owner;
89		};
90
91		void EnsureClientInitialized();
92		void EnsureThreadInitialized();
93		void EnsureTableReady(const std::string& table_name, int64_t embedding_dim);
94		int64_t DefaultEmbeddingDimOrThrow() const;
95		std::size_t MaxGetKeysPerRpc() const;
96		std::size_t MaxPutKeysPerRpc() const;
97		std::size_t MaxInFlightGetRpcs() const;
98		std::vector<ShardChunk> BuildChunks(base::ConstArray<uint64_t> keys) const;
99		void
100		WaitShardRpcsCooperatively(const std::vector<PendingShardRpc>& shard_rpcs);
101		int SubmitGetParameter(base::ConstArray<uint64_t> keys,
102		float* values,
103		bool isAsync,
104		int async_req_id);
105		bool QueryRPCFinished(int rpc_id);
106		void WaitRPCFinish(int rpc_id);
107		void RevokeRPCResource(int rpc_id);
108		const float* BorrowPrefetchResult(const PrefetchState& state,
109		std::int32_t* status_code,
110		std::size_t* response_bytes);
111		PrefetchState GetPrefetchState(uint64_t prefetch_id);
112		void MarkPrefetchConsumed(uint64_t prefetch_id);
113
114		json config_;
115		std::mutex init_mu_;
116		std::mutex thread_init_mu_;
117		std::mutex state_mu_;
118		bool initialized_ = false;
119		std::unordered_set<std::thread::id> initialized_threads_;
120		std::vector<std::unique_ptr<petps::PetPSClient>> shard_clients_;
121		BaseParameterClient* client_ = nullptr;
122		int num_shards_ = 1;
123		std::string hash_method_ = "city_hash";
124		std::unordered_map<int, int> shard_to_client_index_;
125		int batch_rpc_id_acc_ = -1;
126		mutable std::mutex batches_mu_;
127		std::unordered_map<int, BatchRequest> batches_;
128		std::unordered_map<std::string, TableState> tables_;
129		std::unordered_map<uint64_t, PrefetchState> prefetches_;
130		uint64_t next_prefetch_id_ = 1;
131		std::unordered_map<uint64_t, PendingUpdate> pending_updates_;
132		uint64_t next_update_id_ = 1;
133		};
134
135		} // namespace recstore
136