GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 90.3% 56 / 0 / 62
Functions: 100.0% 3 / 0 / 3
Branches: 55.8% 29 / 0 / 52
ps/rdma/shard_routing.h
Line Branch Exec Source
1 #pragma once
2
3 // Shared shard-routing primitives for RDMA parameter clients.
4 //
5 // Both AllShardsParameterClientWrapper (allshards_ps_client.h) and
6 // RDMAPSClientAdapter (rdma_ps_client_adapter.h) fan a batch of keys out to
7 // per-shard RPCs and reassemble the responses into a single caller buffer.
8 // The structs and the partition/finalize logic below are identical between the
9 // two, so they live here instead of being copy-pasted.
10
11 #include <algorithm>
12 #include <cstddef>
13 #include <cstdint>
14 #include <cstring>
15 #include <stdexcept>
16 #include <string>
17 #include <unordered_map>
18 #include <vector>
19
20 #include "base/array.h"
21 #include "base/hash.h"
22 #include "base/log.h"
23 #include "ps/rdma/rdma_common.h"
24 #include "ps/rdma/rdma_status.h"
25
26 namespace recstore {
27 namespace shard_routing {
28
29 // One in-flight per-shard GET RPC and where its rows land in the caller batch.
30 struct PendingShardRpc {
31 int shard_id = 0; // Logical shard this RPC belongs to.
32 int client_index = 0; // Underlying client selected for this shard.
33 int rpc_id = -1; // RPC id returned by the shard-local client.
34 std::vector<std::size_t> original_positions; // Positions in caller's batch.
35 void* recv_buffer = nullptr; // Caller-visible response buffer.
36 std::size_t key_count = 0; // Keys in this shard chunk.
37 };
38
39 // A caller batch spread across one or more PendingShardRpc.
40 struct BatchRequest {
41 float* user_buffer = nullptr; // Final output buffer owned by caller.
42 bool assembled = false; // True once shard RPCs are merged.
43 std::size_t total_key_count = 0; // Total keys across all shards.
44 std::int32_t status_code =
45 static_cast<std::int32_t>(petps::RpcStatus::kPending);
46 std::vector<PendingShardRpc> shard_rpcs; // One pending RPC per shard chunk.
47 };
48
49 // A contiguous group of keys routed to a single shard/client.
50 struct ShardChunk {
51 int shard_id = 0; // Routed shard id.
52 int client_index = 0; // Client that serves this shard.
53 std::vector<uint64_t> keys; // Keys assigned to this shard chunk.
54 std::vector<std::size_t> positions; // Original positions in caller input.
55 };
56
57 // Maps a key to a logical shard using the configured hash method.
58 38 inline int PartitionKey(uint64_t key,
59 int num_shards,
60 const std::string& hash_method) {
61
1/6
✗ Branch 3 not taken.
✓ Branch 4 taken 38 times.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 38 CHECK_GT(num_shards, 0);
62
1/2
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
 38 if (hash_method == "city_hash") {
63 38 return static_cast<int>(GetHash(key) % static_cast<uint64_t>(num_shards));
64 }
65 ✗ if (hash_method == "simple_mod") {
66 ✗ return static_cast<int>(key % static_cast<uint64_t>(num_shards));
67 }
68 ✗ throw std::runtime_error("unsupported shard hash method: " + hash_method);
69 }
70
71 // Splits keys into per-shard chunks, each no larger than max_keys_per_rpc.
72 6 inline std::vector<ShardChunk> BuildChunks(
73 base::ConstArray<uint64_t> keys,
74 int num_shards,
75 const std::string& hash_method,
76 const std::unordered_map<int, int>& shard_to_client_index,
77 std::size_t max_keys_per_rpc) {
78
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 std::vector<std::vector<uint64_t>> shard_keys(num_shards);
79
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 std::vector<std::vector<std::size_t>> shard_positions(num_shards);
80
81
2/2
✓ Branch 1 taken 22 times.
✓ Branch 2 taken 6 times.
 28 for (std::size_t i = 0; i < keys.Size(); ++i) {
82
1/2
✓ Branch 2 taken 22 times.
✗ Branch 3 not taken.
 22 const int shard = PartitionKey(keys[i], num_shards, hash_method);
83
1/2
✓ Branch 3 taken 22 times.
✗ Branch 4 not taken.
 22 shard_keys[static_cast<std::size_t>(shard)].push_back(keys[i]);
84
1/2
✓ Branch 2 taken 22 times.
✗ Branch 3 not taken.
 22 shard_positions[static_cast<std::size_t>(shard)].push_back(i);
85 }
86
87 6 std::vector<ShardChunk> chunks;
88
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 for (int shard = 0; shard < num_shards; ++shard) {
89
1/2
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 12 const int client_index = shard_to_client_index.at(shard);
90 12 for (std::size_t offset = 0;
91
2/2
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 12 times.
 28 offset < shard_keys[static_cast<std::size_t>(shard)].size();
92 16 offset += max_keys_per_rpc) {
93 const std::size_t end =
94 32 std::min(offset + max_keys_per_rpc,
95 16 shard_keys[static_cast<std::size_t>(shard)].size());
96 16 ShardChunk chunk;
97 16 chunk.shard_id = shard;
98 16 chunk.client_index = client_index;
99
1/2
✓ Branch 3 taken 16 times.
✗ Branch 4 not taken.
 32 chunk.keys.assign(
100 16 shard_keys[static_cast<std::size_t>(shard)].begin() + offset,
101 16 shard_keys[static_cast<std::size_t>(shard)].begin() + end);
102
1/2
✓ Branch 3 taken 16 times.
✗ Branch 4 not taken.
 32 chunk.positions.assign(
103 16 shard_positions[static_cast<std::size_t>(shard)].begin() + offset,
104 16 shard_positions[static_cast<std::size_t>(shard)].begin() + end);
105
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 chunks.push_back(std::move(chunk));
106 16 }
107 }
108 12 return chunks;
109 6 }
110
111 // Merges completed shard responses into batch->user_buffer and writes the
112 // trailing batch status word. Returns true iff the whole batch succeeded.
113 4 inline bool FinalizeBatchIfNeeded(BatchRequest* batch, int value_size) {
114
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (batch == nullptr) {
115 ✗ return false;
116 }
117
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (batch->assembled) {
118 ✗ return batch->status_code ==
119 ✗ static_cast<std::int32_t>(petps::RpcStatus::kOk);
120 }
121
122 4 batch->status_code = static_cast<std::int32_t>(petps::RpcStatus::kOk);
123
2/2
✓ Branch 5 taken 12 times.
✓ Branch 6 taken 2 times.
 14 for (const auto& pending : batch->shard_rpcs) {
124 12 const auto* status_word = petps::FixedSlotStatusWord(
125 12 pending.recv_buffer, pending.key_count, value_size);
126
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 10 times.
 12 if (*status_word != static_cast<std::int32_t>(petps::RpcStatus::kOk)) {
127 2 batch->status_code = *status_word;
128 2 break;
129 }
130 }
131
132 4 const int embedding_dim = value_size / sizeof(float);
133
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
 4 if (batch->status_code == static_cast<std::int32_t>(petps::RpcStatus::kOk)) {
134
2/2
✓ Branch 5 taken 8 times.
✓ Branch 6 taken 2 times.
 10 for (const auto& pending : batch->shard_rpcs) {
135 8 const float* shard_values =
136 static_cast<const float*>(pending.recv_buffer);
137
2/2
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 8 times.
 22 for (std::size_t i = 0; i < pending.original_positions.size(); ++i) {
138 28 std::memcpy(
139 14 batch->user_buffer + pending.original_positions[i] * embedding_dim,
140 14 shard_values + i * embedding_dim,
141 value_size);
142 }
143 }
144 }
145
146 4 auto* batch_status_word = reinterpret_cast<std::int32_t*>(
147 4 reinterpret_cast<char*>(batch->user_buffer) +
148 4 batch->total_key_count * static_cast<std::size_t>(value_size));
149 4 *batch_status_word = batch->status_code;
150 4 batch->assembled = true;
151 4 return batch->status_code == static_cast<std::int32_t>(petps::RpcStatus::kOk);
152 }
153
154 } // namespace shard_routing
155 } // namespace recstore
156