network.cc 30 KB
Newer Older
1
2
3
4
5
/*!
 *  Copyright (c) 2018 by Contributors
 * \file graph/network.cc
 * \brief DGL networking related APIs
 */
6
#include "./network.h"
7

Chao Ma's avatar
Chao Ma committed
8
9
#include <stdlib.h>

10
#include <dgl/runtime/container.h>
11
#include <dgl/runtime/ndarray.h>
12
#include <dgl/runtime/parallel_for.h>
13
#include <dgl/packed_func_ext.h>
14
15
16
17
18
#include <dgl/immutable_graph.h>
#include <dgl/nodeflow.h>

#include <unordered_map>

19
20
21
22
#include "../rpc/network/communicator.h"
#include "../rpc/network/socket_communicator.h"
#include "../rpc/network/msg_queue.h"
#include "../rpc/network/common.h"
23

24
using dgl::network::StringPrintf;
25
using namespace dgl::runtime;
26

Chao Ma's avatar
Chao Ma committed
27
28
const bool AUTO_FREE = true;

29
30
31
namespace dgl {
namespace network {

32
33
34
static void NaiveDeleter(DLManagedTensor* managed_tensor) {
  delete [] managed_tensor->dl_tensor.shape;
  delete [] managed_tensor->dl_tensor.strides;
Jinjing Zhou's avatar
Jinjing Zhou committed
35
  free(managed_tensor->dl_tensor.data);
36
37
38
  delete managed_tensor;
}

Chao Ma's avatar
Chao Ma committed
39
40
41
NDArray CreateNDArrayFromRaw(std::vector<int64_t> shape,
                             DLDataType dtype,
                             DLContext ctx,
Chao Ma's avatar
Chao Ma committed
42
43
                             void* raw,
                             bool auto_free) {
Chao Ma's avatar
Chao Ma committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
  DLTensor tensor;
  tensor.ctx = ctx;
  tensor.ndim = static_cast<int>(shape.size());
  tensor.dtype = dtype;
  tensor.shape = new int64_t[tensor.ndim];
  for (int i = 0; i < tensor.ndim; ++i) {
    tensor.shape[i] = shape[i];
  }
  tensor.strides = new int64_t[tensor.ndim];
  for (int i = 0; i < tensor.ndim; ++i) {
    tensor.strides[i] = 1;
  }
  for (int i = tensor.ndim - 2; i >= 0; --i) {
    tensor.strides[i] = tensor.shape[i+1] * tensor.strides[i+1];
  }
  tensor.data = raw;
  DLManagedTensor *managed_tensor = new DLManagedTensor();
  managed_tensor->dl_tensor = tensor;
Chao Ma's avatar
Chao Ma committed
62
63
64
  if (auto_free) {
    managed_tensor->deleter = NaiveDeleter;
  }
Chao Ma's avatar
Chao Ma committed
65
66
67
68
  return NDArray::FromDLPack(managed_tensor);
}

void ArrayMeta::AddArray(const NDArray& array) {
Chao Ma's avatar
Chao Ma committed
69
70
  // Get data type of current NDArray
  data_type_.push_back(array->dtype);
71
72
73
74
75
76
77
78
79
  // We first write the ndim to the data_shape_
  data_shape_.push_back(static_cast<int64_t>(array->ndim));
  // Then we write the data shape
  for (int i = 0; i < array->ndim; ++i) {
    data_shape_.push_back(array->shape[i]);
  }
  ndarray_count_++;
}

Chao Ma's avatar
Chao Ma committed
80
char* ArrayMeta::Serialize(int64_t* size) {
81
82
83
84
85
86
87
  char* buffer = nullptr;
  int64_t buffer_size = 0;
  buffer_size += sizeof(msg_type_);
  if (ndarray_count_ != 0) {
    buffer_size += sizeof(ndarray_count_);
    buffer_size += sizeof(data_shape_.size());
    buffer_size += sizeof(int64_t) * data_shape_.size();
Chao Ma's avatar
Chao Ma committed
88
89
90
    // we don't need to write data_type_.size()
    // because it equals to ndarray_count_ * 3
    buffer_size += sizeof(DLDataType) * data_type_.size();
91
  }
92
  // In the future, we should have a better memory management as
Chao Ma's avatar
Chao Ma committed
93
  // allocating a large chunk of memory can be very expensive.
94
95
96
97
98
99
100
101
102
  buffer = new char[buffer_size];
  char* pointer = buffer;
  // Write msg_type_
  *(reinterpret_cast<int*>(pointer)) = msg_type_;
  pointer += sizeof(msg_type_);
  if (ndarray_count_ != 0) {
    // Write ndarray_count_
    *(reinterpret_cast<int*>(pointer)) = ndarray_count_;
    pointer += sizeof(ndarray_count_);
Chao Ma's avatar
Chao Ma committed
103
104
105
106
107
    // Write data type
    memcpy(pointer,
        reinterpret_cast<DLDataType*>(data_type_.data()),
        sizeof(DLDataType) * data_type_.size());
    pointer += (sizeof(DLDataType) * data_type_.size());
108
109
110
111
112
113
114
    // Write size of data_shape_
    *(reinterpret_cast<size_t*>(pointer)) = data_shape_.size();
    pointer += sizeof(data_shape_.size());
    // Write data of data_shape_
    memcpy(pointer,
        reinterpret_cast<char*>(data_shape_.data()),
        sizeof(int64_t) * data_shape_.size());
115
  }
116
117
  *size = buffer_size;
  return buffer;
118
119
}

Chao Ma's avatar
Chao Ma committed
120
void ArrayMeta::Deserialize(char* buffer, int64_t size) {
121
122
123
124
125
126
127
128
129
130
  int64_t data_size = 0;
  // Read mesg_type_
  msg_type_ = *(reinterpret_cast<int*>(buffer));
  buffer += sizeof(int);
  data_size += sizeof(int);
  if (data_size < size) {
    // Read ndarray_count_
    ndarray_count_ = *(reinterpret_cast<int*>(buffer));
    buffer += sizeof(int);
    data_size += sizeof(int);
Chao Ma's avatar
Chao Ma committed
131
132
133
134
135
136
    // Read data type
    data_type_.resize(ndarray_count_);
    memcpy(data_type_.data(), buffer,
        ndarray_count_ * sizeof(DLDataType));
    buffer += ndarray_count_ * sizeof(DLDataType);
    data_size += ndarray_count_ * sizeof(DLDataType);
137
138
139
140
141
142
143
144
145
    // Read size of data_shape_
    size_t count = *(reinterpret_cast<size_t*>(buffer));
    buffer += sizeof(size_t);
    data_size += sizeof(size_t);
    data_shape_.resize(count);
    // Read data of data_shape_
    memcpy(data_shape_.data(), buffer,
        count * sizeof(int64_t));
    data_size += count * sizeof(int64_t);
146
  }
147
  CHECK_EQ(data_size, size);
148
149
}

Chao Ma's avatar
Chao Ma committed
150
151
152
153
154
155
156
157
158
char* KVStoreMsg::Serialize(int64_t* size) {
  char* buffer = nullptr;
  int64_t buffer_size = 0;
  buffer_size += sizeof(this->msg_type);
  buffer_size += sizeof(this->rank);
  if (!this->name.empty()) {
    buffer_size += sizeof(this->name.size());
    buffer_size += this->name.size();
  }
159
  // In the future, we should have a better memory management as
Chao Ma's avatar
Chao Ma committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
  // allocating a large chunk of memory can be very expensive.
  buffer = new char[buffer_size];
  char* pointer = buffer;
  // write msg_type
  *(reinterpret_cast<int*>(pointer)) = this->msg_type;
  pointer += sizeof(this->msg_type);
  // write rank
  *(reinterpret_cast<int*>(pointer)) = this->rank;
  pointer += sizeof(this->rank);
  // write name
  if (!this->name.empty()) {
    *(reinterpret_cast<size_t*>(pointer)) = this->name.size();
    pointer += sizeof(size_t);
    memcpy(pointer, this->name.c_str(), this->name.size());
  }
  *size = buffer_size;
  return buffer;
}

void KVStoreMsg::Deserialize(char* buffer, int64_t size) {
  int64_t data_size = 0;
  // Read msg_type
  this->msg_type = *(reinterpret_cast<int*>(buffer));
  buffer += sizeof(int);
  data_size += sizeof(int);
  // Read rank
  this->rank = *(reinterpret_cast<int*>(buffer));
  buffer += sizeof(int);
  data_size += sizeof(int);
  if (data_size < size) {
    // Read name
    size_t name_size = *(reinterpret_cast<size_t*>(buffer));
    buffer += sizeof(name_size);
    data_size += sizeof(name_size);
    this->name.assign(buffer, name_size);
    data_size += name_size;
  }
  CHECK_EQ(data_size, size);
}

200
201
202
////////////////////////////////// Basic Networking Components ////////////////////////////////


203
204
DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
205
    std::string type = args[0];
206
    int64_t msg_queue_size = args[1];
207
208
    network::Sender* sender = nullptr;
    if (type == "socket") {
209
      sender = new network::SocketSender(msg_queue_size);
210
211
    } else {
      LOG(FATAL) << "Unknown communicator type: " << type;
212
    }
213
    CommunicatorHandle chandle = static_cast<CommunicatorHandle>(sender);
214
215
216
    *rv = chandle;
  });

217
218
219
DGL_REGISTER_GLOBAL("network._CAPI_DGLReceiverCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    std::string type = args[0];
220
    int64_t msg_queue_size = args[1];
221
222
    network::Receiver* receiver = nullptr;
    if (type == "socket") {
223
      receiver = new network::SocketReceiver(msg_queue_size);
224
225
226
227
228
229
230
    } else {
      LOG(FATAL) << "Unknown communicator type: " << type;
    }
    CommunicatorHandle chandle = static_cast<CommunicatorHandle>(receiver);
    *rv = chandle;
  });

231
DGL_REGISTER_GLOBAL("network._CAPI_DGLFinalizeSender")
232
.set_body([] (DGLArgs args, DGLRetValue* rv) {
233
234
235
236
237
    CommunicatorHandle chandle = args[0];
    network::Sender* sender = static_cast<network::Sender*>(chandle);
    sender->Finalize();
  });

238
239
240
241
242
243
244
DGL_REGISTER_GLOBAL("network._CAPI_DGLFinalizeReceiver")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
    receiver->Finalize();
  });

245
246
247
248
249
250
251
DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderAddReceiver")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
    std::string ip = args[1];
    int port = args[2];
    int recv_id = args[3];
    network::Sender* sender = static_cast<network::Sender*>(chandle);
252
253
254
255
256
257
258
    std::string addr;
    if (sender->Type() == "socket") {
      addr = StringPrintf("socket://%s:%d", ip.c_str(), port);
    } else {
      LOG(FATAL) << "Unknown communicator type: " << sender->Type();
    }
    sender->AddReceiver(addr.c_str(), recv_id);
259
260
261
262
263
264
265
266
  });

DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderConnect")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
    network::Sender* sender = static_cast<network::Sender*>(chandle);
    if (sender->Connect() == false) {
      LOG(FATAL) << "Sender connection failed.";
267
268
269
    }
  });

270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
DGL_REGISTER_GLOBAL("network._CAPI_DGLReceiverWait")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
    std::string ip = args[1];
    int port = args[2];
    int num_sender = args[3];
    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
    std::string addr;
    if (receiver->Type() == "socket") {
      addr = StringPrintf("socket://%s:%d", ip.c_str(), port);
    } else {
      LOG(FATAL) << "Unknown communicator type: " << receiver->Type();
    }
    if (receiver->Wait(addr.c_str(), num_sender) == false) {
      LOG(FATAL) << "Wait sender socket failed.";
    }
  });


////////////////////////// Distributed Sampler Components ////////////////////////////////


DGL_REGISTER_GLOBAL("network._CAPI_SenderSendNodeFlow")
293
294
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
295
    int recv_id = args[1];
296
    GraphRef g = args[2];
297
298
299
300
    NDArray node_mapping = args[3];
    NDArray edge_mapping = args[4];
    NDArray layer_offsets = args[5];
    NDArray flow_offsets = args[6];
301
302
    auto ptr = std::dynamic_pointer_cast<ImmutableGraph>(g.sptr());
    CHECK(ptr) << "only immutable graph is allowed in send/recv";
303
    auto csr = ptr->GetInCSR();
304
305
306
307
    // Create a message for the meta data of ndarray
    NDArray indptr = csr->indptr();
    NDArray indice = csr->indices();
    NDArray edge_ids = csr->edge_ids();
Chao Ma's avatar
Chao Ma committed
308
309
310
311
312
313
314
315
    ArrayMeta meta(kNodeFlowMsg);
    meta.AddArray(node_mapping);
    meta.AddArray(edge_mapping);
    meta.AddArray(layer_offsets);
    meta.AddArray(flow_offsets);
    meta.AddArray(indptr);
    meta.AddArray(indice);
    meta.AddArray(edge_ids);
316
317
    // send meta message
    int64_t size = 0;
Chao Ma's avatar
Chao Ma committed
318
    char* data = meta.Serialize(&size);
319
320
321
322
323
    network::Sender* sender = static_cast<network::Sender*>(chandle);
    Message send_msg;
    send_msg.data = data;
    send_msg.size = size;
    send_msg.deallocator = DefaultMessageDeleter;
Chao Ma's avatar
Chao Ma committed
324
    CHECK_EQ(sender->Send(send_msg, recv_id), ADD_SUCCESS);
325
326
327
328
    // send node_mapping
    Message node_mapping_msg;
    node_mapping_msg.data = static_cast<char*>(node_mapping->data);
    node_mapping_msg.size = node_mapping.GetSize();
329
330
    // capture the array in the closure
    node_mapping_msg.deallocator = [node_mapping](Message*) {};
Chao Ma's avatar
Chao Ma committed
331
    CHECK_EQ(sender->Send(node_mapping_msg, recv_id), ADD_SUCCESS);
332
333
334
335
    // send edege_mapping
    Message edge_mapping_msg;
    edge_mapping_msg.data = static_cast<char*>(edge_mapping->data);
    edge_mapping_msg.size = edge_mapping.GetSize();
336
337
    // capture the array in the closure
    edge_mapping_msg.deallocator = [edge_mapping](Message*) {};
Chao Ma's avatar
Chao Ma committed
338
    CHECK_EQ(sender->Send(edge_mapping_msg, recv_id), ADD_SUCCESS);
339
340
341
342
    // send layer_offsets
    Message layer_offsets_msg;
    layer_offsets_msg.data = static_cast<char*>(layer_offsets->data);
    layer_offsets_msg.size = layer_offsets.GetSize();
343
344
    // capture the array in the closure
    layer_offsets_msg.deallocator = [layer_offsets](Message*) {};
Chao Ma's avatar
Chao Ma committed
345
    CHECK_EQ(sender->Send(layer_offsets_msg, recv_id), ADD_SUCCESS);
346
347
348
349
    // send flow_offset
    Message flow_offsets_msg;
    flow_offsets_msg.data = static_cast<char*>(flow_offsets->data);
    flow_offsets_msg.size = flow_offsets.GetSize();
350
351
    // capture the array in the closure
    flow_offsets_msg.deallocator = [flow_offsets](Message*) {};
Chao Ma's avatar
Chao Ma committed
352
    CHECK_EQ(sender->Send(flow_offsets_msg, recv_id), ADD_SUCCESS);
353
354
355
356
    // send csr->indptr
    Message indptr_msg;
    indptr_msg.data = static_cast<char*>(indptr->data);
    indptr_msg.size = indptr.GetSize();
357
358
    // capture the array in the closure
    indptr_msg.deallocator = [indptr](Message*) {};
Chao Ma's avatar
Chao Ma committed
359
    CHECK_EQ(sender->Send(indptr_msg, recv_id), ADD_SUCCESS);
360
361
362
363
    // send csr->indices
    Message indices_msg;
    indices_msg.data = static_cast<char*>(indice->data);
    indices_msg.size = indice.GetSize();
364
365
    // capture the array in the closure
    indices_msg.deallocator = [indice](Message*) {};
Chao Ma's avatar
Chao Ma committed
366
    CHECK_EQ(sender->Send(indices_msg, recv_id), ADD_SUCCESS);
367
368
    // send csr->edge_ids
    Message edge_ids_msg;
369
370
371
372
    edge_ids_msg.data = static_cast<char*>(edge_ids->data);
    edge_ids_msg.size = edge_ids.GetSize();
    // capture the array in the closure
    edge_ids_msg.deallocator = [edge_ids](Message*) {};
Chao Ma's avatar
Chao Ma committed
373
    CHECK_EQ(sender->Send(edge_ids_msg, recv_id), ADD_SUCCESS);
374
375
  });

376
DGL_REGISTER_GLOBAL("network._CAPI_SenderSendSamplerEndSignal")
377
378
379
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
    int recv_id = args[1];
Chao Ma's avatar
Chao Ma committed
380
    ArrayMeta meta(kFinalMsg);
381
    int64_t size = 0;
Chao Ma's avatar
Chao Ma committed
382
    char* data = meta.Serialize(&size);
383
    network::Sender* sender = static_cast<network::Sender*>(chandle);
384
385
    Message send_msg = {data, size};
    send_msg.deallocator = DefaultMessageDeleter;
Chao Ma's avatar
Chao Ma committed
386
    CHECK_EQ(sender->Send(send_msg, recv_id), ADD_SUCCESS);
387
388
  });

389
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
390
391
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
392
    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
393
394
    int send_id = 0;
    Message recv_msg;
Chao Ma's avatar
Chao Ma committed
395
396
    CHECK_EQ(receiver->Recv(&recv_msg, &send_id), REMOVE_SUCCESS);
    ArrayMeta meta(recv_msg.data, recv_msg.size);
397
    recv_msg.deallocator(&recv_msg);
Chao Ma's avatar
Chao Ma committed
398
399
    if (meta.msg_type() == kNodeFlowMsg) {
      CHECK_EQ(meta.ndarray_count() * 2, meta.data_shape_.size());
400
      NodeFlow nf = NodeFlow::Create();
401
402
      // node_mapping
      Message array_0;
Chao Ma's avatar
Chao Ma committed
403
404
405
406
407
408
      CHECK_EQ(receiver->RecvFrom(&array_0, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[0], 1);
      nf->node_mapping = CreateNDArrayFromRaw(
        {meta.data_shape_[1]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
409
410
        array_0.data,
        AUTO_FREE);
411
412
      // edge_mapping
      Message array_1;
Chao Ma's avatar
Chao Ma committed
413
414
415
416
417
418
      CHECK_EQ(receiver->RecvFrom(&array_1, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[2], 1);
      nf->edge_mapping = CreateNDArrayFromRaw(
        {meta.data_shape_[3]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
419
420
        array_1.data,
        AUTO_FREE);
421
422
      // layer_offset
      Message array_2;
Chao Ma's avatar
Chao Ma committed
423
424
425
426
427
428
      CHECK_EQ(receiver->RecvFrom(&array_2, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[4], 1);
      nf->layer_offsets = CreateNDArrayFromRaw(
        {meta.data_shape_[5]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
429
430
        array_2.data,
        AUTO_FREE);
431
432
      // flow_offset
      Message array_3;
Chao Ma's avatar
Chao Ma committed
433
434
435
436
437
438
      CHECK_EQ(receiver->RecvFrom(&array_3, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[6], 1);
      nf->flow_offsets = CreateNDArrayFromRaw(
        {meta.data_shape_[7]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
439
440
        array_3.data,
        AUTO_FREE);
441
442
      // CSR indptr
      Message array_4;
Chao Ma's avatar
Chao Ma committed
443
444
445
446
447
448
      CHECK_EQ(receiver->RecvFrom(&array_4, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[8], 1);
      NDArray indptr = CreateNDArrayFromRaw(
        {meta.data_shape_[9]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
449
450
        array_4.data,
        AUTO_FREE);
451
452
      // CSR indice
      Message array_5;
Chao Ma's avatar
Chao Ma committed
453
454
455
456
457
458
      CHECK_EQ(receiver->RecvFrom(&array_5, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[10], 1);
      NDArray indice = CreateNDArrayFromRaw(
        {meta.data_shape_[11]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
459
460
        array_5.data,
        AUTO_FREE);
461
462
      // CSR edge_ids
      Message array_6;
Chao Ma's avatar
Chao Ma committed
463
464
465
466
467
468
      CHECK_EQ(receiver->RecvFrom(&array_6, send_id), REMOVE_SUCCESS);
      CHECK_EQ(meta.data_shape_[12], 1);
      NDArray edge_ids = CreateNDArrayFromRaw(
        {meta.data_shape_[13]},
        DLDataType{kDLInt, 64, 1},
        DLContext{kDLCPU, 0},
Chao Ma's avatar
Chao Ma committed
469
470
        array_6.data,
        AUTO_FREE);
471
472
      // Create CSR
      CSRPtr csr(new CSR(indptr, indice, edge_ids));
473
      nf->graph = GraphPtr(new ImmutableGraph(csr, nullptr));
474
      *rv = nf;
Chao Ma's avatar
Chao Ma committed
475
476
    } else if (meta.msg_type() == kFinalMsg) {
      *rv = meta.msg_type();
477
    } else {
Chao Ma's avatar
Chao Ma committed
478
479
480
481
482
483
484
485
      LOG(FATAL) << "Unknown message type: " << meta.msg_type();
    }
  });


////////////////////////// Distributed KVStore Components ////////////////////////////////


Chao Ma's avatar
Chao Ma committed
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
static void send_kv_message(network::Sender* sender,
                            KVStoreMsg* kv_msg,
                            int recv_id,
                            bool auto_free) {
  int64_t kv_size = 0;
  char* kv_data = kv_msg->Serialize(&kv_size);
  // Send kv_data
  Message send_kv_msg;
  send_kv_msg.data = kv_data;
  send_kv_msg.size = kv_size;
  if (auto_free) {
    send_kv_msg.deallocator = DefaultMessageDeleter;
  }
  CHECK_EQ(sender->Send(send_kv_msg, recv_id), ADD_SUCCESS);
  if (kv_msg->msg_type != kFinalMsg &&
      kv_msg->msg_type != kBarrierMsg &&
502
503
      kv_msg->msg_type != kIPIDMsg &&
      kv_msg->msg_type != kGetShapeMsg) {
Chao Ma's avatar
Chao Ma committed
504
505
    // Send ArrayMeta
    ArrayMeta meta(kv_msg->msg_type);
506
507
    if (kv_msg->msg_type != kInitMsg &&
        kv_msg->msg_type != kGetShapeBackMsg) {
508
509
510
      meta.AddArray(kv_msg->id);
    }
    if (kv_msg->msg_type != kPullMsg &&
511
512
        kv_msg->msg_type != kInitMsg &&
        kv_msg->msg_type != kGetShapeBackMsg) {
Chao Ma's avatar
Chao Ma committed
513
514
      meta.AddArray(kv_msg->data);
    }
515
516
517
518
519
    if (kv_msg->msg_type != kPullMsg &&
        kv_msg->msg_type != kPushMsg &&
        kv_msg->msg_type != kPullBackMsg) {
      meta.AddArray(kv_msg->shape);
    }
Chao Ma's avatar
Chao Ma committed
520
521
522
523
524
525
526
527
528
529
    int64_t meta_size = 0;
    char* meta_data = meta.Serialize(&meta_size);
    Message send_meta_msg;
    send_meta_msg.data = meta_data;
    send_meta_msg.size = meta_size;
    if (auto_free) {
      send_meta_msg.deallocator = DefaultMessageDeleter;
    }
    CHECK_EQ(sender->Send(send_meta_msg, recv_id), ADD_SUCCESS);
    // Send ID NDArray
530
531
532
    if (kv_msg->msg_type != kInitMsg &&
        kv_msg->msg_type != kGetShapeMsg &&
        kv_msg->msg_type != kGetShapeBackMsg) {
533
534
535
536
537
538
539
540
541
      Message send_id_msg;
      send_id_msg.data = static_cast<char*>(kv_msg->id->data);
      send_id_msg.size = kv_msg->id.GetSize();
      NDArray id = kv_msg->id;
      if (auto_free) {
        send_id_msg.deallocator = [id](Message*) {};
      }
      CHECK_EQ(sender->Send(send_id_msg, recv_id), ADD_SUCCESS);
    }
Chao Ma's avatar
Chao Ma committed
542
    // Send data NDArray
543
    if (kv_msg->msg_type != kPullMsg &&
544
545
546
        kv_msg->msg_type != kInitMsg &&
        kv_msg->msg_type != kGetShapeMsg &&
        kv_msg->msg_type != kGetShapeBackMsg) {
Chao Ma's avatar
Chao Ma committed
547
548
549
550
551
552
553
554
555
      Message send_data_msg;
      send_data_msg.data = static_cast<char*>(kv_msg->data->data);
      send_data_msg.size = kv_msg->data.GetSize();
      NDArray data = kv_msg->data;
      if (auto_free) {
        send_data_msg.deallocator = [data](Message*) {};
      }
      CHECK_EQ(sender->Send(send_data_msg, recv_id), ADD_SUCCESS);
    }
556
557
558
559
560
561
562
563
564
565
566
567
568
    // Send shape NDArray
    if (kv_msg->msg_type != kPullMsg &&
        kv_msg->msg_type != kPushMsg &&
        kv_msg->msg_type != kPullBackMsg) {
      Message send_shape_msg;
      send_shape_msg.data = static_cast<char*>(kv_msg->shape->data);
      send_shape_msg.size = kv_msg->shape.GetSize();
      NDArray shape = kv_msg->shape;
      if (auto_free) {
        send_shape_msg.deallocator = [shape](Message*) {};
      }
      CHECK_EQ(sender->Send(send_shape_msg, recv_id), ADD_SUCCESS);
    }
Chao Ma's avatar
Chao Ma committed
569
570
571
572
573
574
575
576
577
578
579
580
581
  }
}

static KVStoreMsg* recv_kv_message(network::Receiver* receiver) {
  KVStoreMsg *kv_msg = new KVStoreMsg();
  // Recv kv_Msg
  Message recv_kv_msg;
  int send_id;
  CHECK_EQ(receiver->Recv(&recv_kv_msg, &send_id), REMOVE_SUCCESS);
  kv_msg->Deserialize(recv_kv_msg.data, recv_kv_msg.size);
  recv_kv_msg.deallocator(&recv_kv_msg);
  if (kv_msg->msg_type == kFinalMsg ||
      kv_msg->msg_type == kBarrierMsg ||
582
583
      kv_msg->msg_type == kIPIDMsg ||
      kv_msg->msg_type == kGetShapeMsg) {
Chao Ma's avatar
Chao Ma committed
584
585
586
587
588
589
590
591
    return kv_msg;
  }
  // Recv ArrayMeta
  Message recv_meta_msg;
  CHECK_EQ(receiver->RecvFrom(&recv_meta_msg, send_id), REMOVE_SUCCESS);
  ArrayMeta meta(recv_meta_msg.data, recv_meta_msg.size);
  recv_meta_msg.deallocator(&recv_meta_msg);
  // Recv ID NDArray
592
593
  if (kv_msg->msg_type != kInitMsg &&
      kv_msg->msg_type != kGetShapeBackMsg) {
594
595
596
597
598
    Message recv_id_msg;
    CHECK_EQ(receiver->RecvFrom(&recv_id_msg, send_id), REMOVE_SUCCESS);
    CHECK_EQ(meta.data_shape_[0], 1);
    kv_msg->id = CreateNDArrayFromRaw(
      {meta.data_shape_[1]},
Chao Ma's avatar
Chao Ma committed
599
      meta.data_type_[0],
600
601
602
603
      DLContext{kDLCPU, 0},
      recv_id_msg.data,
      AUTO_FREE);
  }
Chao Ma's avatar
Chao Ma committed
604
  // Recv Data NDArray
605
  if (kv_msg->msg_type != kPullMsg &&
606
607
      kv_msg->msg_type != kInitMsg &&
      kv_msg->msg_type != kGetShapeBackMsg) {
Chao Ma's avatar
Chao Ma committed
608
609
    Message recv_data_msg;
    CHECK_EQ(receiver->RecvFrom(&recv_data_msg, send_id), REMOVE_SUCCESS);
610
611
    int64_t ndim = meta.data_shape_[2];
    CHECK_GE(ndim, 1);
Chao Ma's avatar
Chao Ma committed
612
    std::vector<int64_t> vec_shape;
613
614
    for (int i = 0; i < ndim; ++i) {
      vec_shape.push_back(meta.data_shape_[3+i]);
Chao Ma's avatar
Chao Ma committed
615
616
617
    }
    kv_msg->data = CreateNDArrayFromRaw(
      vec_shape,
Chao Ma's avatar
Chao Ma committed
618
      meta.data_type_[1],
Chao Ma's avatar
Chao Ma committed
619
620
621
622
      DLContext{kDLCPU, 0},
      recv_data_msg.data,
      AUTO_FREE);
  }
623
624
625
626
627
628
629
630
631
632
633
634
635
636
  // Recv Shape
  if (kv_msg->msg_type != kPullMsg &&
      kv_msg->msg_type != kPushMsg &&
      kv_msg->msg_type != kPullBackMsg) {
    Message recv_shape_msg;
    CHECK_EQ(receiver->RecvFrom(&recv_shape_msg, send_id), REMOVE_SUCCESS);
    int64_t ndim = meta.data_shape_[0];
    CHECK_GE(ndim, 1);
    std::vector<int64_t> vec_shape;
    for (int i = 0; i < ndim; ++i) {
      vec_shape.push_back(meta.data_shape_[1+i]);
    }
    kv_msg->shape = CreateNDArrayFromRaw(
      vec_shape,
Chao Ma's avatar
Chao Ma committed
637
      meta.data_type_[0],
638
639
640
641
      DLContext{kDLCPU, 0},
      recv_shape_msg.data,
      AUTO_FREE);
  }
Chao Ma's avatar
Chao Ma committed
642
643
644
  return kv_msg;
}

Chao Ma's avatar
Chao Ma committed
645
646
DGL_REGISTER_GLOBAL("network._CAPI_SenderSendKVMsg")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
647
648
649
    int args_count = 0;
    CommunicatorHandle chandle = args[args_count++];
    int recv_id = args[args_count++];
Chao Ma's avatar
Chao Ma committed
650
    KVStoreMsg kv_msg;
651
652
    kv_msg.msg_type = args[args_count++];
    kv_msg.rank = args[args_count++];
Chao Ma's avatar
Chao Ma committed
653
654
    network::Sender* sender = static_cast<network::Sender*>(chandle);
    if (kv_msg.msg_type != kFinalMsg && kv_msg.msg_type != kBarrierMsg) {
655
      std::string name = args[args_count++];
Chao Ma's avatar
Chao Ma committed
656
      kv_msg.name = name;
657
      if (kv_msg.msg_type != kIPIDMsg &&
658
659
660
          kv_msg.msg_type != kInitMsg &&
          kv_msg.msg_type != kGetShapeMsg &&
          kv_msg.msg_type != kGetShapeBackMsg) {
661
662
        kv_msg.id = args[args_count++];
      }
663
664
      if (kv_msg.msg_type != kPullMsg &&
          kv_msg.msg_type != kIPIDMsg &&
665
666
667
          kv_msg.msg_type != kInitMsg &&
          kv_msg.msg_type != kGetShapeMsg &&
          kv_msg.msg_type != kGetShapeBackMsg) {
668
        kv_msg.data = args[args_count++];
Chao Ma's avatar
Chao Ma committed
669
      }
670
671
672
      if (kv_msg.msg_type != kIPIDMsg &&
          kv_msg.msg_type != kPullMsg &&
          kv_msg.msg_type != kPushMsg &&
673
674
          kv_msg.msg_type != kPullBackMsg &&
          kv_msg.msg_type != kGetShapeMsg) {
675
676
        kv_msg.shape = args[args_count++];
      }
Chao Ma's avatar
Chao Ma committed
677
    }
Chao Ma's avatar
Chao Ma committed
678
    send_kv_message(sender, &kv_msg, recv_id, AUTO_FREE);
Chao Ma's avatar
Chao Ma committed
679
680
681
682
683
684
  });

DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CommunicatorHandle chandle = args[0];
    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
Chao Ma's avatar
Chao Ma committed
685
    *rv = recv_kv_message(receiver);
Chao Ma's avatar
Chao Ma committed
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
  });

DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgType")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    *rv = msg->msg_type;
  });

DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgRank")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    *rv = msg->rank;
  });

DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgName")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    *rv = msg->name;
  });

DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    *rv = msg->id;
  });

DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgData")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    *rv = msg->data;
721
722
  });

723
724
725
726
727
728
729
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgShape")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    *rv = msg->shape;
  });

730
DGL_REGISTER_GLOBAL("network._CAPI_DeleteKVMsg")
731
.set_body([] (DGLArgs args, DGLRetValue* rv) {
732
733
734
    KVMsgHandle chandle = args[0];
    network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
    delete msg;
735
736
  });

Chao Ma's avatar
Chao Ma committed
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
DGL_REGISTER_GLOBAL("network._CAPI_FastPull")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    std::string name = args[0];
    int local_machine_id = args[1];
    int machine_count = args[2];
    int group_count = args[3];
    int client_id = args[4];
    NDArray ID = args[5];
    NDArray pb = args[6];
    NDArray local_data = args[7];
    CommunicatorHandle chandle_sender = args[8];
    CommunicatorHandle chandle_receiver = args[9];
    std::string str_flag = args[10];
    network::Sender* sender = static_cast<network::Sender*>(chandle_sender);
    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle_receiver);
    int64_t ID_size = ID.GetSize() / sizeof(int64_t);
    int64_t* ID_data = static_cast<int64_t*>(ID->data);
    int64_t* pb_data = static_cast<int64_t*>(pb->data);
    char* local_data_char = static_cast<char*>(local_data->data);
    std::vector<int64_t> local_ids;
    std::vector<int64_t> local_ids_orginal;
    std::vector<int64_t> local_data_shape;
    std::vector<std::vector<int64_t> > remote_ids(machine_count);
    std::vector<std::vector<int64_t> > remote_ids_original(machine_count);
    unsigned int seed = 314;
    int row_size = 1;
    for (int i = 0; i < local_data->ndim; ++i) {
      local_data_shape.push_back(local_data->shape[i]);
      if (i != 0) {
        row_size *= local_data->shape[i];
      }
    }
769
    row_size *= (local_data->dtype.bits / 8);
770
771
772
    size_t data_size = local_data.GetSize();
    CHECK_GT(local_data_shape.size(), 0);
    CHECK_EQ(row_size * local_data_shape[0], data_size);
Chao Ma's avatar
Chao Ma committed
773
774
775
776
777
778
779
780
781
    // Get local id and remote id
    if (str_flag.compare("has_g2l") == 0) {
      NDArray g2l = args[11];
      int64_t* g2l_data = static_cast<int64_t*>(g2l->data);
      for (int64_t i = 0; i < ID_size; ++i) {
        int64_t id = ID_data[i];
        int64_t part_id = pb_data[id];
        if (part_id == local_machine_id) {
          int64_t local_id = g2l_data[id];
782
783
          CHECK_LT(local_id, local_data_shape[0]);
          CHECK_GE(local_id, 0);
Chao Ma's avatar
Chao Ma committed
784
785
786
          local_ids.push_back(local_id);
          local_ids_orginal.push_back(i);
        } else {
787
          CHECK_LT(part_id, machine_count) << "invalid partition ID";
Chao Ma's avatar
Chao Ma committed
788
789
790
791
792
793
794
795
796
          remote_ids[part_id].push_back(id);
          remote_ids_original[part_id].push_back(i);
        }
      }
    } else {
      for (int64_t i = 0; i < ID_size; ++i) {
        int64_t id = ID_data[i];
        int64_t part_id = pb_data[id];
        if (part_id == local_machine_id) {
797
798
          CHECK_LT(id, local_data_shape[0]);
          CHECK_GE(id, 0);
Chao Ma's avatar
Chao Ma committed
799
800
801
802
803
804
805
806
807
          local_ids.push_back(id);
          local_ids_orginal.push_back(i);
        } else {
          remote_ids[part_id].push_back(id);
          remote_ids_original[part_id].push_back(i);
        }
      }
    }
    int msg_count = 0;
808
    for (size_t i = 0; i < remote_ids.size(); ++i) {
Chao Ma's avatar
Chao Ma committed
809
810
811
812
813
814
      if (remote_ids[i].size() != 0) {
        KVStoreMsg kv_msg;
        kv_msg.msg_type = MessageType::kPullMsg;
        kv_msg.rank = client_id;
        kv_msg.name = name;
        kv_msg.id = CreateNDArrayFromRaw({static_cast<int64_t>(remote_ids[i].size())},
Chao Ma's avatar
Chao Ma committed
815
                                         ID->dtype,
Chao Ma's avatar
Chao Ma committed
816
817
818
819
820
821
822
823
                                         DLContext{kDLCPU, 0},
                                         remote_ids[i].data(),
                                         !AUTO_FREE);
        int lower = i*group_count;
        int higher = (i+1)*group_count-1;
#ifndef _WIN32  // windows does not support rand_r()
        int s_id = (rand_r(&seed) % (higher-lower+1))+lower;
        send_kv_message(sender, &kv_msg, s_id, !AUTO_FREE);
824
825
#else
        LOG(FATAL) << "KVStore does not support Windows yet.";
Chao Ma's avatar
Chao Ma committed
826
827
828
829
830
#endif
        msg_count++;
      }
    }
    char *return_data = new char[ID_size*row_size];
831
    const int64_t local_ids_size = local_ids.size();
Chao Ma's avatar
Chao Ma committed
832
    // Copy local data
833
834
835
836
837
838
839
840
841
842
    runtime::parallel_for(0, local_ids_size, [&](size_t b, size_t e) {
      for (auto i = b; i < e; ++i) {
        CHECK_GE(ID_size*row_size, local_ids_orginal[i] * row_size + row_size);
        CHECK_GE(data_size, local_ids[i] * row_size + row_size);
        CHECK_GE(local_ids[i], 0);
        memcpy(return_data + local_ids_orginal[i] * row_size,
               local_data_char + local_ids[i] * row_size,
               row_size);
      }
    });
Chao Ma's avatar
Chao Ma committed
843
844
845
846
847
848
    // Recv remote message
    for (int i = 0; i < msg_count; ++i) {
      KVStoreMsg *kv_msg = recv_kv_message(receiver);
      int64_t id_size = kv_msg->id.GetSize() / sizeof(int64_t);
      int part_id = kv_msg->rank / group_count;
      char* data_char = static_cast<char*>(kv_msg->data->data);
849
      for (int64_t n = 0; n < id_size; ++n) {
Chao Ma's avatar
Chao Ma committed
850
851
852
853
854
855
856
857
858
859
        memcpy(return_data + remote_ids_original[part_id][n] * row_size,
               data_char + n * row_size,
               row_size);
      }
      delete kv_msg;
    }
    // Get final tensor
    local_data_shape[0] = ID_size;
    NDArray res_tensor = CreateNDArrayFromRaw(
                          local_data_shape,
Chao Ma's avatar
Chao Ma committed
860
                          local_data->dtype,
Chao Ma's avatar
Chao Ma committed
861
862
863
864
865
866
                          DLContext{kDLCPU, 0},
                          return_data,
                          AUTO_FREE);
    *rv = res_tensor;
  });

867
868
}  // namespace network
}  // namespace dgl