ps/rdma/rdma_distributed_executor.cc
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "ps/rdma/rdma_distributed_executor.h" | ||
| 2 | |||
| 3 | #include <exception> | ||
| 4 | #include <cstring> | ||
| 5 | #include <limits> | ||
| 6 | #include <stdexcept> | ||
| 7 | #include <thread> | ||
| 8 | #include <utility> | ||
| 9 | |||
| 10 | namespace recstore { | ||
| 11 | |||
| 12 | ✗ | RdmaDistributedExecutor::RdmaDistributedExecutor( | |
| 13 | std::vector<petps::PetPSClient*> clients, | ||
| 14 | int value_size, | ||
| 15 | int first_batch_id, | ||
| 16 | ✗ | int batch_id_step) | |
| 17 | ✗ | : clients_(std::move(clients)), | |
| 18 | ✗ | value_size_(value_size), | |
| 19 | ✗ | next_id_(first_batch_id), | |
| 20 | ✗ | id_step_(batch_id_step) { | |
| 21 | ✗ | if (clients_.empty() || value_size_ <= 0 || | |
| 22 | ✗ | (id_step_ != 1 && id_step_ != -1)) { | |
| 23 | ✗ | throw std::invalid_argument("invalid RDMA distributed executor config"); | |
| 24 | } | ||
| 25 | ✗ | } | |
| 26 | |||
| 27 | ✗ | RdmaDistributedExecutor::~RdmaDistributedExecutor() { | |
| 28 | ✗ | for (auto& [batch_id, batch] : batches_) { | |
| 29 | (void)batch_id; | ||
| 30 | try { | ||
| 31 | ✗ | ReleaseShardRpcs(&batch.shard_rpcs); | |
| 32 | ✗ | } catch (const std::exception& error) { | |
| 33 | ✗ | LOG(ERROR) << "failed to release RDMA batch during shutdown: " | |
| 34 | ✗ | << error.what(); | |
| 35 | ✗ | } | |
| 36 | ✗ | RecycleReceiveBuffers(&batch.receive_buffers); | |
| 37 | } | ||
| 38 | ✗ | for (auto& [update_id, update] : pending_updates_) { | |
| 39 | (void)update_id; | ||
| 40 | ✗ | for (const auto& [client_index, rpc_id] : update.shard_rpcs) { | |
| 41 | try { | ||
| 42 | ✗ | ClientAt(client_index)->RevokeRPCResource(rpc_id); | |
| 43 | ✗ | } catch (const std::exception& error) { | |
| 44 | ✗ | LOG(ERROR) << "failed to release RDMA update during shutdown: " | |
| 45 | ✗ | << error.what(); | |
| 46 | ✗ | } | |
| 47 | } | ||
| 48 | } | ||
| 49 | ✗ | } | |
| 50 | |||
| 51 | ✗ | petps::PetPSClient* RdmaDistributedExecutor::ClientAt(int client_index) const { | |
| 52 | ✗ | if (client_index < 0 || client_index >= static_cast<int>(clients_.size()) || | |
| 53 | ✗ | clients_[static_cast<std::size_t>(client_index)] == nullptr) { | |
| 54 | ✗ | throw std::runtime_error("RDMA batch references an invalid shard client"); | |
| 55 | } | ||
| 56 | ✗ | return clients_[static_cast<std::size_t>(client_index)]; | |
| 57 | } | ||
| 58 | |||
| 59 | std::vector<char> | ||
| 60 | ✗ | RdmaDistributedExecutor::AcquireReceiveBuffer(std::size_t bytes) const { | |
| 61 | { | ||
| 62 | ✗ | std::lock_guard<std::mutex> guard(recv_pool_mu_); | |
| 63 | ✗ | for (auto it = recv_pool_.begin(); it != recv_pool_.end(); ++it) { | |
| 64 | ✗ | if (it->size() < bytes) { | |
| 65 | ✗ | continue; | |
| 66 | } | ||
| 67 | ✗ | std::vector<char> buffer = std::move(*it); | |
| 68 | ✗ | recv_pool_bytes_ -= buffer.size(); | |
| 69 | ✗ | recv_pool_.erase(it); | |
| 70 | ✗ | return buffer; | |
| 71 | ✗ | } | |
| 72 | ✗ | } | |
| 73 | // Cold path: first call for this size. The server overwrites the whole | ||
| 74 | // response before reporting kOk, so this zero-fill only exists because | ||
| 75 | // std::vector cannot grow without initializing. | ||
| 76 | ✗ | return std::vector<char>(bytes, 0); | |
| 77 | } | ||
| 78 | |||
| 79 | ✗ | void RdmaDistributedExecutor::RecycleReceiveBuffers( | |
| 80 | std::vector<std::vector<char>>* buffers) const { | ||
| 81 | ✗ | if (buffers == nullptr || buffers->empty()) { | |
| 82 | ✗ | return; | |
| 83 | } | ||
| 84 | // Idle buffers are capped by count and by bytes: peak count tracks in-flight | ||
| 85 | // chunks (bounded by max_keys_per_request), while the byte budget keeps large | ||
| 86 | // value sizes from turning the pool into a second copy of the table. | ||
| 87 | ✗ | constexpr std::size_t kMaxPooledReceiveBuffers = 64; | |
| 88 | ✗ | constexpr std::size_t kMaxPooledReceiveBytes = 64ull * 1024 * 1024; | |
| 89 | ✗ | std::lock_guard<std::mutex> guard(recv_pool_mu_); | |
| 90 | ✗ | for (auto& buffer : *buffers) { | |
| 91 | ✗ | if (recv_pool_.size() >= kMaxPooledReceiveBuffers || | |
| 92 | ✗ | recv_pool_bytes_ + buffer.size() > kMaxPooledReceiveBytes) { | |
| 93 | ✗ | break; | |
| 94 | } | ||
| 95 | ✗ | recv_pool_bytes_ += buffer.size(); | |
| 96 | ✗ | recv_pool_.push_back(std::move(buffer)); | |
| 97 | } | ||
| 98 | ✗ | buffers->clear(); | |
| 99 | ✗ | } | |
| 100 | |||
| 101 | ✗ | int RdmaDistributedExecutor::RegisterBatch(shard_routing::BatchRequest batch) { | |
| 102 | ✗ | std::unique_lock<std::mutex> guard(batches_mu_); | |
| 103 | ✗ | if (next_id_ < std::numeric_limits<int>::min() || | |
| 104 | ✗ | next_id_ > std::numeric_limits<int>::max()) { | |
| 105 | ✗ | guard.unlock(); | |
| 106 | ✗ | ReleaseShardRpcs(&batch.shard_rpcs); | |
| 107 | ✗ | RecycleReceiveBuffers(&batch.receive_buffers); | |
| 108 | ✗ | throw std::runtime_error("RDMA distributed batch handle space exhausted"); | |
| 109 | } | ||
| 110 | ✗ | const int batch_id = static_cast<int>(next_id_); | |
| 111 | ✗ | next_id_ += id_step_; | |
| 112 | ✗ | if (batches_.find(batch_id) != batches_.end()) { | |
| 113 | ✗ | guard.unlock(); | |
| 114 | ✗ | ReleaseShardRpcs(&batch.shard_rpcs); | |
| 115 | ✗ | RecycleReceiveBuffers(&batch.receive_buffers); | |
| 116 | ✗ | throw std::runtime_error("duplicate RDMA distributed batch handle"); | |
| 117 | } | ||
| 118 | ✗ | if (!batches_.emplace(batch_id, std::move(batch)).second) { | |
| 119 | ✗ | throw std::runtime_error("duplicate RDMA distributed batch handle"); | |
| 120 | } | ||
| 121 | ✗ | return batch_id; | |
| 122 | ✗ | } | |
| 123 | |||
| 124 | ✗ | int RdmaDistributedExecutor::SubmitGetParameter( | |
| 125 | const std::vector<shard_routing::ShardChunk>& chunks, | ||
| 126 | std::size_t total_key_count, | ||
| 127 | int value_size, | ||
| 128 | float* values, | ||
| 129 | bool is_async, | ||
| 130 | int async_req_id, | ||
| 131 | std::size_t max_in_flight, | ||
| 132 | bool write_status_word) { | ||
| 133 | ✗ | if (values == nullptr || value_size <= 0 || max_in_flight == 0) { | |
| 134 | ✗ | throw std::invalid_argument("invalid RDMA GET request"); | |
| 135 | } | ||
| 136 | |||
| 137 | ✗ | shard_routing::BatchRequest batch; | |
| 138 | ✗ | batch.user_buffer = values; | |
| 139 | ✗ | batch.total_key_count = total_key_count; | |
| 140 | ✗ | batch.value_size = value_size; | |
| 141 | ✗ | batch.write_status_word = write_status_word; | |
| 142 | ✗ | if (write_status_word) { | |
| 143 | ✗ | shard_routing::WriteBatchStatus(&batch, value_size); | |
| 144 | } | ||
| 145 | ✗ | batch.receive_buffers.reserve(chunks.size()); | |
| 146 | ✗ | batch.shard_rpcs.reserve(chunks.size()); | |
| 147 | ✗ | std::vector<shard_routing::PendingShardRpc> window; | |
| 148 | ✗ | window.reserve(std::min(max_in_flight, chunks.size())); | |
| 149 | |||
| 150 | ✗ | auto drain_window = [this, &batch, &window]() { | |
| 151 | ✗ | WaitShardRpcs(window); | |
| 152 | ✗ | ReleaseShardRpcs(&window); | |
| 153 | ✗ | for (auto& pending : window) { | |
| 154 | ✗ | batch.shard_rpcs.push_back(std::move(pending)); | |
| 155 | } | ||
| 156 | ✗ | window.clear(); | |
| 157 | ✗ | }; | |
| 158 | |||
| 159 | try { | ||
| 160 | ✗ | for (std::size_t chunk_index = 0; chunk_index < chunks.size(); | |
| 161 | ++chunk_index) { | ||
| 162 | ✗ | const auto& chunk = chunks[chunk_index]; | |
| 163 | ✗ | if (chunk.keys.empty() || chunk.keys.size() != chunk.positions.size()) { | |
| 164 | ✗ | throw std::runtime_error("RDMA GET chunk shape mismatch"); | |
| 165 | } | ||
| 166 | ✗ | for (const std::size_t position : chunk.positions) { | |
| 167 | ✗ | if (position >= total_key_count) { | |
| 168 | ✗ | throw std::invalid_argument( | |
| 169 | ✗ | "RDMA GET chunk position is out of range"); | |
| 170 | } | ||
| 171 | } | ||
| 172 | ✗ | auto* client = ClientAt(chunk.client_index); | |
| 173 | ✗ | batch.receive_buffers.push_back(AcquireReceiveBuffer( | |
| 174 | ✗ | chunk.keys.size() * static_cast<std::size_t>(value_size) + | |
| 175 | sizeof(std::int32_t))); | ||
| 176 | ✗ | void* recv = batch.receive_buffers.back().data(); | |
| 177 | ✗ | shard_routing::PendingShardRpc pending{ | |
| 178 | ✗ | chunk.shard_id, | |
| 179 | ✗ | chunk.client_index, | |
| 180 | -1, | ||
| 181 | ✗ | chunk.positions, | |
| 182 | recv, | ||
| 183 | ✗ | chunk.keys.size(), | |
| 184 | ✗ | }; | |
| 185 | ✗ | pending.rpc_id = client->GetParameter( | |
| 186 | ✗ | base::ConstArray<std::uint64_t>(chunk.keys), | |
| 187 | static_cast<float*>(recv), | ||
| 188 | /*isAsync=*/true, | ||
| 189 | async_req_id, | ||
| 190 | value_size / static_cast<int>(sizeof(float))); | ||
| 191 | ✗ | if (pending.rpc_id < 0) { | |
| 192 | ✗ | throw std::runtime_error("failed to submit RDMA GET"); | |
| 193 | } | ||
| 194 | ✗ | window.push_back(std::move(pending)); | |
| 195 | ✗ | if (window.size() >= max_in_flight && chunk_index + 1 < chunks.size()) { | |
| 196 | ✗ | drain_window(); | |
| 197 | } | ||
| 198 | ✗ | } | |
| 199 | ✗ | } catch (...) { | |
| 200 | ✗ | const auto error = std::current_exception(); | |
| 201 | // Accepted RPCs must finish or retire their lanes before slots are reused. | ||
| 202 | try { | ||
| 203 | ✗ | WaitShardRpcs(window); | |
| 204 | ✗ | } catch (...) { | |
| 205 | ✗ | } | |
| 206 | try { | ||
| 207 | ✗ | ReleaseShardRpcs(&window); | |
| 208 | ✗ | } catch (...) { | |
| 209 | ✗ | } | |
| 210 | try { | ||
| 211 | ✗ | ReleaseShardRpcs(&batch.shard_rpcs); | |
| 212 | ✗ | } catch (...) { | |
| 213 | ✗ | } | |
| 214 | try { | ||
| 215 | ✗ | RecycleReceiveBuffers(&batch.receive_buffers); | |
| 216 | ✗ | } catch (...) { | |
| 217 | ✗ | } | |
| 218 | ✗ | std::rethrow_exception(error); | |
| 219 | ✗ | } | |
| 220 | ✗ | for (auto& pending : window) { | |
| 221 | ✗ | batch.shard_rpcs.push_back(std::move(pending)); | |
| 222 | } | ||
| 223 | |||
| 224 | ✗ | const int batch_id = RegisterBatch(std::move(batch)); | |
| 225 | ✗ | if (!is_async) { | |
| 226 | try { | ||
| 227 | ✗ | WaitBatch(batch_id); | |
| 228 | ✗ | } catch (...) { | |
| 229 | ✗ | const auto error = std::current_exception(); | |
| 230 | try { | ||
| 231 | ✗ | ReleaseBatch(batch_id); | |
| 232 | ✗ | } catch (...) { | |
| 233 | // Preserve the wait failure; the caller never received this handle. | ||
| 234 | ✗ | } | |
| 235 | ✗ | std::rethrow_exception(error); | |
| 236 | ✗ | } | |
| 237 | } | ||
| 238 | ✗ | return batch_id; | |
| 239 | ✗ | } | |
| 240 | |||
| 241 | ✗ | bool RdmaDistributedExecutor::FinalizeBatch( | |
| 242 | shard_routing::BatchRequest* batch) const { | ||
| 243 | ✗ | if (batch->assembled) { | |
| 244 | ✗ | return batch->status_code == | |
| 245 | ✗ | static_cast<std::int32_t>(petps::RpcStatus::kOk); | |
| 246 | } | ||
| 247 | ✗ | batch->status_code = | |
| 248 | ✗ | shard_routing::DecodeBatchStatus(*batch, batch->value_size); | |
| 249 | ✗ | if (batch->status_code == static_cast<std::int32_t>(petps::RpcStatus::kOk)) { | |
| 250 | ✗ | shard_routing::MergeBatchRows(batch, batch->value_size); | |
| 251 | } | ||
| 252 | ✗ | if (batch->write_status_word) { | |
| 253 | ✗ | shard_routing::WriteBatchStatus(batch, batch->value_size); | |
| 254 | } | ||
| 255 | ✗ | batch->assembled = true; | |
| 256 | ✗ | return batch->status_code == static_cast<std::int32_t>(petps::RpcStatus::kOk); | |
| 257 | } | ||
| 258 | |||
| 259 | ✗ | bool RdmaDistributedExecutor::QueryBatchFinished(int batch_id) { | |
| 260 | ✗ | std::lock_guard<std::mutex> guard(batches_mu_); | |
| 261 | ✗ | auto it = batches_.find(batch_id); | |
| 262 | ✗ | if (it == batches_.end()) { | |
| 263 | ✗ | throw std::runtime_error("unknown RDMA distributed batch handle"); | |
| 264 | } | ||
| 265 | ✗ | for (const auto& pending : it->second.shard_rpcs) { | |
| 266 | ✗ | if (pending.released) { | |
| 267 | ✗ | continue; | |
| 268 | } | ||
| 269 | ✗ | if (!ClientAt(pending.client_index)->QueryRPCFinished(pending.rpc_id)) { | |
| 270 | ✗ | return false; | |
| 271 | } | ||
| 272 | } | ||
| 273 | ✗ | for (const auto& pending : it->second.shard_rpcs) { | |
| 274 | ✗ | if (!pending.released) { | |
| 275 | ✗ | ClientAt(pending.client_index)->WaitRPCFinish(pending.rpc_id); | |
| 276 | } | ||
| 277 | } | ||
| 278 | ✗ | return FinalizeBatch(&it->second); | |
| 279 | ✗ | } | |
| 280 | |||
| 281 | ✗ | const float* RdmaDistributedExecutor::BorrowBatchResult( | |
| 282 | int batch_id, std::int32_t* status_code, std::size_t* response_bytes) { | ||
| 283 | ✗ | std::lock_guard<std::mutex> guard(batches_mu_); | |
| 284 | ✗ | const auto it = batches_.find(batch_id); | |
| 285 | ✗ | if (it == batches_.end()) { | |
| 286 | ✗ | throw std::runtime_error("unknown RDMA distributed batch handle"); | |
| 287 | } | ||
| 288 | ✗ | const auto& batch = it->second; | |
| 289 | ✗ | if (batch.shard_rpcs.size() != 1 || batch.shard_rpcs.front().released || | |
| 290 | ✗ | batch.shard_rpcs.front().key_count != batch.total_key_count) { | |
| 291 | ✗ | return nullptr; | |
| 292 | } | ||
| 293 | ✗ | const auto& pending = batch.shard_rpcs.front(); | |
| 294 | ✗ | for (std::size_t i = 0; i < pending.original_positions.size(); ++i) { | |
| 295 | ✗ | if (pending.original_positions[i] != i) { | |
| 296 | ✗ | return nullptr; | |
| 297 | } | ||
| 298 | } | ||
| 299 | ✗ | auto* client = ClientAt(pending.client_index); | |
| 300 | ✗ | std::size_t key_count = 0; | |
| 301 | ✗ | const float* payload = client->BorrowGetResultPayload( | |
| 302 | ✗ | pending.rpc_id, &key_count, response_bytes, status_code); | |
| 303 | ✗ | return key_count == batch.total_key_count ? payload : nullptr; | |
| 304 | ✗ | } | |
| 305 | |||
| 306 | ✗ | std::uint64_t RdmaDistributedExecutor::SubmitPrefetch( | |
| 307 | const std::vector<shard_routing::ShardChunk>& chunks, | ||
| 308 | std::size_t key_count, | ||
| 309 | std::int64_t embedding_dim, | ||
| 310 | int value_size, | ||
| 311 | std::size_t max_in_flight) { | ||
| 312 | ✗ | if (key_count == 0 || embedding_dim <= 0) { | |
| 313 | ✗ | throw std::invalid_argument("invalid RDMA prefetch request"); | |
| 314 | } | ||
| 315 | ✗ | const std::size_t value_count = | |
| 316 | ✗ | key_count * static_cast<std::size_t>(embedding_dim); | |
| 317 | ✗ | auto buffer = std::make_shared<std::vector<float>>(value_count + 1, 0.0f); | |
| 318 | ✗ | const int batch_id = SubmitGetParameter( | |
| 319 | chunks, key_count, value_size, buffer->data(), true, 0, max_in_flight); | ||
| 320 | ✗ | std::lock_guard<std::mutex> guard(prefetches_mu_); | |
| 321 | ✗ | const std::uint64_t prefetch_id = next_prefetch_id_++; | |
| 322 | ✗ | prefetches_.emplace( | |
| 323 | prefetch_id, | ||
| 324 | ✗ | PendingPrefetch{buffer, | |
| 325 | batch_id, | ||
| 326 | ✗ | static_cast<std::int64_t>(key_count), | |
| 327 | embedding_dim, | ||
| 328 | value_size}); | ||
| 329 | ✗ | return prefetch_id; | |
| 330 | ✗ | } | |
| 331 | |||
| 332 | RdmaDistributedExecutor::PendingPrefetch | ||
| 333 | ✗ | RdmaDistributedExecutor::GetPrefetch(std::uint64_t prefetch_id) const { | |
| 334 | ✗ | std::lock_guard<std::mutex> guard(prefetches_mu_); | |
| 335 | ✗ | const auto it = prefetches_.find(prefetch_id); | |
| 336 | ✗ | if (it == prefetches_.end()) { | |
| 337 | ✗ | throw std::runtime_error("unknown or consumed RDMA prefetch handle"); | |
| 338 | } | ||
| 339 | ✗ | return it->second; | |
| 340 | ✗ | } | |
| 341 | |||
| 342 | ✗ | void RdmaDistributedExecutor::ErasePrefetch(std::uint64_t prefetch_id) { | |
| 343 | ✗ | std::lock_guard<std::mutex> guard(prefetches_mu_); | |
| 344 | ✗ | if (prefetches_.erase(prefetch_id) == 0) { | |
| 345 | ✗ | throw std::runtime_error("unknown or consumed RDMA prefetch handle"); | |
| 346 | } | ||
| 347 | ✗ | } | |
| 348 | |||
| 349 | ✗ | bool RdmaDistributedExecutor::IsPrefetchDone(std::uint64_t prefetch_id) { | |
| 350 | ✗ | return QueryBatchFinished(GetPrefetch(prefetch_id).batch_id); | |
| 351 | } | ||
| 352 | |||
| 353 | ✗ | void RdmaDistributedExecutor::WaitForPrefetch(std::uint64_t prefetch_id) { | |
| 354 | ✗ | WaitBatch(GetPrefetch(prefetch_id).batch_id); | |
| 355 | ✗ | } | |
| 356 | |||
| 357 | std::int64_t | ||
| 358 | ✗ | RdmaDistributedExecutor::PrefetchEmbeddingDim(std::uint64_t prefetch_id) const { | |
| 359 | ✗ | return GetPrefetch(prefetch_id).embedding_dim; | |
| 360 | } | ||
| 361 | |||
| 362 | RdmaDistributedExecutor::PrefetchResult | ||
| 363 | ✗ | RdmaDistributedExecutor::ReadPrefetch(std::uint64_t prefetch_id) { | |
| 364 | ✗ | const PendingPrefetch state = GetPrefetch(prefetch_id); | |
| 365 | ✗ | PrefetchResult result; | |
| 366 | ✗ | result.buffer = state.buffer; | |
| 367 | ✗ | result.key_count = state.key_count; | |
| 368 | ✗ | result.embedding_dim = state.embedding_dim; | |
| 369 | ✗ | result.response_bytes = 0; | |
| 370 | ✗ | result.status_code = static_cast<std::int32_t>(petps::RpcStatus::kOk); | |
| 371 | ✗ | result.payload = BorrowBatchResult( | |
| 372 | ✗ | state.batch_id, &result.status_code, &result.response_bytes); | |
| 373 | ✗ | if (result.payload == nullptr) { | |
| 374 | ✗ | WaitBatch(state.batch_id); | |
| 375 | ✗ | const auto* status_word = petps::FixedSlotStatusWord( | |
| 376 | ✗ | state.buffer->data(), | |
| 377 | ✗ | static_cast<std::size_t>(state.key_count), | |
| 378 | ✗ | state.value_size); | |
| 379 | ✗ | result.status_code = *status_word; | |
| 380 | ✗ | result.response_bytes = static_cast<std::size_t>(state.key_count) * | |
| 381 | ✗ | static_cast<std::size_t>(state.value_size); | |
| 382 | ✗ | result.payload = state.buffer->data(); | |
| 383 | } | ||
| 384 | ✗ | return result; | |
| 385 | ✗ | } | |
| 386 | |||
| 387 | ✗ | void RdmaDistributedExecutor::ReleasePrefetch(std::uint64_t prefetch_id) { | |
| 388 | ✗ | const PendingPrefetch state = GetPrefetch(prefetch_id); | |
| 389 | ✗ | ErasePrefetch(prefetch_id); | |
| 390 | ✗ | ReleaseBatch(state.batch_id); | |
| 391 | ✗ | } | |
| 392 | |||
| 393 | ✗ | void RdmaDistributedExecutor::WaitShardRpcs( | |
| 394 | const std::vector<shard_routing::PendingShardRpc>& shard_rpcs) const { | ||
| 395 | // Every RPC in this window is already submitted; waiting preserves overlap. | ||
| 396 | ✗ | std::exception_ptr first_error; | |
| 397 | ✗ | for (const auto& pending : shard_rpcs) { | |
| 398 | ✗ | if (pending.released) { | |
| 399 | ✗ | continue; | |
| 400 | } | ||
| 401 | try { | ||
| 402 | ✗ | ClientAt(pending.client_index)->WaitRPCFinish(pending.rpc_id); | |
| 403 | ✗ | } catch (...) { | |
| 404 | ✗ | if (first_error == nullptr) { | |
| 405 | ✗ | first_error = std::current_exception(); | |
| 406 | } | ||
| 407 | ✗ | } | |
| 408 | } | ||
| 409 | // Later RPCs must also finish or retire before the caller releases the | ||
| 410 | // window. | ||
| 411 | ✗ | if (first_error != nullptr) { | |
| 412 | ✗ | std::rethrow_exception(first_error); | |
| 413 | } | ||
| 414 | ✗ | } | |
| 415 | |||
| 416 | ✗ | void RdmaDistributedExecutor::ReleaseShardRpcs( | |
| 417 | std::vector<shard_routing::PendingShardRpc>* shard_rpcs) const { | ||
| 418 | ✗ | if (shard_rpcs == nullptr) { | |
| 419 | ✗ | throw std::invalid_argument("RDMA shard RPC list must not be null"); | |
| 420 | } | ||
| 421 | ✗ | std::exception_ptr first_error; | |
| 422 | ✗ | for (auto& pending : *shard_rpcs) { | |
| 423 | ✗ | if (pending.released) { | |
| 424 | ✗ | continue; | |
| 425 | } | ||
| 426 | ✗ | pending.released = true; | |
| 427 | try { | ||
| 428 | ✗ | ClientAt(pending.client_index)->RevokeRPCResource(pending.rpc_id); | |
| 429 | ✗ | } catch (...) { | |
| 430 | ✗ | if (first_error == nullptr) { | |
| 431 | ✗ | first_error = std::current_exception(); | |
| 432 | } | ||
| 433 | ✗ | } | |
| 434 | } | ||
| 435 | ✗ | if (first_error != nullptr) { | |
| 436 | ✗ | std::rethrow_exception(first_error); | |
| 437 | } | ||
| 438 | ✗ | } | |
| 439 | |||
| 440 | ✗ | int RdmaDistributedExecutor::PutParameter( | |
| 441 | const std::vector<shard_routing::ShardChunk>& chunks, | ||
| 442 | const base::RecTensor& values) const { | ||
| 443 | ✗ | for (const auto& chunk : chunks) { | |
| 444 | ✗ | const int64_t dim = values.shape(1); | |
| 445 | base::RecTensor chunk_values( | ||
| 446 | ✗ | {static_cast<int64_t>(chunk.positions.size()), dim}, | |
| 447 | ✗ | base::DataType::FLOAT32); | |
| 448 | ✗ | std::size_t chunk_row = 0; | |
| 449 | ✗ | for (const std::size_t position : chunk.positions) { | |
| 450 | ✗ | if (position >= static_cast<std::size_t>(values.shape(0))) { | |
| 451 | ✗ | throw std::invalid_argument("RDMA PUT chunk position is out of range"); | |
| 452 | } | ||
| 453 | ✗ | std::memcpy( | |
| 454 | ✗ | chunk_values.data_as<float>() + | |
| 455 | ✗ | chunk_row++ * static_cast<std::size_t>(dim), | |
| 456 | ✗ | values.data_as<float>() + position * static_cast<std::size_t>(dim), | |
| 457 | ✗ | static_cast<std::size_t>(dim) * sizeof(float)); | |
| 458 | } | ||
| 459 | const int rc = | ||
| 460 | ✗ | ClientAt(chunk.client_index) | |
| 461 | ✗ | ->PutParameter( | |
| 462 | ✗ | base::ConstArray<std::uint64_t>(chunk.keys), chunk_values); | |
| 463 | ✗ | if (rc != 0) { | |
| 464 | ✗ | return rc; | |
| 465 | } | ||
| 466 | ✗ | } | |
| 467 | ✗ | return 0; | |
| 468 | } | ||
| 469 | |||
| 470 | ✗ | int RdmaDistributedExecutor::UpdateParameter( | |
| 471 | const std::string& table_name, | ||
| 472 | const std::vector<shard_routing::ShardChunk>& chunks, | ||
| 473 | const base::RecTensor& grads) const { | ||
| 474 | ✗ | for (const auto& chunk : chunks) { | |
| 475 | ✗ | const int64_t dim = grads.shape(1); | |
| 476 | base::RecTensor chunk_grads( | ||
| 477 | ✗ | {static_cast<int64_t>(chunk.positions.size()), dim}, | |
| 478 | ✗ | base::DataType::FLOAT32); | |
| 479 | ✗ | std::size_t chunk_row = 0; | |
| 480 | ✗ | for (const std::size_t position : chunk.positions) { | |
| 481 | ✗ | if (position >= static_cast<std::size_t>(grads.shape(0))) { | |
| 482 | ✗ | throw std::invalid_argument( | |
| 483 | ✗ | "RDMA UPDATE chunk position is out of range"); | |
| 484 | } | ||
| 485 | ✗ | std::memcpy( | |
| 486 | ✗ | chunk_grads.data_as<float>() + | |
| 487 | ✗ | chunk_row++ * static_cast<std::size_t>(dim), | |
| 488 | ✗ | grads.data_as<float>() + position * static_cast<std::size_t>(dim), | |
| 489 | ✗ | static_cast<std::size_t>(dim) * sizeof(float)); | |
| 490 | } | ||
| 491 | const int rc = | ||
| 492 | ✗ | ClientAt(chunk.client_index) | |
| 493 | ✗ | ->UpdateParameter(table_name, | |
| 494 | ✗ | base::ConstArray<std::uint64_t>(chunk.keys), | |
| 495 | chunk_grads); | ||
| 496 | ✗ | if (rc != 0) { | |
| 497 | ✗ | return rc; | |
| 498 | } | ||
| 499 | ✗ | } | |
| 500 | ✗ | return 0; | |
| 501 | } | ||
| 502 | |||
| 503 | ✗ | int RdmaDistributedExecutor::WaitUpdateRpcs( | |
| 504 | std::vector<std::pair<int, int>>* shard_rpcs) const { | ||
| 505 | ✗ | int result = 0; | |
| 506 | ✗ | std::size_t index = 0; | |
| 507 | try { | ||
| 508 | ✗ | for (; index < shard_rpcs->size(); ++index) { | |
| 509 | ✗ | const auto [client_index, rpc_id] = (*shard_rpcs)[index]; | |
| 510 | ✗ | auto* client = ClientAt(client_index); | |
| 511 | ✗ | if (client->WaitUpdateParameter(rpc_id) != 0) { | |
| 512 | ✗ | result = -1; | |
| 513 | } | ||
| 514 | } | ||
| 515 | ✗ | } catch (...) { | |
| 516 | ✗ | std::exception_ptr first_error = std::current_exception(); | |
| 517 | ✗ | for (++index; index < shard_rpcs->size(); ++index) { | |
| 518 | try { | ||
| 519 | ✗ | const auto [client_index, rpc_id] = (*shard_rpcs)[index]; | |
| 520 | ✗ | ClientAt(client_index)->RevokeRPCResource(rpc_id); | |
| 521 | ✗ | } catch (...) { | |
| 522 | // Preserve the operation error while still releasing later shards. | ||
| 523 | ✗ | } | |
| 524 | } | ||
| 525 | ✗ | shard_rpcs->clear(); | |
| 526 | ✗ | std::rethrow_exception(first_error); | |
| 527 | ✗ | } | |
| 528 | ✗ | shard_rpcs->clear(); | |
| 529 | ✗ | return result; | |
| 530 | } | ||
| 531 | |||
| 532 | ✗ | std::uint64_t RdmaDistributedExecutor::SubmitUpdateParameterFlat( | |
| 533 | const std::string& table_name, | ||
| 534 | base::ConstArray<std::uint64_t> keys, | ||
| 535 | const float* grads, | ||
| 536 | std::size_t embedding_dim, | ||
| 537 | const std::vector<shard_routing::ShardChunk>& chunks, | ||
| 538 | std::size_t max_in_flight) { | ||
| 539 | ✗ | if ((keys.Size() > 0 && grads == nullptr) || embedding_dim == 0 || | |
| 540 | max_in_flight == 0) { | ||
| 541 | ✗ | throw std::invalid_argument("invalid RDMA flat UPDATE request"); | |
| 542 | } | ||
| 543 | |||
| 544 | ✗ | std::vector<std::pair<int, int>> pending; | |
| 545 | ✗ | auto drain = [this, &pending]() { | |
| 546 | ✗ | if (WaitUpdateRpcs(&pending) != 0) { | |
| 547 | ✗ | throw std::runtime_error("RDMA embedding update failed"); | |
| 548 | } | ||
| 549 | ✗ | }; | |
| 550 | |||
| 551 | try { | ||
| 552 | ✗ | for (const auto& chunk : chunks) { | |
| 553 | ✗ | auto* client = ClientAt(chunk.client_index); | |
| 554 | ✗ | if (chunk.keys.size() != chunk.positions.size()) { | |
| 555 | ✗ | throw std::runtime_error("RDMA UPDATE chunk shape mismatch"); | |
| 556 | } | ||
| 557 | ✗ | bool contiguous = !chunk.positions.empty(); | |
| 558 | ✗ | for (std::size_t i = 1; i < chunk.positions.size(); ++i) { | |
| 559 | ✗ | contiguous = | |
| 560 | ✗ | contiguous && chunk.positions[i] == chunk.positions.front() + i; | |
| 561 | } | ||
| 562 | const int rpc_id = | ||
| 563 | contiguous | ||
| 564 | ✗ | ? client->SubmitUpdateParameterFlat( | |
| 565 | table_name, | ||
| 566 | ✗ | base::ConstArray<std::uint64_t>(chunk.keys), | |
| 567 | ✗ | grads + chunk.positions.front() * embedding_dim, | |
| 568 | embedding_dim) | ||
| 569 | ✗ | : client->SubmitUpdateParameterFlatGather( | |
| 570 | table_name, | ||
| 571 | keys.Data(), | ||
| 572 | grads, | ||
| 573 | ✗ | keys.Size(), | |
| 574 | embedding_dim, | ||
| 575 | chunk.positions.data(), | ||
| 576 | ✗ | chunk.positions.size()); | |
| 577 | ✗ | if (rpc_id < 0) { | |
| 578 | ✗ | throw std::runtime_error("failed to submit RDMA embedding update"); | |
| 579 | } | ||
| 580 | ✗ | pending.emplace_back(chunk.client_index, rpc_id); | |
| 581 | ✗ | if (pending.size() >= max_in_flight) { | |
| 582 | ✗ | drain(); | |
| 583 | } | ||
| 584 | } | ||
| 585 | ✗ | } catch (...) { | |
| 586 | try { | ||
| 587 | ✗ | drain(); | |
| 588 | ✗ | } catch (...) { | |
| 589 | // Preserve the operation error that initiated cleanup. | ||
| 590 | ✗ | } | |
| 591 | ✗ | throw; | |
| 592 | ✗ | } | |
| 593 | |||
| 594 | ✗ | std::lock_guard<std::mutex> guard(updates_mu_); | |
| 595 | ✗ | const std::uint64_t update_id = next_update_id_++; | |
| 596 | ✗ | pending_updates_.emplace( | |
| 597 | ✗ | update_id, PendingUpdate{std::move(pending), std::this_thread::get_id()}); | |
| 598 | ✗ | return update_id; | |
| 599 | ✗ | } | |
| 600 | |||
| 601 | ✗ | int RdmaDistributedExecutor::WaitUpdateParameterFlat(std::uint64_t update_id) { | |
| 602 | ✗ | PendingUpdate update; | |
| 603 | { | ||
| 604 | ✗ | std::lock_guard<std::mutex> guard(updates_mu_); | |
| 605 | ✗ | const auto it = pending_updates_.find(update_id); | |
| 606 | ✗ | if (it == pending_updates_.end()) { | |
| 607 | ✗ | throw std::runtime_error( | |
| 608 | ✗ | "unknown or already consumed RDMA update handle"); | |
| 609 | } | ||
| 610 | ✗ | if (it->second.owner != std::this_thread::get_id()) { | |
| 611 | ✗ | throw std::runtime_error( | |
| 612 | ✗ | "RDMA update handle must be waited by its submitting thread"); | |
| 613 | } | ||
| 614 | ✗ | update = std::move(it->second); | |
| 615 | ✗ | pending_updates_.erase(it); | |
| 616 | ✗ | } | |
| 617 | ✗ | return WaitUpdateRpcs(&update.shard_rpcs); | |
| 618 | ✗ | } | |
| 619 | |||
| 620 | ✗ | bool RdmaDistributedExecutor::WaitBatch(int batch_id) { | |
| 621 | ✗ | std::vector<shard_routing::PendingShardRpc> shard_rpcs; | |
| 622 | { | ||
| 623 | ✗ | std::lock_guard<std::mutex> guard(batches_mu_); | |
| 624 | ✗ | auto it = batches_.find(batch_id); | |
| 625 | ✗ | if (it == batches_.end()) { | |
| 626 | ✗ | throw std::runtime_error("unknown RDMA distributed batch handle"); | |
| 627 | } | ||
| 628 | ✗ | if (it->second.assembled) { | |
| 629 | ✗ | return it->second.status_code == | |
| 630 | ✗ | static_cast<std::int32_t>(petps::RpcStatus::kOk); | |
| 631 | } | ||
| 632 | ✗ | shard_rpcs = it->second.shard_rpcs; | |
| 633 | ✗ | } | |
| 634 | ✗ | WaitShardRpcs(shard_rpcs); | |
| 635 | ✗ | std::lock_guard<std::mutex> guard(batches_mu_); | |
| 636 | ✗ | auto it = batches_.find(batch_id); | |
| 637 | ✗ | if (it == batches_.end()) { | |
| 638 | ✗ | throw std::runtime_error("RDMA distributed batch released while waiting"); | |
| 639 | } | ||
| 640 | ✗ | return FinalizeBatch(&it->second); | |
| 641 | ✗ | } | |
| 642 | |||
| 643 | ✗ | bool RdmaDistributedExecutor::HasGetSlotCapacity( | |
| 644 | const std::vector<shard_routing::ShardChunk>& chunks, | ||
| 645 | std::size_t max_in_flight) const { | ||
| 646 | ✗ | if (max_in_flight == 0) { | |
| 647 | ✗ | return false; | |
| 648 | } | ||
| 649 | ✗ | std::vector<std::size_t> required(clients_.size(), 0); | |
| 650 | ✗ | for (const auto& chunk : chunks) { | |
| 651 | ✗ | auto& count = required.at(static_cast<std::size_t>(chunk.client_index)); | |
| 652 | // Each submission window has at most max_in_flight RPCs across all shards. | ||
| 653 | ✗ | if (count < max_in_flight) { | |
| 654 | ✗ | ++count; | |
| 655 | } | ||
| 656 | } | ||
| 657 | ✗ | for (std::size_t client = 0; client < required.size(); ++client) { | |
| 658 | ✗ | if (required[client] > 0 && !ClientAt(static_cast<int>(client)) | |
| 659 | ✗ | ->HasAvailableSlots(required[client])) { | |
| 660 | ✗ | return false; | |
| 661 | } | ||
| 662 | } | ||
| 663 | ✗ | return true; | |
| 664 | ✗ | } | |
| 665 | |||
| 666 | ✗ | void RdmaDistributedExecutor::DrainBatchSlots() { | |
| 667 | ✗ | std::lock_guard<std::mutex> guard(batches_mu_); | |
| 668 | ✗ | for (auto& entry : batches_) { | |
| 669 | ✗ | auto& batch = entry.second; | |
| 670 | ✗ | WaitShardRpcs(batch.shard_rpcs); | |
| 671 | ✗ | FinalizeBatch(&batch); | |
| 672 | ✗ | ReleaseShardRpcs(&batch.shard_rpcs); | |
| 673 | } | ||
| 674 | ✗ | } | |
| 675 | |||
| 676 | ✗ | void RdmaDistributedExecutor::ReleaseBatch(int batch_id) { | |
| 677 | ✗ | shard_routing::BatchRequest batch; | |
| 678 | { | ||
| 679 | ✗ | std::lock_guard<std::mutex> guard(batches_mu_); | |
| 680 | ✗ | auto it = batches_.find(batch_id); | |
| 681 | ✗ | if (it == batches_.end()) { | |
| 682 | ✗ | throw std::runtime_error( | |
| 683 | ✗ | "unknown or released RDMA distributed batch handle"); | |
| 684 | } | ||
| 685 | ✗ | batch = std::move(it->second); | |
| 686 | ✗ | batches_.erase(it); | |
| 687 | ✗ | } | |
| 688 | ✗ | ReleaseShardRpcs(&batch.shard_rpcs); | |
| 689 | ✗ | RecycleReceiveBuffers(&batch.receive_buffers); | |
| 690 | ✗ | } | |
| 691 | |||
| 692 | } // namespace recstore | ||
| 693 |