linkers_socket.cpp 6.79 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
#ifdef USE_SOCKET

3
#include <LightGBM/config.h>
Guolin Ke's avatar
Guolin Ke committed
4
5
6
#include <LightGBM/utils/common.h>
#include <LightGBM/utils/text_reader.h>

7
8
#include <string>
#include <chrono>
Guolin Ke's avatar
Guolin Ke committed
9
#include <cstring>
10
#include <thread>
Guolin Ke's avatar
Guolin Ke committed
11
12
13
#include <unordered_map>
#include <unordered_set>
#include <vector>
14
15

#include "linkers.h"
Guolin Ke's avatar
Guolin Ke committed
16
17
18

namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
19
Linkers::Linkers(Config config) {
20
  is_init_ = false;
Guolin Ke's avatar
Guolin Ke committed
21
22
23
24
25
26
27
  // start up socket
  TcpSocket::Startup();
  network_time_ = std::chrono::duration<double, std::milli>(0);
  num_machines_ = config.num_machines;
  local_listen_port_ = config.local_listen_port;
  socket_timeout_ = config.time_out;
  rank_ = -1;
zhangyafeikimi's avatar
zhangyafeikimi committed
28
  // parse clients from file
29
  ParseMachineList(config.machines, config.machine_list_filename);
Guolin Ke's avatar
Guolin Ke committed
30
31
32
33
34
35
36
37
38
39
40
41
42

  if (rank_ == -1) {
    // get ip list of local machine
    std::unordered_set<std::string> local_ip_list = TcpSocket::GetLocalIpList();
    // get local rank
    for (size_t i = 0; i < client_ips_.size(); ++i) {
      if (local_ip_list.count(client_ips_[i]) > 0 && client_ports_[i] == local_listen_port_) {
        rank_ = static_cast<int>(i);
        break;
      }
    }
  }
  if (rank_ == -1) {
43
    Log::Fatal("Machine list file doesn't contain the local machine");
Guolin Ke's avatar
Guolin Ke committed
44
45
  }
  // construct listener
Guolin Ke's avatar
Guolin Ke committed
46
  listener_ = std::unique_ptr<TcpSocket>(new TcpSocket());
Guolin Ke's avatar
Guolin Ke committed
47
48
49
50
51
  TryBind(local_listen_port_);

  for (int i = 0; i < num_machines_; ++i) {
    linkers_.push_back(nullptr);
  }
52

Guolin Ke's avatar
Guolin Ke committed
53
54
55
56
57
58
59
60
  // construct communication topo
  bruck_map_ = BruckMap::Construct(rank_, num_machines_);
  recursive_halving_map_ = RecursiveHalvingMap::Construct(rank_, num_machines_);

  // construct linkers
  Construct();
  // free listener
  listener_->Close();
61
  is_init_ = true;
Guolin Ke's avatar
Guolin Ke committed
62
63
64
}

Linkers::~Linkers() {
65
66
67
68
69
  if (is_init_) {
    for (size_t i = 0; i < linkers_.size(); ++i) {
      if (linkers_[i] != nullptr) {
        linkers_[i]->Close();
      }
Guolin Ke's avatar
Guolin Ke committed
70
    }
71
72
    TcpSocket::Finalize();
    Log::Info("Finished linking network in %f seconds", network_time_ * 1e-3);
Guolin Ke's avatar
Guolin Ke committed
73
74
75
  }
}

76
77
78
79
80
81
82
83
84
85
86
void Linkers::ParseMachineList(const std::string& machines, const std::string& filename) {
  std::vector<std::string> lines;
  if (machines.empty()) {
    TextReader<size_t> machine_list_reader(filename.c_str(), false);
    machine_list_reader.ReadAllLines();
    if (machine_list_reader.Lines().empty()) {
      Log::Fatal("Machine list file %s doesn't exist", filename.c_str());
    }
    lines = machine_list_reader.Lines();
  } else {
    lines = Common::Split(machines.c_str(), ',');
Guolin Ke's avatar
Guolin Ke committed
87
  }
88
  for (auto& line : lines) {
Guolin Ke's avatar
Guolin Ke committed
89
90
91
92
93
94
    line = Common::Trim(line);
    if (line.find_first_of("rank=") != std::string::npos) {
      std::vector<std::string> str_after_split = Common::Split(line.c_str(), '=');
      Common::Atoi(str_after_split[1].c_str(), &rank_);
      continue;
    }
95
    std::vector<std::string> str_after_split = Common::Split(line.c_str(), ' ');
Guolin Ke's avatar
Guolin Ke committed
96
    if (str_after_split.size() != 2) {
97
98
99
100
      str_after_split = Common::Split(line.c_str(), ':');
      if (str_after_split.size() != 2) {
        continue;
      }
Guolin Ke's avatar
Guolin Ke committed
101
102
    }
    if (client_ips_.size() >= static_cast<size_t>(num_machines_)) {
103
      Log::Warning("machine_list size is larger than the parameter num_machines, ignoring redundant entries");
Guolin Ke's avatar
Guolin Ke committed
104
105
106
107
108
109
110
      break;
    }
    str_after_split[0] = Common::Trim(str_after_split[0]);
    str_after_split[1] = Common::Trim(str_after_split[1]);
    client_ips_.push_back(str_after_split[0]);
    client_ports_.push_back(atoi(str_after_split[1].c_str()));
  }
Guolin Ke's avatar
Guolin Ke committed
111
  if (client_ips_.empty()) {
112
113
    Log::Fatal("Cannot find any ip and port.\n"
               "Please check machine_list_filename or machines parameter");
114
  }
Guolin Ke's avatar
Guolin Ke committed
115
  if (client_ips_.size() != static_cast<size_t>(num_machines_)) {
116
    Log::Warning("World size is larger than the machine_list size, change world size to %d", client_ips_.size());
Guolin Ke's avatar
Guolin Ke committed
117
118
119
120
121
    num_machines_ = static_cast<int>(client_ips_.size());
  }
}

void Linkers::TryBind(int port) {
122
  Log::Info("Trying to bind port %d...", port);
Guolin Ke's avatar
Guolin Ke committed
123
  if (listener_->Bind(port)) {
124
    Log::Info("Binding port %d succeeded", port);
Guolin Ke's avatar
Guolin Ke committed
125
  } else {
126
    Log::Fatal("Binding port %d failed", port);
Guolin Ke's avatar
Guolin Ke committed
127
128
129
130
  }
}

void Linkers::SetLinker(int rank, const TcpSocket& socket) {
Guolin Ke's avatar
Guolin Ke committed
131
  linkers_[rank].reset(new TcpSocket(socket));
Guolin Ke's avatar
Guolin Ke committed
132
133
134
135
136
  // set timeout
  linkers_[rank]->SetTimeout(socket_timeout_ * 1000 * 60);
}

void Linkers::ListenThread(int incoming_cnt) {
137
  Log::Info("Listening...");
Guolin Ke's avatar
Guolin Ke committed
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
  char buffer[100];
  int connected_cnt = 0;
  while (connected_cnt < incoming_cnt) {
    // accept incoming socket
    TcpSocket handler = listener_->Accept();
    if (handler.IsClosed()) {
      continue;
    }
    // receive rank
    int read_cnt = 0;
    int size_of_int = static_cast<int>(sizeof(int));
    while (read_cnt < size_of_int) {
      int cur_read_cnt = handler.Recv(buffer + read_cnt, size_of_int - read_cnt);
      read_cnt += cur_read_cnt;
    }
    int* ptr_in_rank = reinterpret_cast<int*>(buffer);
    int in_rank = *ptr_in_rank;
    // add new socket
    SetLinker(in_rank, handler);
    ++connected_cnt;
  }
}

void Linkers::Construct() {
  // save ranks that need to connect with
  std::unordered_map<int, int> need_connect;
Guolin Ke's avatar
Guolin Ke committed
164
165
166
  for (int i = 0; i < num_machines_; ++i) {
    if (i != rank_) {
      need_connect[i] = 1;
Guolin Ke's avatar
Guolin Ke committed
167
168
169
170
171
172
173
174
175
176
177
178
179
    }
  }
  int need_connect_cnt = 0;
  int incoming_cnt = 0;
  for (auto it = need_connect.begin(); it != need_connect.end(); ++it) {
    int machine_rank = it->first;
    if (machine_rank >= 0 && machine_rank != rank_) {
      ++need_connect_cnt;
    }
    if (machine_rank < rank_) {
      ++incoming_cnt;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
180

Guolin Ke's avatar
Guolin Ke committed
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
  // start listener
  listener_->SetTimeout(socket_timeout_);
  listener_->Listen(incoming_cnt);
  std::thread listen_thread(&Linkers::ListenThread, this, incoming_cnt);
  const int connect_fail_retry_cnt = 20;
  const int connect_fail_delay_time = 10 * 1000;  // 10s
  // start connect
  for (auto it = need_connect.begin(); it != need_connect.end(); ++it) {
    int out_rank = it->first;
    // let smaller rank connect to larger rank
    if (out_rank > rank_) {
      TcpSocket cur_socket;
      for (int i = 0; i < connect_fail_retry_cnt; ++i) {
        if (cur_socket.Connect(client_ips_[out_rank].c_str(), client_ports_[out_rank])) {
          break;
        } else {
197
          Log::Warning("Connecting to rank %d failed, waiting for %d milliseconds", out_rank, connect_fail_delay_time);
Guolin Ke's avatar
Guolin Ke committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
          std::this_thread::sleep_for(std::chrono::milliseconds(connect_fail_delay_time));
        }
      }
      // send local rank
      cur_socket.Send(reinterpret_cast<const char*>(&rank_), sizeof(rank_));
      SetLinker(out_rank, cur_socket);
    }
  }
  // wait for listener
  listen_thread.join();
  // print connected linkers
  PrintLinkers();
}

bool Linkers::CheckLinker(int rank) {
  if (linkers_[rank] == nullptr || linkers_[rank]->IsClosed()) {
    return false;
  }
  return true;
}

void Linkers::PrintLinkers() {
  for (int i = 0; i < num_machines_; ++i) {
    if (CheckLinker(i)) {
222
      Log::Info("Connected to rank %d", i);
Guolin Ke's avatar
Guolin Ke committed
223
224
225
226
227
228
229
    }
  }
}

}  // namespace LightGBM

#endif  // USE_SOCKET