update

Former-commit-id: 97b0367c3d63c40a2c2ba6a4b272796ed706088a

oneflow/core/comm_network/epoll/epoll_comm_network.cpp

@@ -71,7 +71,8 @@ void EpollCommNet::SendActorMsg(int64_t dst_machine_id, const ActorMsg& actor_ms
   SocketMsg msg;
   msg.msg_type = SocketMsgType::kActor;
   msg.actor_msg = actor_msg;
-  int32_t link_i = std::uniform_int_distribution<int32_t>(0, epoll_conf_.link_num())(random_gen_);
+  int32_t link_i =
+      std::uniform_int_distribution<int32_t>(0, epoll_conf_.link_num() - 1)(random_gen_);
   GetSocketHelper(dst_machine_id, link_i)->AsyncWrite(msg);
 }
 
@@ -90,8 +91,10 @@ void EpollCommNet::SendSocketMsg(int64_t dst_machine_id, const SocketMsg& total_
     total_byte_size -= offset;
     SocketMsg msg;
     msg.msg_type = total_msg.msg_type;
-    msg.request_read_msg.src_token = NewMemDesc(src_mem_desc->mem_ptr + link_i * offset, byte_size);
-    msg.request_read_msg.dst_token = NewMemDesc(dst_mem_desc->mem_ptr + link_i * offset, byte_size);
+    msg.request_read_msg.src_token = total_msg.request_read_msg.src_token;
+    msg.request_read_msg.dst_token = total_msg.request_read_msg.dst_token;
+    msg.request_read_msg.offset = link_i * offset;
+    msg.request_read_msg.byte_size = byte_size;
     msg.request_read_msg.read_id = total_msg.request_read_msg.read_id;
     msg.request_read_msg.part_num = part_num;
     GetSocketHelper(dst_machine_id, link_i)->AsyncWrite(msg);
@@ -189,7 +192,7 @@ void EpollCommNet::InitSockets() {
 }
 
 SocketHelper* EpollCommNet::GetSocketHelper(int64_t machine_id, int32_t link_index) {
-  int sockfd = machine_id2sockfds_.at(machine_id * epoll_conf_.link_num() + link_index);
+  int sockfd = machine_id2sockfds_[machine_id * epoll_conf_.link_num() + link_index];
   return sockfd2helper_.at(sockfd);
 }
 
@@ -201,13 +204,14 @@ void EpollCommNet::DoRead(void* read_id, int64_t src_machine_id, void* src_token
   msg.request_write_msg.dst_machine_id = Global<MachineCtx>::Get()->this_machine_id();
   msg.request_write_msg.dst_token = dst_token;
   msg.request_write_msg.read_id = read_id;
-  int32_t link_i = std::uniform_int_distribution<int32_t>(0, epoll_conf_.link_num())(random_gen_);
+  int32_t link_i =
+      std::uniform_int_distribution<int32_t>(0, epoll_conf_.link_num() - 1)(random_gen_);
   GetSocketHelper(src_machine_id, link_i)->AsyncWrite(msg);
 }
 
 void EpollCommNet::PartReadDone(void* read_id, int32_t part_num) {
-  int32_t& part_read_done_cnt = read_id2part_done_cnt_.at(read_id);
   std::unique_lock<std::mutex> lck(part_done_cnt_mtx_);
+  int32_t& part_read_done_cnt = read_id2part_done_cnt_.at(read_id);
   part_read_done_cnt++;
   if (part_read_done_cnt == part_num) {
     ReadDone(read_id);

oneflow/core/comm_network/epoll/socket_message.h

@@ -40,6 +40,8 @@ struct RequestWriteMsg {
 struct RequestReadMsg {
   void* src_token;
   void* dst_token;
+  int64_t offset;
+  int64_t byte_size;
   void* read_id;
   int32_t part_num;
 };

oneflow/core/comm_network/epoll/socket_read_helper.cpp

@@ -82,8 +82,8 @@ void SocketReadHelper::SetStatusWhenRequestWriteMsgHeadDone() {
 
 void SocketReadHelper::SetStatusWhenRequestReadMsgHeadDone() {
   auto mem_desc = static_cast<const SocketMemDesc*>(cur_msg_.request_read_msg.dst_token);
-  read_ptr_ = reinterpret_cast<char*>(mem_desc->mem_ptr);
-  read_size_ = mem_desc->byte_size;
+  read_ptr_ = reinterpret_cast<char*>(mem_desc->mem_ptr) + cur_msg_.request_read_msg.offset;
+  read_size_ = cur_msg_.request_read_msg.byte_size;
   cur_read_handle_ = &SocketReadHelper::MsgBodyReadHandle;
 }
 

oneflow/core/comm_network/epoll/socket_write_helper.cpp

@@ -116,8 +116,9 @@ void SocketWriteHelper::SetStatusWhenRequestWriteMsgHeadDone() {
 void SocketWriteHelper::SetStatusWhenRequestReadMsgHeadDone() {
   const void* src_token = cur_msg_.request_read_msg.src_token;
   auto src_mem_desc = static_cast<const SocketMemDesc*>(src_token);
-  write_ptr_ = reinterpret_cast<const char*>(src_mem_desc->mem_ptr);
-  write_size_ = src_mem_desc->byte_size;
+  write_ptr_ =
+      reinterpret_cast<const char*>(src_mem_desc->mem_ptr) + cur_msg_.request_read_msg.offset;
+  write_size_ = cur_msg_.request_read_msg.byte_size;
   cur_write_handle_ = &SocketWriteHelper::MsgBodyWriteHandle;
 }