GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 8.3% 64 / 0 / 771
Functions: 15.9% 7 / 0 / 44
Branches: 4.5% 45 / 0 / 1010
ps/rdma/rdma_ps_client_adapter.cc
Line Branch Exec Source
1 #include "ps/rdma/rdma_ps_client_adapter.h"
2
3 #include <algorithm>
4 #include <atomic>
5 #include <chrono>
6 #include <cstddef>
7 #include <cstdint>
8 #include <cstdlib>
9 #include <cstring>
10 #include <fstream>
11 #include <limits>
12 #include <memory>
13 #include <stdexcept>
14 #include <thread>
15 #include <utility>
16
17 #include <folly/portability/GFlags.h>
18 #include <folly/init/Init.h>
19
20 #include "framework/common/ps_client_config_adapter.h"
21 #include "ps/base/config.h"
22 #include "ps/rdma/rdma_common.h"
23 #include "ps/rdma/rc_options.h"
24
25 DECLARE_int32(global_id);
26 DECLARE_int32(num_server_processes);
27 DECLARE_int32(num_client_processes);
28 DECLARE_int32(value_size);
29 DECLARE_int32(max_kv_num_per_request);
30 DECLARE_int32(rdma_rc_client_id_base);
31 DECLARE_int32(rdma_rc_num_logical_clients);
32 DECLARE_int32(rdma_control_plane_timeout_ms);
33 DECLARE_string(rdma_get_response_mode);
34 DECLARE_string(rdma_transport_mode);
35 DEFINE_string(rdma_transport_mode, "rc_write", "RDMA transport mode: rc_write");
36 DEFINE_bool(rdma_adapter_skip_prefetch_result_copy,
37 false,
38 "Benchmark-only option to skip copying RDMA prefetch results into "
39 "the GetPrefetchResultFlat output vector");
40
41 namespace recstore {
42
43 namespace detail {
44
45 34 bool TryParseIntEnv(const char* env_name, int* parsed_value) {
46 34 const char* value = std::getenv(env_name);
47
3/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 34 if (value == nullptr || *value == '\0') {
48 22 return false;
49 }
50 12 char* end = nullptr;
51 12 const long parsed = std::strtol(value, &end, 10);
52
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 12 if (end == value || *end != '\0') {
53 ✗ return false;
54 }
55 12 *parsed_value = static_cast<int>(parsed);
56 12 return true;
57 }
58
59 } // namespace detail
60
61 namespace {
62
63 ✗ bool AdapterProfileEnabled() {
64 ✗ const char* value = std::getenv("RECSTORE_RDMA_ADAPTER_PROFILE");
65 ✗ return value != nullptr && std::string(value) != "0";
66 }
67
68 14 void SetIntFlagFromEnv(const char* env_name, int32_t* flag_value) {
69 14 int parsed = 0;
70
2/4
✓ Branch 1 taken 14 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 14 times.
 14 if (detail::TryParseIntEnv(env_name, &parsed)) {
71 ✗ *flag_value = static_cast<int32_t>(parsed);
72 }
73 14 }
74
75 2 void ApplyRdmaFlagsFromEnv() {
76
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (const char* value = std::getenv("RECSTORE_RDMA_RC_NAMESPACE")) {
77 ✗ FLAGS_rdma_rc_namespace = value;
78 }
79
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 if (const char* value = std::getenv("RECSTORE_RDMA_CONTROL_PLANE_HOST")) {
80 ✗ FLAGS_rdma_control_plane_host = value;
81 }
82
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 if (const char* value = std::getenv("RECSTORE_RDMA_GET_RESPONSE_MODE")) {
83
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 const std::string mode(value);
84
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
 2 if (mode != "direct_sg" && mode != "staging_copy") {
85 ✗ throw std::runtime_error(
86 ✗ "RECSTORE_RDMA_GET_RESPONSE_MODE must be direct_sg or staging_copy");
87 }
88
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 FLAGS_rdma_get_response_mode = mode;
89 2 }
90 2 SetIntFlagFromEnv(
91 "RECSTORE_RDMA_CONTROL_PLANE_PORT", &FLAGS_rdma_control_plane_port);
92 2 SetIntFlagFromEnv("RECSTORE_RDMA_CONTROL_PLANE_TIMEOUT_MS",
93 &FLAGS_rdma_control_plane_timeout_ms);
94 2 SetIntFlagFromEnv(
95 "RECSTORE_RDMA_WAIT_TIMEOUT_MS", &FLAGS_rdma_wait_timeout_ms);
96 2 SetIntFlagFromEnv("RECSTORE_RDMA_RC_QPS_PER_CLIENT_PER_SHARD",
97 &FLAGS_rdma_rc_qps_per_client_per_shard);
98 2 SetIntFlagFromEnv(
99 "RECSTORE_RDMA_RC_SLOTS_PER_QP", &FLAGS_rdma_rc_slots_per_qp);
100 2 SetIntFlagFromEnv("RECSTORE_RDMA_RC_SERVER_COROUTINES_PER_THREAD",
101 &FLAGS_rdma_rc_server_coroutines_per_thread);
102 2 SetIntFlagFromEnv(
103 "RECSTORE_RDMA_RC_SERVER_GET_WORKERS", &FLAGS_rdma_rc_server_get_workers);
104 2 }
105
106 ✗ std::int64_t NsSince(std::chrono::steady_clock::time_point start,
107 std::chrono::steady_clock::time_point end) {
108 ✗ return std::chrono::duration_cast<std::chrono::nanoseconds>(end - start)
109 ✗ .count();
110 }
111
112 ✗ int ValueSizeHintFromBaseKvConfig(const json& base_kv_config,
113 int fallback_value_size) {
114 ✗ if (!base_kv_config.is_object()) {
115 ✗ return fallback_value_size;
116 }
117 ✗ if (!base_kv_config.contains("value") ||
118 ✗ !base_kv_config["value"].is_object()) {
119 ✗ return fallback_value_size;
120 }
121 ✗ return base_kv_config["value"].value(
122 ✗ "default_value_size_hint", fallback_value_size);
123 }
124
125 ✗ std::vector<std::string> ReadProcessArgv() {
126 ✗ std::ifstream cmdline("/proc/self/cmdline", std::ios::binary);
127 ✗ std::vector<std::string> argv;
128 ✗ if (!cmdline.is_open()) {
129 ✗ return argv;
130 }
131
132 ✗ std::string current;
133 ✗ char ch = '\0';
134 ✗ while (cmdline.get(ch)) {
135 ✗ if (ch == '\0') {
136 ✗ if (!current.empty()) {
137 ✗ argv.push_back(current);
138 ✗ current.clear();
139 }
140 ✗ continue;
141 }
142 ✗ current.push_back(ch);
143 }
144 ✗ if (!current.empty()) {
145 ✗ argv.push_back(current);
146 }
147 ✗ return argv;
148 ✗ }
149 } // namespace
150
151 6 EmbeddedRdmaClientIdentity ResolveEmbeddedRdmaClientIdentity(int num_shards) {
152
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (num_shards <= 0) {
153 ✗ throw std::runtime_error("embedded RDMA num_shards must be positive");
154 }
155
156 6 int client_index = 0;
157
4/6
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
 10 if (!detail::TryParseIntEnv("RECSTORE_RDMA_OS_CLIENT_INDEX", &client_index) &&
158
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
 4 !detail::TryParseIntEnv("RANK", &client_index)) {
159 ✗ detail::TryParseIntEnv("LOCAL_RANK", &client_index);
160 }
161
162 6 int num_client_processes = 1;
163
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 if (!detail::TryParseIntEnv(
164
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 10 "RECSTORE_RDMA_NUM_CLIENT_PROCESSES", &num_client_processes) &&
165
2/4
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
 4 !detail::TryParseIntEnv("WORLD_SIZE", &num_client_processes)) {
166 ✗ detail::TryParseIntEnv("LOCAL_WORLD_SIZE", &num_client_processes);
167 }
168
169
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (client_index < 0) {
170 ✗ throw std::runtime_error("embedded RDMA client index must be non-negative");
171 }
172
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (num_client_processes <= 0) {
173 ✗ throw std::runtime_error(
174 ✗ "embedded RDMA num_client_processes must be positive");
175 }
176
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 if (client_index >= num_client_processes) {
177
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 throw std::runtime_error(
178 4 "embedded RDMA client index out of range for num_client_processes");
179 }
180
181 4 EmbeddedRdmaClientIdentity identity;
182 4 identity.client_index = client_index;
183 4 identity.num_client_processes = num_client_processes;
184 4 identity.global_id = num_shards + client_index;
185 4 return identity;
186 }
187
188 std::vector<RDMAPSClientAdapter::ShardChunk>
189 ✗ RDMAPSClientAdapter::BuildChunks(base::ConstArray<uint64_t> keys) const {
190 return shard_routing::BuildChunks(
191 keys,
192 ✗ num_shards_,
193 ✗ hash_method_,
194 ✗ shard_to_client_index_,
195 ✗ MaxGetKeysPerRpc());
196 }
197
198 ✗ void RDMAPSClientAdapter::WaitShardRpcsCooperatively(
199 const std::vector<PendingShardRpc>& shard_rpcs) {
200 ✗ std::vector<bool> finished(shard_rpcs.size(), false);
201 ✗ std::size_t remaining = shard_rpcs.size();
202 ✗ while (remaining > 0) {
203 ✗ bool made_progress = false;
204 ✗ for (std::size_t i = 0; i < shard_rpcs.size(); ++i) {
205 ✗ if (finished[i]) {
206 ✗ continue;
207 }
208 ✗ const auto& pending = shard_rpcs[i];
209 auto& client =
210 ✗ shard_clients_[static_cast<std::size_t>(pending.client_index)];
211 ✗ if (!client->QueryRPCFinished(pending.rpc_id)) {
212 ✗ continue;
213 }
214 ✗ client->WaitRPCFinish(pending.rpc_id);
215 ✗ finished[i] = true;
216 ✗ made_progress = true;
217 ✗ --remaining;
218 }
219 ✗ if (!made_progress) {
220 ✗ std::this_thread::yield();
221 }
222 }
223 ✗ }
224
225 2 void InitializeRdmaProcessRuntime() {
226 static std::once_flag init_once;
227
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 std::call_once(init_once, []() {
228 // Python entrypoints pass application CLI flags that are not gflags.
229 // Passing them to folly::init makes gflags abort before the RDMA client can
230 // start.
231
2/4
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
 8 std::vector<std::string> argv_strings = {"recstore_rdma_client"};
232 2 std::vector<char*> argv_storage;
233
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 argv_storage.reserve(argv_strings.size() + 1);
234
2/2
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 2 times.
 4 for (auto& arg : argv_strings) {
235
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 argv_storage.push_back(arg.data());
236 }
237
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 argv_storage.push_back(nullptr);
238
239 2 int argc = static_cast<int>(argv_strings.size());
240 2 char** argv = argv_storage.data();
241
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 folly::init(&argc, &argv);
242
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 ApplyRdmaFlagsFromEnv();
243 2 });
244 2 }
245
246 4 RDMAPSClientAdapter::RDMAPSClientAdapter(json config)
247
2/4
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 14 taken 4 times.
✗ Branch 15 not taken.
 4 : BasePSClient(config), config_(std::move(config)) {}
248
249 ✗ void RDMAPSClientAdapter::EnsureClientInitialized() {
250 ✗ std::lock_guard<std::mutex> guard(init_mu_);
251 ✗ if (initialized_) {
252 ✗ return;
253 }
254
255 const json cache_ps_cfg =
256 ✗ config_.contains("cache_ps") ? config_["cache_ps"] : json::object();
257 const json client_cfg =
258 ✗ config_.contains("client") ? config_["client"] : json::object();
259 ✗ const json dist_cfg = ResolveFrameworkDistributedClientConfig(config_);
260
261 ✗ num_shards_ = dist_cfg.value("num_shards", 1);
262 ✗ hash_method_ = dist_cfg.value("hash_method", "city_hash");
263 ✗ if (FLAGS_global_id < num_shards_) {
264 ✗ const auto identity = ResolveEmbeddedRdmaClientIdentity(num_shards_);
265 ✗ FLAGS_num_server_processes = num_shards_;
266 ✗ FLAGS_num_client_processes = identity.num_client_processes;
267 ✗ FLAGS_global_id = identity.global_id;
268 ✗ if (FLAGS_rdma_rc_num_logical_clients < 0) {
269 ✗ FLAGS_rdma_rc_num_logical_clients = identity.num_client_processes;
270 }
271 ✗ if (FLAGS_rdma_rc_client_id_base < 0) {
272 ✗ FLAGS_rdma_rc_client_id_base = identity.client_index;
273 }
274 ✗ } else if (FLAGS_num_server_processes != num_shards_) {
275 ✗ throw std::runtime_error(
276 ✗ "RDMA num_server_processes must match distributed_client.num_shards");
277 }
278 ✗ FLAGS_value_size =
279 ✗ cache_ps_cfg.contains("base_kv_config")
280 ✗ ? ValueSizeHintFromBaseKvConfig(
281 ✗ cache_ps_cfg["base_kv_config"], FLAGS_value_size)
282 : FLAGS_value_size;
283 ✗ FLAGS_max_kv_num_per_request =
284 ✗ dist_cfg.value("max_keys_per_request", FLAGS_max_kv_num_per_request);
285 ✗ if (const char* mode = std::getenv("RECSTORE_RDMA_TRANSPORT_MODE")) {
286 ✗ FLAGS_rdma_transport_mode = mode;
287 }
288
289 const int logical_client_id =
290 ✗ config_.value("rdma_logical_client_id", FLAGS_rdma_rc_client_id_base);
291
292 ✗ shard_clients_.clear();
293 ✗ shard_to_client_index_.clear();
294 ✗ client_ = nullptr;
295
296 ✗ if (num_shards_ <= 1) {
297 ✗ shard_clients_.push_back(std::make_unique<petps::PetPSClient>(
298 ✗ client_cfg.value("host", std::string("127.0.0.1")),
299 ✗ client_cfg.value("port", 25000),
300 ✗ client_cfg.value("shard", 0),
301 logical_client_id));
302 ✗ client_ = shard_clients_.front().get();
303 ✗ shard_to_client_index_[0] = 0;
304 } else {
305 ✗ const auto servers_it = dist_cfg.find("servers");
306 ✗ if (servers_it == dist_cfg.end() || !servers_it->is_array() ||
307 ✗ servers_it->empty()) {
308 ✗ throw std::runtime_error(
309 "RDMA distributed_client.servers must be provided for multi-shard "
310 ✗ "configuration");
311 }
312
313 ✗ CHECK_EQ(static_cast<int>(servers_it->size()), num_shards_)
314 ✗ << "RDMA distributed_client.servers size must equal num_shards";
315 ✗ for (const auto& server : *servers_it) {
316 ✗ const int shard = server.value("shard", -1);
317 ✗ if (shard < 0) {
318 ✗ throw std::runtime_error(
319 ✗ "RDMA distributed_client.servers[].shard must be explicit");
320 }
321 ✗ shard_clients_.push_back(std::make_unique<petps::PetPSClient>(
322 ✗ server.value("host", std::string("127.0.0.1")),
323 ✗ server.value("port", 25000),
324 shard,
325 logical_client_id));
326 ✗ shard_to_client_index_[shard] =
327 ✗ static_cast<int>(shard_clients_.size() - 1);
328 }
329 }
330
331 ✗ initialized_ = true;
332 ✗ }
333
334 ✗ void RDMAPSClientAdapter::EnsureThreadInitialized() {
335 ✗ EnsureClientInitialized();
336 ✗ const std::thread::id tid = std::this_thread::get_id();
337 ✗ std::lock_guard<std::mutex> guard(thread_init_mu_);
338 ✗ if (initialized_threads_.find(tid) != initialized_threads_.end()) {
339 ✗ return;
340 }
341
342 ✗ if (num_shards_ <= 1) {
343 ✗ if (client_ != nullptr) {
344 ✗ client_->InitThread();
345 }
346 } else {
347 ✗ for (auto& shard_client : shard_clients_) {
348 ✗ shard_client->InitThread();
349 }
350 }
351
352 ✗ initialized_threads_.insert(tid);
353 ✗ }
354
355 ✗ void RDMAPSClientAdapter::EnsureTableReady(const std::string& table_name,
356 int64_t embedding_dim) {
357 ✗ std::lock_guard<std::mutex> guard(state_mu_);
358 ✗ const auto it = tables_.find(table_name);
359 ✗ if (it == tables_.end()) {
360 ✗ throw std::runtime_error("RDMA table is not initialized: " + table_name);
361 }
362 ✗ if (static_cast<int64_t>(it->second.config.embedding_dim) != embedding_dim) {
363 ✗ throw std::runtime_error(
364 ✗ "RDMA embedding dimension mismatch for table " + table_name);
365 }
366 ✗ }
367
368 ✗ int64_t RDMAPSClientAdapter::DefaultEmbeddingDimOrThrow() const {
369 ✗ if (tables_.empty()) {
370 ✗ throw std::runtime_error(
371 ✗ "RDMA table metadata is empty; call InitEmbeddingTable first");
372 }
373 ✗ return static_cast<int64_t>(tables_.begin()->second.config.embedding_dim);
374 }
375
376 ✗ std::size_t RDMAPSClientAdapter::MaxGetKeysPerRpc() const {
377 ✗ const std::size_t response_limited = petps::GetKeysPerRpcByResponseBudget(
378 static_cast<std::size_t>(FLAGS_value_size),
379 static_cast<std::size_t>(FLAGS_rdma_rc_mtu_bytes),
380 ✗ static_cast<std::size_t>(FLAGS_rdma_rc_target_response_mtu));
381 const std::size_t request_limited =
382 ✗ petps::PutPayloadBudget(
383 ✗ static_cast<std::size_t>(FLAGS_rdma_rc_request_slot_bytes)) /
384 ✗ sizeof(std::uint64_t);
385 ✗ std::size_t limit = static_cast<std::size_t>(FLAGS_max_kv_num_per_request);
386 ✗ if (response_limited > 0) {
387 ✗ limit = std::min(limit, response_limited);
388 }
389 ✗ if (request_limited > 0) {
390 ✗ limit = std::min(limit, request_limited);
391 }
392 ✗ return std::max<std::size_t>(limit, 1);
393 }
394
395 ✗ std::size_t RDMAPSClientAdapter::MaxPutKeysPerRpc() const {
396 ✗ const std::size_t payload_budget = petps::PutPayloadBudget(
397 static_cast<std::size_t>(FLAGS_rdma_rc_request_slot_bytes));
398 const std::size_t embedding_dim =
399 ✗ static_cast<std::size_t>(DefaultEmbeddingDimOrThrow());
400 ✗ const std::size_t bytes_per_row =
401 ✗ sizeof(ParameterCompressItem) + embedding_dim * sizeof(float) +
402 sizeof(int);
403 ✗ std::size_t limit = static_cast<std::size_t>(FLAGS_max_kv_num_per_request);
404 ✗ if (payload_budget > sizeof(int) && bytes_per_row > 0) {
405 ✗ const std::size_t request_limited =
406 ✗ (payload_budget - sizeof(int)) / bytes_per_row;
407 ✗ if (request_limited > 0) {
408 ✗ limit = std::min(limit, request_limited);
409 }
410 }
411 ✗ return std::max<std::size_t>(limit, 1);
412 }
413
414 ✗ std::size_t RDMAPSClientAdapter::MaxInFlightGetRpcs() const {
415 const std::size_t qps = static_cast<std::size_t>(
416 ✗ std::max(FLAGS_rdma_rc_qps_per_client_per_shard, 1));
417 const std::size_t slots =
418 ✗ static_cast<std::size_t>(std::max(FLAGS_rdma_rc_slots_per_qp, 1));
419 ✗ return std::max<std::size_t>(qps * slots, 1);
420 }
421
422 RDMAPSClientAdapter::PrefetchState
423 ✗ RDMAPSClientAdapter::GetPrefetchState(uint64_t prefetch_id) {
424 ✗ std::lock_guard<std::mutex> guard(state_mu_);
425 ✗ const auto it = prefetches_.find(prefetch_id);
426 ✗ if (it == prefetches_.end()) {
427 ✗ throw std::runtime_error(
428 ✗ "Unknown RDMA prefetch id: " + std::to_string(prefetch_id));
429 }
430 ✗ return it->second;
431 ✗ }
432
433 ✗ void RDMAPSClientAdapter::MarkPrefetchConsumed(uint64_t prefetch_id) {
434 ✗ std::lock_guard<std::mutex> guard(state_mu_);
435 ✗ prefetches_.erase(prefetch_id);
436 ✗ }
437
438 ✗ bool RDMAPSClientAdapter::QueryRPCFinished(int rpc_id) {
439 ✗ if (rpc_id >= 0 && num_shards_ <= 1) {
440 ✗ return client_ != nullptr ? client_->QueryRPCFinished(rpc_id) : true;
441 }
442
443 ✗ std::lock_guard<std::mutex> guard(batches_mu_);
444 ✗ auto it = batches_.find(rpc_id);
445 ✗ CHECK(it != batches_.end());
446
447 ✗ for (const auto& pending : it->second.shard_rpcs) {
448 ✗ if (!shard_clients_[static_cast<std::size_t>(pending.client_index)]
449 ✗ ->QueryRPCFinished(pending.rpc_id)) {
450 ✗ return false;
451 }
452 }
453
454 ✗ return shard_routing::FinalizeBatchIfNeeded(&it->second, FLAGS_value_size);
455 ✗ }
456
457 ✗ void RDMAPSClientAdapter::WaitRPCFinish(int rpc_id) {
458 ✗ if (rpc_id >= 0 && num_shards_ <= 1) {
459 ✗ if (client_ != nullptr) {
460 ✗ client_->WaitRPCFinish(rpc_id);
461 }
462 ✗ return;
463 }
464
465 ✗ std::vector<PendingShardRpc> shard_rpcs;
466 {
467 ✗ std::lock_guard<std::mutex> guard(batches_mu_);
468 ✗ auto it = batches_.find(rpc_id);
469 ✗ CHECK(it != batches_.end());
470 ✗ if (it->second.assembled) {
471 ✗ return;
472 }
473 ✗ shard_rpcs = it->second.shard_rpcs;
474 ✗ }
475
476 ✗ WaitShardRpcsCooperatively(shard_rpcs);
477
478 {
479 ✗ std::lock_guard<std::mutex> guard(batches_mu_);
480 ✗ auto it = batches_.find(rpc_id);
481 ✗ CHECK(it != batches_.end());
482 ✗ shard_routing::FinalizeBatchIfNeeded(&it->second, FLAGS_value_size);
483 ✗ }
484 ✗ }
485
486 ✗ void RDMAPSClientAdapter::RevokeRPCResource(int rpc_id) {
487 ✗ if (rpc_id >= 0 && num_shards_ <= 1) {
488 ✗ if (client_ != nullptr) {
489 ✗ client_->RevokeRPCResource(rpc_id);
490 }
491 ✗ return;
492 }
493
494 ✗ std::lock_guard<std::mutex> guard(batches_mu_);
495 ✗ auto it = batches_.find(rpc_id);
496 ✗ CHECK(it != batches_.end());
497
498 ✗ for (const auto& pending : it->second.shard_rpcs) {
499 ✗ shard_clients_[static_cast<std::size_t>(pending.client_index)]
500 ✗ ->RevokeRPCResource(pending.rpc_id);
501 }
502
503 ✗ batches_.erase(it);
504 ✗ }
505
506 ✗ const float* RDMAPSClientAdapter::BorrowPrefetchResult(
507 const PrefetchState& state,
508 std::int32_t* status_code,
509 std::size_t* response_bytes) {
510 ✗ if (!state.borrowed_response || client_ == nullptr) {
511 ✗ return nullptr;
512 }
513 ✗ auto* pet_client = dynamic_cast<petps::PetPSClient*>(client_);
514 ✗ if (pet_client == nullptr) {
515 ✗ return nullptr;
516 }
517 ✗ std::size_t key_count = 0;
518 ✗ const float* payload = pet_client->BorrowGetResultPayload(
519 ✗ state.rpc_id, &key_count, response_bytes, status_code);
520 ✗ if (payload == nullptr ||
521 ✗ key_count != static_cast<std::size_t>(state.key_count)) {
522 ✗ return nullptr;
523 }
524 ✗ return payload;
525 }
526
527 ✗ int RDMAPSClientAdapter::SubmitGetParameter(
528 base::ConstArray<uint64_t> keys,
529 float* values,
530 bool isAsync,
531 int async_req_id) {
532 ✗ EnsureThreadInitialized();
533 ✗ if (keys.Size() == 0) {
534 ✗ auto* status =
535 reinterpret_cast<std::int32_t*>(reinterpret_cast<char*>(values));
536 ✗ *status = static_cast<std::int32_t>(petps::RpcStatus::kOk);
537 ✗ return 0;
538 }
539
540 ✗ BatchRequest batch;
541 ✗ batch.user_buffer = values;
542 ✗ batch.total_key_count = keys.Size();
543 auto* batch_status_word =
544 ✗ petps::FixedSlotStatusWord(values, keys.Size(), FLAGS_value_size);
545 ✗ *batch_status_word = static_cast<std::int32_t>(petps::RpcStatus::kPending);
546
547 ✗ if (num_shards_ <= 1) {
548 ✗ if (client_ == nullptr) {
549 ✗ return -1;
550 }
551 ✗ const std::size_t max_keys_per_rpc = MaxGetKeysPerRpc();
552 ✗ const std::size_t max_in_flight = MaxInFlightGetRpcs();
553 ✗ const std::size_t total_keys = keys.Size();
554 ✗ if (total_keys <= max_keys_per_rpc) {
555 ✗ return client_->GetParameter(keys, values, isAsync, async_req_id);
556 }
557 ✗ std::vector<PendingShardRpc> window;
558 ✗ window.reserve(max_in_flight);
559 ✗ auto drain_and_release_window = [this, &window, &batch]() {
560 // Large model batches can split into more GET RPCs than the RC slot pool.
561 // Keep submission bounded by waiting and freeing each window before
562 // acquiring more slots.
563 ✗ for (const auto& pending : window) {
564 ✗ client_->WaitRPCFinish(pending.rpc_id);
565 }
566 ✗ for (const auto& pending : window) {
567 ✗ batch.shard_rpcs.push_back(pending);
568 ✗ client_->RevokeRPCResource(pending.rpc_id);
569 }
570 ✗ window.clear();
571 ✗ };
572 ✗ for (std::size_t offset = 0; offset < total_keys;
573 ✗ offset += max_keys_per_rpc) {
574 ✗ const std::size_t end = std::min(offset + max_keys_per_rpc, total_keys);
575 ✗ std::vector<uint64_t> key_slice;
576 ✗ key_slice.reserve(end - offset);
577 ✗ std::vector<std::size_t> positions;
578 ✗ positions.reserve(end - offset);
579 ✗ for (std::size_t i = offset; i < end; ++i) {
580 ✗ key_slice.push_back(keys[i]);
581 ✗ positions.push_back(i);
582 }
583 ✗ void* recv = client_->GetReceiveBuffer(
584 ✗ key_slice.size() * static_cast<std::size_t>(FLAGS_value_size) +
585 sizeof(std::int32_t));
586 ✗ const int rpc_id = client_->GetParameter(
587 base::ConstArray<uint64_t>(key_slice),
588 static_cast<float*>(recv),
589 isAsync,
590 async_req_id);
591 ✗ window.push_back(PendingShardRpc{
592 0,
593 0,
594 rpc_id,
595 ✗ std::move(positions),
596 recv,
597 ✗ key_slice.size(),
598 });
599 ✗ if (window.size() >= max_in_flight) {
600 ✗ drain_and_release_window();
601 }
602 ✗ }
603 ✗ if (!window.empty()) {
604 ✗ drain_and_release_window();
605 }
606 ✗ } else {
607 ✗ const std::size_t max_in_flight = MaxInFlightGetRpcs();
608 ✗ std::vector<PendingShardRpc> window;
609 ✗ window.reserve(max_in_flight);
610 ✗ auto drain_and_release_window = [this, &window, &batch]() {
611 ✗ WaitShardRpcsCooperatively(window);
612 ✗ for (const auto& pending : window) {
613 ✗ batch.shard_rpcs.push_back(pending);
614 ✗ shard_clients_[static_cast<std::size_t>(pending.client_index)]
615 ✗ ->RevokeRPCResource(pending.rpc_id);
616 }
617 ✗ window.clear();
618 ✗ };
619 ✗ for (const auto& chunk : BuildChunks(keys)) {
620 ✗ BaseParameterClient* client = shard_clients_[chunk.client_index].get();
621 ✗ void* recv = client->GetReceiveBuffer(
622 ✗ chunk.keys.size() * static_cast<std::size_t>(FLAGS_value_size) +
623 sizeof(std::int32_t));
624 ✗ const int rpc_id = client->GetParameter(
625 ✗ base::ConstArray<uint64_t>(chunk.keys),
626 static_cast<float*>(recv),
627 isAsync,
628 async_req_id);
629 ✗ window.push_back(PendingShardRpc{
630 ✗ chunk.shard_id,
631 ✗ chunk.client_index,
632 rpc_id,
633 ✗ chunk.positions,
634 recv,
635 ✗ chunk.keys.size(),
636 });
637 ✗ if (window.size() >= max_in_flight) {
638 ✗ drain_and_release_window();
639 }
640 ✗ }
641 ✗ if (!window.empty()) {
642 ✗ drain_and_release_window();
643 }
644 ✗ }
645
646 ✗ int batch_id = 0;
647 {
648 ✗ std::lock_guard<std::mutex> guard(batches_mu_);
649 ✗ batch_id = batch_rpc_id_acc_--;
650 ✗ if (batch_id >= 0) {
651 ✗ throw std::runtime_error("rdma batch rpc id exhausted negative range");
652 }
653 ✗ batches_[batch_id] = std::move(batch);
654 ✗ }
655 ✗ if (!isAsync) {
656 ✗ WaitRPCFinish(batch_id);
657 }
658 ✗ return batch_id;
659 ✗ }
660
661 ✗ int RDMAPSClientAdapter::GetParameter(const base::ConstArray<uint64_t>& keys,
662 float* values) {
663 ✗ EnsureThreadInitialized();
664 ✗ if (keys.Size() == 0) {
665 ✗ return 0;
666 }
667
668 const std::size_t response_bytes =
669 ✗ petps::FixedSlotResponseBytes(keys.Size(), FLAGS_value_size);
670 ✗ float* recv = nullptr;
671 ✗ if (num_shards_ > 1) {
672 ✗ if (shard_clients_.empty()) {
673 ✗ return -1;
674 }
675 recv = static_cast<float*>(
676 ✗ shard_clients_.front()->GetReceiveBuffer(response_bytes));
677 } else {
678 ✗ recv = static_cast<float*>(client_->GetReceiveBuffer(response_bytes));
679 }
680
681 ✗ const int rpc_id = SubmitGetParameter(keys, recv, false, 0);
682 ✗ WaitRPCFinish(rpc_id);
683 const auto* status_word =
684 ✗ petps::FixedSlotStatusWord(recv, keys.Size(), FLAGS_value_size);
685 ✗ if (*status_word != static_cast<std::int32_t>(petps::RpcStatus::kOk)) {
686 ✗ RevokeRPCResource(rpc_id);
687 ✗ return -1;
688 }
689
690 ✗ std::memcpy(
691 ✗ values, recv, keys.Size() * static_cast<std::size_t>(FLAGS_value_size));
692 ✗ RevokeRPCResource(rpc_id);
693 ✗ return 0;
694 }
695
696 ✗ int RDMAPSClientAdapter::PutParameter(
697 const base::ConstArray<uint64_t>& keys,
698 const std::vector<std::vector<float>>& values) {
699 ✗ EnsureThreadInitialized();
700 ✗ if (keys.Size() != values.size()) {
701 ✗ return -1;
702 }
703 ✗ const std::size_t max_keys_per_rpc = MaxPutKeysPerRpc();
704 ✗ if (num_shards_ <= 1) {
705 ✗ if (client_ == nullptr) {
706 ✗ return -1;
707 }
708 ✗ const std::size_t key_count = static_cast<std::size_t>(keys.Size());
709 ✗ for (std::size_t offset = 0; offset < key_count;
710 ✗ offset += max_keys_per_rpc) {
711 ✗ const std::size_t end = std::min(offset + max_keys_per_rpc, key_count);
712 ✗ std::vector<uint64_t> key_slice;
713 ✗ key_slice.reserve(end - offset);
714 ✗ for (std::size_t i = offset; i < end; ++i) {
715 ✗ key_slice.push_back(keys[i]);
716 }
717 std::vector<std::vector<float>> value_slice(
718 ✗ values.begin() + static_cast<std::ptrdiff_t>(offset),
719 ✗ values.begin() + static_cast<std::ptrdiff_t>(end));
720 ✗ const int rc = client_->PutParameter(key_slice, value_slice);
721 ✗ if (rc != 0) {
722 ✗ return rc;
723 }
724 ✗ }
725 ✗ return 0;
726 }
727 ✗ if (keys.Size() == 0) {
728 ✗ return 0;
729 }
730
731 ✗ std::vector<std::vector<uint64_t>> shard_keys(num_shards_);
732 ✗ std::vector<std::vector<std::vector<float>>> shard_values(num_shards_);
733
734 ✗ for (std::size_t i = 0; i < keys.Size(); ++i) {
735 const int shard =
736 ✗ shard_routing::PartitionKey(keys[i], num_shards_, hash_method_);
737 ✗ shard_keys[static_cast<std::size_t>(shard)].push_back(keys[i]);
738 ✗ shard_values[static_cast<std::size_t>(shard)].push_back(values[i]);
739 }
740
741 ✗ for (int shard = 0; shard < num_shards_; ++shard) {
742 ✗ const int client_index = shard_to_client_index_.at(shard);
743 ✗ for (std::size_t offset = 0;
744 ✗ offset < shard_keys[static_cast<std::size_t>(shard)].size();
745 ✗ offset += max_keys_per_rpc) {
746 const std::size_t end =
747 ✗ std::min(offset + max_keys_per_rpc,
748 ✗ shard_keys[static_cast<std::size_t>(shard)].size());
749 std::vector<uint64_t> key_slice(
750 ✗ shard_keys[static_cast<std::size_t>(shard)].begin() + offset,
751 ✗ shard_keys[static_cast<std::size_t>(shard)].begin() + end);
752 std::vector<std::vector<float>> value_slice(
753 ✗ shard_values[static_cast<std::size_t>(shard)].begin() + offset,
754 ✗ shard_values[static_cast<std::size_t>(shard)].begin() + end);
755 int rc =
756 ✗ shard_clients_[static_cast<std::size_t>(client_index)]->PutParameter(
757 key_slice, value_slice);
758 ✗ if (rc != 0) {
759 ✗ return rc;
760 }
761 ✗ }
762 }
763 ✗ return 0;
764 ✗ }
765
766 ✗ int RDMAPSClientAdapter::UpdateParameter(
767 const std::string& table_name,
768 const base::ConstArray<uint64_t>& keys,
769 const std::vector<std::vector<float>>* grads) {
770 ✗ if (grads == nullptr) {
771 ✗ return -1;
772 }
773 ✗ if (grads->empty()) {
774 ✗ return 0;
775 }
776 ✗ EnsureThreadInitialized();
777 ✗ const std::size_t max_keys_per_rpc = MaxPutKeysPerRpc();
778 ✗ if (num_shards_ <= 1) {
779 ✗ if (client_ == nullptr) {
780 ✗ return -1;
781 }
782 ✗ const std::size_t key_count = static_cast<std::size_t>(keys.Size());
783 ✗ for (std::size_t offset = 0; offset < key_count;
784 ✗ offset += max_keys_per_rpc) {
785 ✗ const std::size_t end = std::min(offset + max_keys_per_rpc, key_count);
786 ✗ std::vector<uint64_t> key_slice;
787 ✗ key_slice.reserve(end - offset);
788 ✗ for (std::size_t i = offset; i < end; ++i) {
789 ✗ key_slice.push_back(keys[i]);
790 }
791 std::vector<std::vector<float>> grad_slice(
792 ✗ grads->begin() + static_cast<std::ptrdiff_t>(offset),
793 ✗ grads->begin() + static_cast<std::ptrdiff_t>(end));
794 ✗ const int rc = client_->UpdateParameter(
795 table_name, base::ConstArray<uint64_t>(key_slice), &grad_slice);
796 ✗ if (rc != 0) {
797 ✗ return rc;
798 }
799 ✗ }
800 ✗ return 0;
801 }
802 ✗ if (keys.Size() != grads->size()) {
803 ✗ return -1;
804 }
805 ✗ if (keys.Size() == 0) {
806 ✗ return 0;
807 }
808
809 ✗ std::vector<std::vector<uint64_t>> shard_keys(num_shards_);
810 ✗ std::vector<std::vector<std::vector<float>>> shard_grads(num_shards_);
811
812 ✗ for (std::size_t i = 0; i < keys.Size(); ++i) {
813 const int shard =
814 ✗ shard_routing::PartitionKey(keys[i], num_shards_, hash_method_);
815 ✗ shard_keys[static_cast<std::size_t>(shard)].push_back(keys[i]);
816 ✗ shard_grads[static_cast<std::size_t>(shard)].push_back((*grads)[i]);
817 }
818
819 ✗ for (int shard = 0; shard < num_shards_; ++shard) {
820 ✗ if (shard_keys[static_cast<std::size_t>(shard)].empty()) {
821 ✗ continue;
822 }
823 ✗ const int client_index = shard_to_client_index_.at(shard);
824 ✗ for (std::size_t offset = 0;
825 ✗ offset < shard_keys[static_cast<std::size_t>(shard)].size();
826 ✗ offset += max_keys_per_rpc) {
827 const std::size_t end =
828 ✗ std::min(offset + max_keys_per_rpc,
829 ✗ shard_keys[static_cast<std::size_t>(shard)].size());
830 std::vector<uint64_t> key_slice(
831 ✗ shard_keys[static_cast<std::size_t>(shard)].begin() + offset,
832 ✗ shard_keys[static_cast<std::size_t>(shard)].begin() + end);
833 std::vector<std::vector<float>> grad_slice(
834 ✗ shard_grads[static_cast<std::size_t>(shard)].begin() + offset,
835 ✗ shard_grads[static_cast<std::size_t>(shard)].begin() + end);
836 const int rc =
837 ✗ shard_clients_[static_cast<std::size_t>(client_index)]
838 ✗ ->UpdateParameter(table_name,
839 base::ConstArray<uint64_t>(key_slice),
840 &grad_slice);
841 ✗ if (rc != 0) {
842 ✗ return rc;
843 }
844 ✗ }
845 }
846 ✗ return 0;
847 ✗ }
848
849 ✗ int RDMAPSClientAdapter::UpdateParameterFlat(
850 const std::string& table_name,
851 const base::ConstArray<uint64_t>& keys,
852 const float* grads,
853 int64_t num_rows,
854 int64_t embedding_dim) {
855 ✗ const uint64_t update_id = SubmitUpdateParameterFlatAsync(
856 table_name, keys, grads, num_rows, embedding_dim);
857 ✗ return WaitUpdateParameterFlat(update_id);
858 }
859
860 ✗ uint64_t RDMAPSClientAdapter::SubmitUpdateParameterFlatAsync(
861 const std::string& table_name,
862 const base::ConstArray<uint64_t>& keys,
863 const float* grads,
864 int64_t num_rows,
865 int64_t embedding_dim) {
866 ✗ EnsureTableReady(table_name, embedding_dim);
867 ✗ if (num_rows < 0 || (num_rows > 0 && grads == nullptr)) {
868 ✗ throw std::invalid_argument("RDMA update has invalid rows or gradients");
869 }
870 ✗ if (keys.Size() != static_cast<std::size_t>(num_rows)) {
871 ✗ throw std::invalid_argument("RDMA update key and gradient rows differ");
872 }
873 ✗ EnsureThreadInitialized();
874 ✗ const std::size_t max_keys_per_rpc = MaxPutKeysPerRpc();
875 ✗ const std::size_t dim = static_cast<std::size_t>(embedding_dim);
876 ✗ std::vector<std::pair<int, int>> pending;
877
878 ✗ auto wait_pending = [this, &pending]() {
879 ✗ int result = 0;
880 ✗ for (const auto& [client_index, rpc_id] : pending) {
881 ✗ if (shard_clients_[static_cast<std::size_t>(client_index)]
882 ✗ ->WaitUpdateParameter(rpc_id) != 0) {
883 ✗ result = -1;
884 }
885 }
886 ✗ pending.clear();
887 ✗ return result;
888 ✗ };
889
890 try {
891 ✗ if (num_shards_ <= 1) {
892 ✗ for (std::size_t offset = 0; offset < keys.Size();
893 ✗ offset += max_keys_per_rpc) {
894 const std::size_t count =
895 ✗ std::min(max_keys_per_rpc, keys.Size() - offset);
896 ✗ const int rpc_id = shard_clients_.front()->SubmitUpdateParameterFlat(
897 table_name,
898 ✗ base::ConstArray<uint64_t>(keys.Data() + offset, count),
899 ✗ grads + offset * dim,
900 ✗ dim);
901 ✗ if (rpc_id < 0) {
902 ✗ throw std::runtime_error("Failed to submit RDMA embedding update");
903 }
904 ✗ pending.emplace_back(0, rpc_id);
905 ✗ if (pending.size() >= MaxInFlightGetRpcs() && wait_pending() != 0) {
906 ✗ throw std::runtime_error("RDMA embedding update failed");
907 }
908 }
909 } else {
910 ✗ std::vector<std::vector<std::size_t>> shard_rows(num_shards_);
911 const std::size_t rows_per_shard =
912 ✗ (keys.Size() + static_cast<std::size_t>(num_shards_) - 1) /
913 ✗ static_cast<std::size_t>(num_shards_);
914 ✗ for (auto& rows : shard_rows) {
915 ✗ rows.reserve(rows_per_shard);
916 }
917 ✗ for (std::size_t row = 0; row < keys.Size(); ++row) {
918 const int shard =
919 ✗ shard_routing::PartitionKey(keys[row], num_shards_, hash_method_);
920 ✗ shard_rows[static_cast<std::size_t>(shard)].push_back(row);
921 }
922
923 ✗ pending.reserve(static_cast<std::size_t>(num_shards_));
924 ✗ for (int shard = 0; shard < num_shards_; ++shard) {
925 const auto& rows_for_shard =
926 ✗ shard_rows[static_cast<std::size_t>(shard)];
927 ✗ const int client_index = shard_to_client_index_.at(shard);
928 ✗ for (std::size_t offset = 0; offset < rows_for_shard.size();
929 ✗ offset += max_keys_per_rpc) {
930 const std::size_t count =
931 ✗ std::min(max_keys_per_rpc, rows_for_shard.size() - offset);
932 const int rpc_id =
933 ✗ shard_clients_[static_cast<std::size_t>(client_index)]
934 ✗ ->SubmitUpdateParameterFlatGather(
935 table_name,
936 keys.Data(),
937 grads,
938 ✗ keys.Size(),
939 dim,
940 ✗ rows_for_shard.data() + offset,
941 ✗ count);
942 ✗ if (rpc_id < 0) {
943 ✗ throw std::runtime_error("Failed to submit sharded RDMA update");
944 }
945 ✗ pending.emplace_back(client_index, rpc_id);
946 ✗ if (pending.size() >= MaxInFlightGetRpcs() && wait_pending() != 0) {
947 ✗ throw std::runtime_error("Sharded RDMA embedding update failed");
948 }
949 }
950 }
951 ✗ }
952 ✗ } catch (...) {
953 ✗ wait_pending();
954 ✗ throw;
955 ✗ }
956
957 ✗ std::lock_guard<std::mutex> guard(state_mu_);
958 ✗ const uint64_t update_id = next_update_id_++;
959 ✗ pending_updates_.emplace(
960 update_id,
961 ✗ PendingUpdate{std::move(pending), std::this_thread::get_id()});
962 ✗ return update_id;
963 ✗ }
964
965 ✗ int RDMAPSClientAdapter::WaitUpdateParameterFlat(uint64_t update_id) {
966 ✗ PendingUpdate update;
967 {
968 ✗ std::lock_guard<std::mutex> guard(state_mu_);
969 ✗ const auto it = pending_updates_.find(update_id);
970 ✗ if (it == pending_updates_.end()) {
971 ✗ throw std::runtime_error("Unknown or already consumed RDMA update handle");
972 }
973 ✗ if (it->second.owner != std::this_thread::get_id()) {
974 ✗ throw std::runtime_error(
975 ✗ "RDMA update handle must be waited by its submitting thread");
976 }
977 ✗ update = std::move(it->second);
978 ✗ pending_updates_.erase(it);
979 ✗ }
980
981 ✗ int result = 0;
982 ✗ for (const auto& [client_index, rpc_id] : update.shard_rpcs) {
983 ✗ if (shard_clients_[static_cast<std::size_t>(client_index)]
984 ✗ ->WaitUpdateParameter(rpc_id) != 0) {
985 ✗ result = -1;
986 }
987 }
988 ✗ return result;
989 ✗ }
990
991 ✗ int RDMAPSClientAdapter::InitEmbeddingTable(
992 const std::string& table_name, const EmbeddingTableConfig& config) {
993 ✗ EnsureThreadInitialized();
994 ✗ if (num_shards_ <= 1) {
995 ✗ if (client_ == nullptr) {
996 ✗ return -1;
997 }
998 ✗ const int init_rc = client_->InitEmbeddingTable(
999 ✗ table_name, config.num_embeddings, config.embedding_dim);
1000 ✗ if (init_rc != 0) {
1001 ✗ return init_rc;
1002 }
1003 } else {
1004 ✗ for (auto& shard_client : shard_clients_) {
1005 ✗ const int rc = shard_client->InitEmbeddingTable(
1006 ✗ table_name, config.num_embeddings, config.embedding_dim);
1007 ✗ if (rc != 0) {
1008 ✗ return rc;
1009 }
1010 }
1011 }
1012
1013 ✗ std::lock_guard<std::mutex> guard(state_mu_);
1014 ✗ const auto [it, inserted] = tables_.emplace(table_name, TableState{config});
1015 ✗ if (!inserted) {
1016 ✗ if (it->second.config.embedding_dim != config.embedding_dim ||
1017 ✗ it->second.config.num_embeddings != config.num_embeddings) {
1018 ✗ return -1;
1019 }
1020 }
1021 ✗ return 0;
1022 ✗ }
1023
1024 ✗ int RDMAPSClientAdapter::AsyncGetParameter(const base::ConstArray<uint64_t>&,
1025 float*) {
1026 ✗ throw std::runtime_error(
1027 ✗ "RDMA adapter AsyncGetParameter not implemented yet");
1028 }
1029
1030 ✗ void RDMAPSClientAdapter::Command(PSCommand) {
1031 ✗ EnsureThreadInitialized();
1032 ✗ if (num_shards_ <= 1) {
1033 ✗ if (client_ == nullptr) {
1034 ✗ throw std::runtime_error("RDMA adapter has no initialized client");
1035 }
1036 ✗ client_->Barrier("rdma_command", 0);
1037 ✗ return;
1038 }
1039 ✗ if (shard_clients_.empty()) {
1040 ✗ throw std::runtime_error("RDMA adapter has no initialized clients");
1041 }
1042 ✗ shard_clients_.front()->Barrier("rdma_command", 0);
1043 }
1044
1045 uint64_t
1046 ✗ RDMAPSClientAdapter::PrefetchParameter(const base::ConstArray<uint64_t>& keys) {
1047 ✗ EnsureThreadInitialized();
1048 ✗ if (keys.Size() == 0) {
1049 ✗ throw std::invalid_argument("RDMA prefetch requires at least one key");
1050 }
1051
1052 ✗ const int64_t embedding_dim = DefaultEmbeddingDimOrThrow();
1053 const std::size_t response_bytes =
1054 ✗ petps::FixedSlotResponseBytes(keys.Size(), FLAGS_value_size);
1055 const bool borrow_single_shard_response =
1056 ✗ num_shards_ <= 1 && keys.Size() <= MaxGetKeysPerRpc();
1057 ✗ const bool batch_response = !borrow_single_shard_response;
1058 auto buffer = std::make_shared<std::vector<float>>(
1059 ✗ response_bytes / sizeof(float));
1060 ✗ auto* status_word = petps::FixedSlotStatusWord(
1061 ✗ buffer->data(), static_cast<std::size_t>(keys.Size()), FLAGS_value_size);
1062 ✗ *status_word = static_cast<std::int32_t>(petps::RpcStatus::kPending);
1063
1064 ✗ const int rpc_id = SubmitGetParameter(keys, buffer->data(), true, 0);
1065
1066 ✗ std::lock_guard<std::mutex> guard(state_mu_);
1067 ✗ const uint64_t prefetch_id = next_prefetch_id_++;
1068 ✗ prefetches_.emplace(
1069 prefetch_id,
1070 ✗ PrefetchState{
1071 buffer,
1072 rpc_id,
1073 ✗ static_cast<int64_t>(keys.Size()),
1074 embedding_dim,
1075 borrow_single_shard_response,
1076 batch_response,
1077 });
1078 ✗ return prefetch_id;
1079 ✗ }
1080
1081 ✗ bool RDMAPSClientAdapter::IsPrefetchDone(uint64_t prefetch_id) {
1082 ✗ EnsureThreadInitialized();
1083 ✗ const PrefetchState state = GetPrefetchState(prefetch_id);
1084 ✗ return QueryRPCFinished(state.rpc_id);
1085 ✗ }
1086
1087 ✗ void RDMAPSClientAdapter::WaitForPrefetch(uint64_t prefetch_id) {
1088 ✗ EnsureThreadInitialized();
1089 ✗ const PrefetchState state = GetPrefetchState(prefetch_id);
1090 try {
1091 ✗ WaitRPCFinish(state.rpc_id);
1092 ✗ } catch (...) {
1093 ✗ RevokeRPCResource(state.rpc_id);
1094 ✗ MarkPrefetchConsumed(prefetch_id);
1095 ✗ throw;
1096 ✗ }
1097 ✗ }
1098
1099 ✗ bool RDMAPSClientAdapter::GetPrefetchResult(
1100 uint64_t prefetch_id, std::vector<std::vector<float>>* values) {
1101 ✗ if (values == nullptr) {
1102 ✗ return false;
1103 }
1104
1105 ✗ const PrefetchState state = GetPrefetchState(prefetch_id);
1106 ✗ std::vector<float> flat;
1107 ✗ int64_t num_rows = 0;
1108 ✗ if (!GetPrefetchResultFlat(
1109 ✗ prefetch_id, &flat, &num_rows, state.embedding_dim)) {
1110 ✗ return false;
1111 }
1112
1113 ✗ petps::CopyFlatRowsToVectors(
1114 ✗ flat.data(),
1115 static_cast<std::size_t>(num_rows),
1116 ✗ static_cast<std::size_t>(state.embedding_dim),
1117 values);
1118 ✗ return true;
1119 ✗ }
1120
1121 ✗ bool RDMAPSClientAdapter::GetPrefetchResultFlat(
1122 uint64_t prefetch_id,
1123 std::vector<float>* values,
1124 int64_t* num_rows,
1125 int64_t embedding_dim) {
1126 ✗ if (values == nullptr || num_rows == nullptr) {
1127 ✗ return false;
1128 }
1129
1130 ✗ const PrefetchState state = GetPrefetchState(prefetch_id);
1131 ✗ if (embedding_dim != state.embedding_dim) {
1132 ✗ return false;
1133 }
1134
1135 ✗ const bool profile_enabled = AdapterProfileEnabled();
1136 ✗ const auto wait_begin = std::chrono::steady_clock::now();
1137 ✗ std::int32_t status_code = static_cast<std::int32_t>(petps::RpcStatus::kOk);
1138 ✗ std::size_t response_bytes = 0;
1139 const float* result_payload =
1140 ✗ BorrowPrefetchResult(state, &status_code, &response_bytes);
1141 ✗ if (result_payload == nullptr) {
1142 ✗ WaitForPrefetch(prefetch_id);
1143 ✗ const auto* status_word = petps::FixedSlotStatusWord(
1144 ✗ state.buffer->data(),
1145 ✗ static_cast<std::size_t>(state.key_count),
1146 FLAGS_value_size);
1147 ✗ status_code = *status_word;
1148 ✗ response_bytes = static_cast<std::size_t>(state.key_count) *
1149 ✗ static_cast<std::size_t>(FLAGS_value_size);
1150 ✗ result_payload = state.buffer->data();
1151 }
1152 ✗ const auto wait_end = std::chrono::steady_clock::now();
1153 ✗ if (status_code != static_cast<std::int32_t>(petps::RpcStatus::kOk)) {
1154 ✗ RevokeRPCResource(state.rpc_id);
1155 ✗ MarkPrefetchConsumed(prefetch_id);
1156 ✗ return false;
1157 }
1158
1159 ✗ const std::size_t value_count =
1160 ✗ static_cast<std::size_t>(state.key_count) *
1161 ✗ static_cast<std::size_t>(state.embedding_dim);
1162 ✗ const auto assign_begin = std::chrono::steady_clock::now();
1163 ✗ if (FLAGS_rdma_adapter_skip_prefetch_result_copy) {
1164 ✗ values->clear();
1165 ✗ } else if (response_bytes == 0) {
1166 ✗ values->clear();
1167 ✗ } else if (result_payload == state.buffer->data()) {
1168 ✗ state.buffer->resize(value_count);
1169 ✗ values->swap(*state.buffer);
1170 } else {
1171 ✗ values->resize(value_count);
1172 ✗ if (value_count > 0) {
1173 ✗ const std::size_t expected_bytes = value_count * sizeof(values->front());
1174 ✗ if (response_bytes < expected_bytes) {
1175 ✗ RevokeRPCResource(state.rpc_id);
1176 ✗ MarkPrefetchConsumed(prefetch_id);
1177 ✗ return false;
1178 }
1179 ✗ std::memcpy(values->data(), result_payload, expected_bytes);
1180 }
1181 }
1182 ✗ const auto assign_end = std::chrono::steady_clock::now();
1183 ✗ *num_rows = state.key_count;
1184 ✗ const auto revoke_begin = std::chrono::steady_clock::now();
1185 ✗ RevokeRPCResource(state.rpc_id);
1186 ✗ MarkPrefetchConsumed(prefetch_id);
1187 ✗ const auto revoke_end = std::chrono::steady_clock::now();
1188 ✗ if (profile_enabled) {
1189 static std::atomic<std::uint64_t> count{0};
1190 static std::atomic<std::uint64_t> wait_ns{0};
1191 static std::atomic<std::uint64_t> assign_ns{0};
1192 static std::atomic<std::uint64_t> revoke_ns{0};
1193 ✗ const std::uint64_t current = count.fetch_add(1) + 1;
1194 ✗ wait_ns.fetch_add(
1195 ✗ static_cast<std::uint64_t>(NsSince(wait_begin, wait_end)));
1196 ✗ assign_ns.fetch_add(
1197 ✗ static_cast<std::uint64_t>(NsSince(assign_begin, assign_end)));
1198 ✗ revoke_ns.fetch_add(
1199 ✗ static_cast<std::uint64_t>(NsSince(revoke_begin, revoke_end)));
1200 ✗ if (current == 1 || current % 512 == 0) {
1201 ✗ const double denom = static_cast<double>(current);
1202 std::cout
1203 << "component=rdma_adapter_prefetch_profile"
1204 ✗ << " batches=" << current
1205 ✗ << " wait_avg_ns=" << static_cast<double>(wait_ns.load()) / denom
1206 ✗ << " assign_avg_ns=" << static_cast<double>(assign_ns.load()) / denom
1207 ✗ << " revoke_avg_ns=" << static_cast<double>(revoke_ns.load()) / denom
1208 ✗ << " value_count=" << value_count << std::endl;
1209 }
1210 }
1211 ✗ return true;
1212 ✗ }
1213
1214 } // namespace recstore
1215