linkers.h

#ifndef LIGHTGBM_NETWORK_LINKERS_H_
#define LIGHTGBM_NETWORK_LINKERS_H_
 
#include <LightGBM/network.h>
#include <LightGBM/meta.h>
#include <LightGBM/config.h>
 
 
#include <algorithm>
#include <chrono>
#include <ctime>
#include <thread>
#include <vector>
#include <string>
#include <memory>
 
#ifdef USE_SOCKET
#include "socket_wrapper.hpp"
#include <LightGBM/utils/common.h>
#endif
 
#ifdef USE_MPI
#include <mpi.h>
#define MPI_SAFE_CALL(mpi_return) CHECK((mpi_return) == MPI_SUCCESS)
#endif
 
namespace LightGBM {
 
class Linkers {
public:
  Linkers() {
    is_init_ = false;
  }
  explicit Linkers(Config config);
  ~Linkers();
  inline void Recv(int rank, char* data, int len) const;
 
  inline void Recv(int rank, char* data, int64_t len) const;
 
  inline void Send(int rank, char* data, int len) const;
 
  inline void Send(int rank, char* data, int64_t len) const;
  inline void SendRecv(int send_rank, char* send_data, int send_len,
                       int recv_rank, char* recv_data, int recv_len);
 
  inline void SendRecv(int send_rank, char* send_data, int64_t send_len,
                       int recv_rank, char* recv_data, int64_t recv_len);
  inline int rank();
  inline int num_machines();
  inline const BruckMap& bruck_map();
  inline const RecursiveHalvingMap& recursive_halving_map();
 
  #ifdef USE_SOCKET
  void TryBind(int port);
  void SetLinker(int rank, const TcpSocket& socket);
  void ListenThread(int incoming_cnt);
  void Construct();
  void ParseMachineList(const std::string& machines, const std::string& filename);
  bool CheckLinker(int rank);
  void PrintLinkers();
 
  #endif  // USE_SOCKET
 
 
private:
  int rank_;
  int num_machines_;
  BruckMap bruck_map_;
  RecursiveHalvingMap recursive_halving_map_;
 
  std::chrono::duration<double, std::milli> network_time_;
 
  bool is_init_;
 
  #ifdef USE_SOCKET
  std::vector<std::string> client_ips_;
  std::vector<int> client_ports_;
  int socket_timeout_;
  int local_listen_port_;
  std::vector<std::unique_ptr<TcpSocket>> linkers_;
  std::unique_ptr<TcpSocket> listener_;
  #endif  // USE_SOCKET
};
 
 
inline int Linkers::rank() {
  return rank_;
}
 
inline int Linkers::num_machines() {
  return num_machines_;
}
 
inline const BruckMap& Linkers::bruck_map() {
  return bruck_map_;
}
 
inline const RecursiveHalvingMap& Linkers::recursive_halving_map() {
  return recursive_halving_map_;
}
 
inline void Linkers::Recv(int rank, char* data, int64_t len) const {
  int64_t used = 0;
  do {
    int cur_size = static_cast<int>(std::min<int64_t>(len - used, INT32_MAX));
    Recv(rank, data + used, cur_size);
    used += cur_size;
  } while (used < len);
}
 
inline void Linkers::Send(int rank, char* data, int64_t len) const {
  int64_t used = 0;
  do {
    int cur_size = static_cast<int>(std::min<int64_t>(len - used, INT32_MAX));
    Send(rank, data + used, cur_size);
    used += cur_size;
  } while (used < len);
}
 
inline void Linkers::SendRecv(int send_rank, char* send_data, int64_t send_len,
                              int recv_rank, char* recv_data, int64_t recv_len) {
  auto start_time = std::chrono::high_resolution_clock::now();
  std::thread send_worker(
    [this, send_rank, send_data, send_len]() {
    Send(send_rank, send_data, send_len);
  });
  Recv(recv_rank, recv_data, recv_len);
  send_worker.join();
  // wait for send complete
  auto end_time = std::chrono::high_resolution_clock::now();
  // output used time on each iteration
  network_time_ += std::chrono::duration<double, std::milli>(end_time - start_time);
}
 
#ifdef USE_SOCKET
 
inline void Linkers::Recv(int rank, char* data, int len) const {
  int recv_cnt = 0;
  while (recv_cnt < len) {
    recv_cnt += linkers_[rank]->Recv(data + recv_cnt,
      // len - recv_cnt
      std::min(len - recv_cnt, SocketConfig::kMaxReceiveSize));
  }
}
 
inline void Linkers::Send(int rank, char* data, int len) const {
  if (len <= 0) {
    return;
  }
  int send_cnt = 0;
  while (send_cnt < len) {
    send_cnt += linkers_[rank]->Send(data + send_cnt, len - send_cnt);
  }
}
 
inline void Linkers::SendRecv(int send_rank, char* send_data, int send_len,
                              int recv_rank, char* recv_data, int recv_len) {
  auto start_time = std::chrono::high_resolution_clock::now();
  if (send_len < SocketConfig::kSocketBufferSize) {
    // if buffer is enough, send will non-blocking
    Send(send_rank, send_data, send_len);
    Recv(recv_rank, recv_data, recv_len);
  } else {
    // if buffer is not enough, use another thread to send, since send will be blocking
    std::thread send_worker(
      [this, send_rank, send_data, send_len]() {
      Send(send_rank, send_data, send_len);
    });
    Recv(recv_rank, recv_data, recv_len);
    send_worker.join();
  }
  // wait for send complete
  auto end_time = std::chrono::high_resolution_clock::now();
  // output used time on each iteration
  network_time_ += std::chrono::duration<double, std::milli>(end_time - start_time);
}
 
#endif  // USE_SOCKET
 
#ifdef USE_MPI
 
inline void Linkers::Recv(int rank, char* data, int len) const {
  MPI_Status status;
  int read_cnt = 0;
  while (read_cnt < len) {
    MPI_SAFE_CALL(MPI_Recv(data + read_cnt, len - read_cnt, MPI_BYTE, rank, MPI_ANY_TAG, MPI_COMM_WORLD, &status));
    int cur_cnt;
    MPI_SAFE_CALL(MPI_Get_count(&status, MPI_BYTE, &cur_cnt));
    read_cnt += cur_cnt;
  }
}
 
inline void Linkers::Send(int rank, char* data, int len) const {
  if (len <= 0) {
    return;
  }
  MPI_Status status;
  MPI_Request send_request;
  MPI_SAFE_CALL(MPI_Isend(data, len, MPI_BYTE, rank, 0, MPI_COMM_WORLD, &send_request));
  MPI_SAFE_CALL(MPI_Wait(&send_request, &status));
}
 
inline void Linkers::SendRecv(int send_rank, char* send_data, int send_len,
                              int recv_rank, char* recv_data, int recv_len) {
  MPI_Request send_request;
  // send first, non-blocking
  MPI_SAFE_CALL(MPI_Isend(send_data, send_len, MPI_BYTE, send_rank, 0, MPI_COMM_WORLD, &send_request));
  // then receive, blocking
  MPI_Status status;
  int read_cnt = 0;
  while (read_cnt < recv_len) {
    MPI_SAFE_CALL(MPI_Recv(recv_data + read_cnt, recv_len - read_cnt, MPI_BYTE, recv_rank, 0, MPI_COMM_WORLD, &status));
    int cur_cnt;
    MPI_SAFE_CALL(MPI_Get_count(&status, MPI_BYTE, &cur_cnt));
    read_cnt += cur_cnt;
  }
  // wait for send complete
  MPI_SAFE_CALL(MPI_Wait(&send_request, &status));
}
 
#endif  // USE_MPI
}  // namespace LightGBM
#endif   // LightGBM_NETWORK_LINKERS_H_