| 270 | } |
| 271 | |
| 272 | void GdrMemoryManager::Run() { |
| 273 | stopped_ = false; |
| 274 | while (!stopped_) { |
| 275 | rdma_cm_id* id = nullptr; |
| 276 | // Accept incoming connections |
| 277 | if (!rdma_get_request(listening_.get(), &id)) { |
| 278 | if (!rdma_accept(id, nullptr)) { |
| 279 | LOG(INFO) << "Accepted new RDMA connection"; |
| 280 | for (int i = 0; i < 1024; i++) { |
| 281 | if (rdma_post_recvv(id, nullptr, nullptr, 0)) { |
| 282 | LOG(ERROR) << strerror(errno) << ": rdma_post_recvv failed"; |
| 283 | EndpointDeleter(id); |
| 284 | continue; |
| 285 | } |
| 286 | } |
| 287 | server_clients_.push_back({id, EndpointDeleter}); |
| 288 | } |
| 289 | } |
| 290 | // Polling server side work completions |
| 291 | for (const auto& client : server_clients_) { |
| 292 | ibv_wc wc[32]; |
| 293 | int ret = ibv_poll_cq(client->recv_cq, 32, wc); |
| 294 | if (ret < 0) { |
| 295 | LOG(ERROR) << "ibv_poll_cq failed"; |
| 296 | continue; |
| 297 | } |
| 298 | for (int i = 0; i < ret; i++) { |
| 299 | if (wc[i].opcode != IBV_WC_RECV_RDMA_WITH_IMM) { |
| 300 | LOG(ERROR) << "Received unknown operation " << wc[i].opcode; |
| 301 | } |
| 302 | if (wc[i].status != 0) { |
| 303 | LOG(ERROR) << ibv_wc_status_str(wc[i].status); |
| 304 | } |
| 305 | TensorKey tensor_key = ntohl(wc[i].imm_data); |
| 306 | |
| 307 | if (rdma_post_recvv(client.get(), nullptr, nullptr, 0)) { |
| 308 | perror("rdma_post_recvv"); |
| 309 | LOG(ERROR) << "rdma_post_recvv failed"; |
| 310 | } |
| 311 | |
| 312 | mutex_lock l(buf_mu_); |
| 313 | auto iter = tensor_buffers_.find(tensor_key); |
| 314 | if (iter == std::end(tensor_buffers_)) { |
| 315 | LOG(ERROR) << "Cannot find tensor buffer for tensor key " |
| 316 | << tensor_key; |
| 317 | } else { |
| 318 | const TensorBuffer* buffer = iter->second; |
| 319 | buffer->Unref(); |
| 320 | tensor_buffers_.erase(iter); |
| 321 | } |
| 322 | } |
| 323 | } |
| 324 | // Polling client side work completions |
| 325 | if (client_mu_.try_lock()) { |
| 326 | for (const auto& client : clients_) { |
| 327 | ibv_wc wc[32]; |
| 328 | int ret = ibv_poll_cq(client.second->send_cq, 32, wc); |
| 329 | for (int i = 0; i < ret; i++) { |