RdmaTensorResponse Represents a single tensor response.
| 306 | // RdmaTensorResponse |
| 307 | // Represents a single tensor response. |
| 308 | class RdmaTensorResponse { |
| 309 | public: |
| 310 | // Creates a response for request message. |
| 311 | RdmaTensorResponse(RdmaChannel* channel, const RdmaMessage& rm) |
| 312 | : channel_(channel), rm_(rm) {} |
| 313 | |
| 314 | void Update(const RdmaMessage& rm) { rm_ = rm; } |
| 315 | |
| 316 | // Start the tensor response sequence. |
| 317 | // |
| 318 | // 1. Find the tensor in the local tag-match table and invoke RecvHandler. |
| 319 | // (Using RecvLocalAsync()). |
| 320 | // 2. Compare the tensor's meta-data to the meta-data in the message (taken |
| 321 | // from the requester's local cache). |
| 322 | // If meta-data changed: |
| 323 | // a. Clone the tensor to be sent later. |
| 324 | // b. Send a meta-data update message and wait for re-request. |
| 325 | // Else: |
| 326 | // a. Send the tensor's content (using direct RDMA write). |
| 327 | void Start(); |
| 328 | |
| 329 | // Resume the response sequence, after a re-request. |
| 330 | // |
| 331 | // 1. Send the tensor's content that was cloned earlier. |
| 332 | void Resume(); |
| 333 | |
| 334 | // Destroy the response's resources and remove it from the pending list. |
| 335 | void Destroy(); |
| 336 | |
| 337 | private: |
| 338 | void RecvHandler(Rendezvous::ParsedKey parsed, |
| 339 | const Rendezvous::Args& send_args, |
| 340 | const Rendezvous::Args& recv_args, const Tensor& in, |
| 341 | bool is_dead); |
| 342 | void Clone(const Tensor& in, const TensorProto& proto, bool is_dead); |
| 343 | void Send(const Tensor& in, const TensorProto& proto, bool is_dead, |
| 344 | const Status& status); |
| 345 | bool TensorMetaDataChanged(const Tensor& in, bool is_dead); |
| 346 | Status PrepareRecvTensor(const Rendezvous::ParsedKey& parsed, |
| 347 | Device** src_dev); |
| 348 | void SendMetaData(const Tensor& in, const TensorProto& proto, bool is_dead); |
| 349 | void SendContent(const Tensor& in, const TensorProto& proto, bool is_dead); |
| 350 | void SendErrorStatus(const Status& status); |
| 351 | |
| 352 | RdmaChannel* channel_; |
| 353 | RdmaMessage rm_; // The request message |
| 354 | Device* src_dev_ = nullptr; |
| 355 | TensorBuffer* src_buffer_ = nullptr; |
| 356 | void* src_addr_ = nullptr; |
| 357 | ibv_mr* mr_ = nullptr; |
| 358 | uint64_t checksum_ = 0; |
| 359 | bool meta_data_changed_ = false; |
| 360 | |
| 361 | // Re-item: |
| 362 | TensorProto* proto_ = nullptr; |
| 363 | Tensor* tensor_ = nullptr; |
| 364 | bool is_dead_ = false; |
| 365 | }; |