GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
98 / 169
58.0%
Branches:
62 / 148
41.9%
Functions:
14 / 23
60.9%
ps/rdma/rdma_protocol.h
Line Branch Exec Source
1 #pragma once
2
3 #include <algorithm>
4 #include <array>
5 #include <atomic>
6 #include <chrono>
7 #include <cstddef>
8 #include <cstdint>
9 #include <cstring>
10 #include <limits>
11 #include <stdexcept>
12 #include <string>
13 #include <string_view>
14 #include <thread>
15
16 #include "base/array.h"
17 #include "base/flatc.h"
18 #include "base/log.h"
19 #include "base/tensor.h"
20 #include "ps/base/parameters.h"
21 #include "ps/rdma/rdma_status.h"
22
23 namespace petps {
24
25 inline constexpr std::uint32_t kRcProtocolMagic = 0x52435053;
26 inline constexpr std::uint16_t kRcProtocolVersion = 1;
27 inline constexpr std::size_t kTableNameBytes = 64;
28 inline constexpr std::uint32_t kRcSlotReady = 1;
29 inline constexpr std::uint32_t kRcSlotDone = 2;
30 inline constexpr std::uint32_t kRcFlagGetDirectSg = 1U << 0;
31 inline constexpr std::uint32_t kRcFlagGetAllowFallbackCopy = 1U << 1;
32
33 enum class RcOp : std::uint16_t {
34 kGet = 1,
35 kPut = 2,
36 kUpdate = 3,
37 kInitTable = 4,
38 kUpdateFlat = 5,
39 kSparseGradientCommit = 6,
40 };
41
42 enum class RcHashMethod : std::uint8_t {
43 kCityHash = 1,
44 kSimpleMod = 2,
45 };
46
47 struct alignas(64) RequestDescriptor {
48 std::uint32_t magic = kRcProtocolMagic;
49 std::uint16_t version = kRcProtocolVersion;
50 std::uint16_t op = static_cast<std::uint16_t>(RcOp::kGet);
51 std::uint64_t seq = 0; // Monotonic lane-local request sequence.
52 std::uint32_t shard_id = 0; // Logical shard targeted by this RPC.
53 std::uint32_t client_id = 0; // Logical client owner of this lane.
54 std::uint32_t qp_index = 0; // Lane index within the client.
55 std::uint32_t key_count = 0; // Number of keys in the payload.
56 std::uint32_t value_size = 0; // Row size in bytes for GET responses.
57 std::uint32_t embedding_dim = 0; // Row size expressed as float count.
58 std::uint32_t payload_offset = 0; // Offset from slot base to payload.
59 std::uint32_t payload_bytes = 0; // Bytes occupied by the payload.
60 std::uint32_t response_bytes = 0; // Bytes expected in the response payload.
61 std::uint32_t reserved0 = 0;
62 std::uint64_t client_response_addr =
63 0; // Optional client response address for verbs RC.
64 std::uint32_t client_response_rkey =
65 0; // Optional client response remote key.
66 std::uint32_t client_status_rkey = 0; // Optional client status remote key.
67 std::uint64_t client_status_addr =
68 0; // Optional client status address for verbs RC.
69 std::uint32_t flags = 0; // Op-specific protocol flags.
70 std::uint32_t reserved1 = 0;
71 std::array<char, kTableNameBytes>
72 table_name{}; // Optional logical table name.
73 };
74
75 struct alignas(64) CommitWord {
76 std::atomic<std::uint64_t> seq{0}; // Mirrors RequestDescriptor::seq.
77 std::atomic<std::uint32_t> state{
78 0}; // READY/DONE state published by client/server.
79 std::uint32_t checksum_or_reserved =
80 0; // Reserved for future integrity checks.
81 };
82
83 struct alignas(64) StatusWord {
84 std::atomic<std::uint64_t> seq{0}; // Mirrors the request seq completed here.
85 std::atomic<std::uint32_t> state{
86 0}; // DONE when the server has finished writing.
87 std::int32_t status = 0; // RpcStatus value returned by the server.
88 std::uint32_t response_bytes = 0; // Payload bytes valid in the response slot.
89 std::uint32_t reserved = 0;
90 };
91
92 static_assert(sizeof(RequestDescriptor) == 192, "RequestDescriptor size");
93 static_assert(alignof(RequestDescriptor) == 64, "RequestDescriptor align");
94 static_assert(alignof(CommitWord) == 64, "CommitWord align");
95 static_assert(alignof(StatusWord) == 64, "StatusWord align");
96
97 ✗ inline std::size_t Align64(std::size_t value) {
98 ✗ return (value + 63U) & ~std::size_t{63U};
99 }
100
101 2 inline std::size_t GetKeysPerRpcByResponseBudget(
102 std::size_t value_size, std::size_t mtu_bytes, std::size_t response_mtu) {
103
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 2 if (value_size == 0 || mtu_bytes == 0 || response_mtu == 0) {
104 ✗ return 0;
105 }
106 2 return (mtu_bytes * response_mtu) / value_size;
107 }
108
109 ✗ inline std::size_t GetRequestBytes(std::size_t key_count) {
110 ✗ return key_count * sizeof(std::uint64_t);
111 }
112
113 inline std::size_t
114 ✗ GetResponseBytes(std::size_t key_count, std::size_t value_size) {
115 ✗ return key_count * value_size;
116 }
117
118 inline std::size_t
119 ✗ FixedSlotResponseBytes(std::size_t key_count, std::size_t value_size) {
120 ✗ return GetResponseBytes(key_count, value_size) + sizeof(std::int32_t);
121 }
122
123 ✗ inline std::size_t InitTablePayloadBytes() { return sizeof(std::uint64_t) * 3; }
124
125 ✗ inline std::size_t PutPayloadBudget(std::size_t request_slot_bytes) {
126 ✗ if (request_slot_bytes <=
127
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 ✗ Align64(sizeof(RequestDescriptor)) + Align64(sizeof(CommitWord))) {
128 ✗ return 0;
129 }
130 ✗ return request_slot_bytes - Align64(sizeof(RequestDescriptor)) -
131 ✗ Align64(sizeof(CommitWord));
132 }
133
134 inline std::size_t ParameterReaderBytes(const ParameterCompressReader& reader) {
135 return static_cast<std::size_t>(reader.byte_size());
136 }
137
138 6 inline std::size_t PutPayloadBytes(
139 base::ConstArray<std::uint64_t> keys,
140 const float* values,
141 int64_t embedding_dim,
142 std::string* payload,
143 std::string* error = nullptr) {
144
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (payload == nullptr) {
145
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
146
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ *error = "payload buffer is null";
147 }
148 ✗ return 0;
149 }
150
4/8
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 6 times.
 6 if (keys.Size() != 0 && (values == nullptr || embedding_dim <= 0)) {
151
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
152
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ *error = "values pointer or embedding dim is invalid";
153 }
154 ✗ return 0;
155 }
156
157
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 ParameterCompressor compressor;
158
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 CompressEmbeddingRows(
159 6 &compressor, keys.Data(), values, keys.Size(), embedding_dim);
160
161 6 payload->clear();
162
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 compressor.ToBlock(payload);
163 6 return payload->size();
164 6 }
165
166 6 inline std::size_t PutPayloadBytes(
167 base::ConstArray<std::uint64_t> keys,
168 const base::RecTensor& values,
169 std::string* payload,
170 std::string* error = nullptr) {
171
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
 6 if (keys.Size() == 0) {
172 ✗ return PutPayloadBytes(keys, nullptr, 0, payload, error);
173 }
174
1/2
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 12 if (values.dtype() != base::DataType::FLOAT32 || values.dim() != 2 ||
175
4/8
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 6 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 6 times.
 18 values.shape(0) != static_cast<int64_t>(keys.Size()) ||
176 6 values.data() == nullptr) {
177
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
178 ✗ *error = "keys and values size mismatch";
179 }
180 ✗ return 0;
181 }
182 6 return PutPayloadBytes(
183 6 keys, values.data_as<float>(), values.shape(1), payload, error);
184 }
185
186 inline std::size_t UpdatePayloadBytes(
187 base::ConstArray<std::uint64_t> keys,
188 const float* values,
189 int64_t embedding_dim,
190 std::string* payload,
191 std::string* error = nullptr) {
192 return PutPayloadBytes(keys, values, embedding_dim, payload, error);
193 }
194
195 2 inline std::size_t UpdatePayloadBytes(
196 base::ConstArray<std::uint64_t> keys,
197 const base::RecTensor& values,
198 std::string* payload,
199 std::string* error = nullptr) {
200 2 return PutPayloadBytes(keys, values, payload, error);
201 }
202
203 inline std::size_t
204 18 FlatUpdatePayloadBytes(std::size_t key_count, std::size_t embedding_dim) {
205 18 constexpr std::size_t kKeyBytes = sizeof(std::uint64_t);
206 18 constexpr std::size_t kFloatBytes = sizeof(float);
207 18 const std::size_t max_size = std::numeric_limits<std::size_t>::max();
208
3/4
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 16 times.
 18 if (embedding_dim == 0 ||
209 embedding_dim > (max_size - kKeyBytes) / kFloatBytes) {
210 2 return 0;
211 }
212 16 const std::size_t row_bytes = kKeyBytes + embedding_dim * kFloatBytes;
213
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 2 times.
 16 return key_count > max_size / row_bytes ? 0 : key_count * row_bytes;
214 }
215
216 4 inline std::size_t UpdatePayloadBytesFlat(
217 base::ConstArray<std::uint64_t> keys,
218 const float* values,
219 std::size_t embedding_dim,
220 std::string* payload,
221 std::string* error = nullptr) {
222
4/8
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
 4 if (payload == nullptr || (keys.Size() > 0 && values == nullptr)) {
223
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
224 ✗ *error = "payload buffer or values is null";
225 }
226 ✗ return 0;
227 }
228
229 const std::size_t payload_bytes =
230 4 FlatUpdatePayloadBytes(keys.Size(), embedding_dim);
231
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (payload_bytes == 0) {
232
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
233 ✗ *error = "invalid flat update shape";
234 }
235 ✗ return 0;
236 }
237 4 const std::size_t key_bytes = keys.Size() * sizeof(std::uint64_t);
238 4 const std::size_t value_bytes = keys.Size() * embedding_dim * sizeof(float);
239 4 payload->resize(payload_bytes);
240
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (key_bytes > 0) {
241 4 std::memcpy(payload->data(), keys.Data(), key_bytes);
242 4 std::memcpy(payload->data() + key_bytes, values, value_bytes);
243 }
244 4 return payload->size();
245 }
246
247 4 inline std::size_t PackFlatUpdatePayloadGather(
248 const std::uint64_t* keys,
249 const float* values,
250 std::size_t num_rows,
251 std::size_t embedding_dim,
252 const std::size_t* row_indices,
253 std::size_t row_count,
254 void* payload,
255 std::size_t payload_capacity,
256 std::string* error = nullptr) {
257
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 if ((row_count > 0 &&
258
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
 4 (keys == nullptr || values == nullptr || row_indices == nullptr)) ||
259 payload == nullptr) {
260
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
261 ✗ *error = "flat update gather input is null";
262 }
263 ✗ return 0;
264 }
265 const std::size_t payload_bytes =
266 4 FlatUpdatePayloadBytes(row_count, embedding_dim);
267
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (payload_bytes == 0 || payload_bytes > payload_capacity) {
268
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
269 ✗ *error = "flat update gather payload does not fit";
270 }
271 ✗ return 0;
272 }
273
274 4 auto* payload_keys = static_cast<std::uint64_t*>(payload);
275 4 auto* payload_values = reinterpret_cast<float*>(
276 4 static_cast<char*>(payload) + row_count * sizeof(std::uint64_t));
277 4 const std::size_t row_bytes = embedding_dim * sizeof(float);
278
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 for (std::size_t row = 0; row < row_count; ++row) {
279 6 const std::size_t source_row = row_indices[row];
280
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 if (source_row >= num_rows) {
281
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (error != nullptr) {
282 2 *error = "flat update gather row index is out of range";
283 }
284 2 return 0;
285 }
286 4 payload_keys[row] = keys[source_row];
287 4 std::memcpy(payload_values + row * embedding_dim,
288 4 values + source_row * embedding_dim,
289 row_bytes);
290 }
291 2 return payload_bytes;
292 }
293
294 ✗ inline bool CopyTableName(std::string_view table_name,
295 std::array<char, kTableNameBytes>* storage) {
296
0/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
 ✗ if (storage == nullptr || table_name.size() >= kTableNameBytes) {
297 ✗ return false;
298 }
299
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ storage->fill('\0');
300 ✗ std::memcpy(storage->data(), table_name.data(), table_name.size());
301 ✗ return true;
302 }
303
304 inline std::string_view
305 ✗ DescriptorTableName(const RequestDescriptor& descriptor) {
306 ✗ return std::string_view(
307 descriptor.table_name.data(),
308
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 ✗ std::find(
309 ✗ descriptor.table_name.begin(), descriptor.table_name.end(), '\0') -
310 ✗ descriptor.table_name.begin());
311 }
312
313 ✗ inline bool ValidateRequestDescriptor(
314 const RequestDescriptor& descriptor,
315 std::size_t request_slot_bytes,
316 std::size_t response_slot_bytes,
317 std::string* error = nullptr) {
318
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (descriptor.magic != kRcProtocolMagic) {
319
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
320 ✗ *error = "bad request magic";
321 }
322 ✗ return false;
323 }
324
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (descriptor.version != kRcProtocolVersion) {
325
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
326 ✗ *error = "bad request version";
327 }
328 ✗ return false;
329 }
330
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (descriptor.payload_offset < sizeof(RequestDescriptor) ||
331 ✗ static_cast<std::size_t>(descriptor.payload_offset) +
332
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ descriptor.payload_bytes >
333 request_slot_bytes) {
334
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
335 ✗ *error = "request payload exceeds slot capacity";
336 }
337 ✗ return false;
338 }
339
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (descriptor.response_bytes > response_slot_bytes) {
340
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ if (error != nullptr) {
341 ✗ *error = "response exceeds slot capacity";
342 }
343 ✗ return false;
344 }
345 ✗ return true;
346 }
347
348 8 inline void ResetStatusWord(StatusWord* status, std::uint64_t seq) {
349 8 status->status = static_cast<std::int32_t>(RpcStatus::kPending);
350 8 status->response_bytes = 0;
351 8 status->seq.store(seq, std::memory_order_release);
352 8 status->state.store(0, std::memory_order_release);
353 8 }
354
355 4812 inline bool StatusWordDone(const StatusWord& status, std::uint64_t seq) {
356
2/2
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 4804 times.
 9632 return status.state.load(std::memory_order_acquire) == kRcSlotDone &&
357
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
 4828 status.seq.load(std::memory_order_acquire) == seq;
358 }
359
360 class RdmaRequestTimeout : public std::runtime_error {
361 public:
362 using std::runtime_error::runtime_error;
363 };
364
365 inline std::int32_t
366 2 WaitStatusWord(const StatusWord& status,
367 std::uint64_t seq,
368 int timeout_ms,
369 int spin_iterations) {
370 2 const auto start = std::chrono::steady_clock::now();
371 2 int spins = 0;
372
1/2
✓ Branch 1 taken 4800 times.
✗ Branch 2 not taken.
 4800 while (!StatusWordDone(status, seq)) {
373
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4800 times.
 4800 if (spins < spin_iterations) {
374 ✗ ++spins;
375 } else {
376 4800 spins = 0;
377 4800 std::this_thread::yield();
378 }
379
1/2
✓ Branch 0 taken 4800 times.
✗ Branch 1 not taken.
 9600 if (timeout_ms > 0 &&
380
1/2
✓ Branch 1 taken 4800 times.
✗ Branch 2 not taken.
 4800 std::chrono::duration_cast<std::chrono::milliseconds>(
381
3/4
✓ Branch 2 taken 4800 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 4798 times.
 9600 std::chrono::steady_clock::now() - start)
382
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4798 times.
 4800 .count() > timeout_ms) {
383
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 throw RdmaRequestTimeout("RDMA request wait timeout");
384 }
385 }
386 ✗ return status.status;
387 }
388
389 class RequestSlotLifecycle {
390 public:
391 14 bool TryAcquire(const StatusWord& status) {
392
6/6
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 8 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 4 times.
✓ Branch 5 taken 2 times.
✓ Branch 6 taken 12 times.
 14 if (state_ == State::kQuarantined && StatusWordDone(status, seq_)) {
393 2 state_ = State::kIdle;
394 2 seq_ = 0;
395 }
396
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 8 times.
 14 if (state_ != State::kIdle) {
397 6 return false;
398 }
399 8 state_ = State::kInFlight;
400 8 return true;
401 }
402
403 2 void Release() {
404 2 state_ = State::kIdle;
405 2 seq_ = 0;
406 2 }
407
408 2 void Quarantine(std::uint64_t seq) {
409 // ponytail: transport teardown recovers slots whose completion never
410 // arrives.
411 2 state_ = State::kQuarantined;
412 2 seq_ = seq;
413 2 }
414
415 4 bool IsQuarantined() const { return state_ == State::kQuarantined; }
416
417 private:
418 enum class State : std::uint8_t { kIdle, kInFlight, kQuarantined };
419
420 State state_ = State::kIdle;
421 std::uint64_t seq_ = 0;
422 };
423
424 } // namespace petps
425