[KVstore] Fast-pull (#1446)

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* fix lint

python/dgl/contrib/dis_kvstore.py

@@ -5,6 +5,7 @@ from ..network import _network_wait, _add_receiver_addr
 from ..network import _receiver_wait, _sender_connect
 from ..network import _send_kv_msg, _recv_kv_msg
 from ..network import _clear_kv_msg
+from ..network import _fast_pull
 from ..network import KVMsgType, KVStoreMsg
 
 from .. import backend as F
@@ -146,6 +147,11 @@ class KVServer(object):
         self._open_file_list = []
         # record for total message count
         self._msg_count = 0
+        # user-defined push handler
+        self._udf_push_handler = None
+        self._udf_push_param = None
+        # user-defined pull handler
+        self._udf_pull_handler = None
 
 
     def __del__(self):
@@ -317,6 +323,8 @@ class KVServer(object):
         # Get connected with all client nodes
         _receiver_wait(self._receiver, self._ip, self._port, self._client_count)
 
+        print("%d clients connected!" % self._client_count)
+
         # recv client address information
         addr_list = []
         for i in range(self._client_count):
@@ -378,14 +386,20 @@ class KVServer(object):
                     local_id = self._data_store[msg.name+'-g2l-'][msg.id]
                 else:
                     local_id = msg.id
-                self._push_handler(msg.name+'-data-', local_id, msg.data, self._data_store)
+                if self._udf_push_handler is not None:
+                    self._udf_push_handler(msg.name+'-data-', local_id, msg.data, self._data_store, self._udf_push_param)
+                else:
+                    self._default_push_handler(msg.name+'-data-', local_id, msg.data, self._data_store)
             # Pull message
             elif msg.type == KVMsgType.PULL:
                 if (msg.name+'-g2l-' in self._has_data) == True:
                     local_id = self._data_store[msg.name+'-g2l-'][msg.id]
                 else:
                     local_id = msg.id
-                res_tensor = self._pull_handler(msg.name+'-data-', local_id, self._data_store)
+                if self._udf_pull_handler is not None:
+                    res_tensor = self._udf_pull_handler(msg.name+'-data-', local_id, self._data_store)
+                else:
+                    res_tensor = self._default_pull_handler(msg.name+'-data-', local_id, self._data_store)
                 back_msg = KVStoreMsg(
                     type=KVMsgType.PULL_BACK,
                     rank=self._server_id,
@@ -500,7 +514,7 @@ class KVServer(object):
         return data_shape
 
 
-    def _push_handler(self, name, ID, data, target):
+    def _default_push_handler(self, name, ID, data, target):
         """Default handler for PUSH message. 
 
         On default, _push_handler perform update operation for the tensor.
@@ -519,7 +533,7 @@ class KVServer(object):
         target[name][ID] = data
 
 
-    def _pull_handler(self, name, ID, target):
+    def _default_pull_handler(self, name, ID, target):
         """Default handler for PULL operation.
 
         On default, _pull_handler perform get operation for the tensor.
@@ -582,6 +596,7 @@ class KVClient(object):
         self._server_namebook = server_namebook
         self._server_count = len(server_namebook)
         self._group_count = server_namebook[0][3]
+        self._machine_count = int(self._server_count / self._group_count)
         # client ID will be assign by server after connecting to server
         self._client_id = -1
         # Get local machine id via server_namebook
@@ -593,6 +608,11 @@ class KVClient(object):
         self._open_file_list = []
         # Gargage_collection
         self._garbage_msg = []
+        # User-defined pull handler
+        self._udf_pull_handler = None
+        # User-defined push handler
+        self._udf_push_handler = None
+        self._udf_push_param = None
         # Used load-balance
         random.seed(time.time())
 
@@ -812,7 +832,10 @@ class KVClient(object):
             start += count[idx]
 
         if local_id is not None: # local push
-            self._push_handler(name+'-data-', local_id, local_data, self._data_store)
+            if self._udf_push_handler is not None:
+                self._udf_push_handler(name+'-data-', local_id, local_data, self._data_store, self._udf_push_param)
+            else:
+                self._default_push_handler(name+'-data-', local_id, local_data, self._data_store)
     
 
     def pull(self, name, id_tensor):
@@ -833,6 +856,21 @@ class KVClient(object):
         assert len(name) > 0, 'name cannot be empty.'
         assert F.ndim(id_tensor) == 1, 'ID must be a vector.'
 
+        if self._udf_pull_handler is None: # Use fast-pull
+            g2l = None
+            if name+'-g2l-' in self._data_store:
+               g2l = self._data_store[name+'-g2l-']
+            return _fast_pull(name, id_tensor,
+                        self._machine_count,
+                        self._group_count,
+                        self._machine_id,
+                        self._client_id,
+                        self._data_store[name+'-part-'], 
+                        g2l, 
+                        self._data_store[name+'-data-'],
+                        self._sender,
+                        self._receiver)
+        else:
             for msg in self._garbage_msg:
                 _clear_kv_msg(msg)
             self._garbage_msg = []
@@ -877,7 +915,7 @@ class KVClient(object):
 
             msg_list = []
             if local_id is not None: # local pull
-            local_data = self._pull_handler(name+'-data-', local_id, self._data_store)
+                local_data = self._udf_pull_handler(name+'-data-', local_id, self._data_store)
                 s_id = random.randint(self._machine_id*self._group_count, (self._machine_id+1)*self._group_count-1)
                 local_msg = KVStoreMsg(
                     type=KVMsgType.PULL_BACK, 
@@ -1082,7 +1120,7 @@ class KVClient(object):
         return elem.rank
 
 
-    def _push_handler(self, name, ID, data, target):
+    def _default_push_handler(self, name, ID, data, target):
         """Default handler for PUSH message. 
 
         On default, _push_handler perform update operation for the tensor.
@@ -1100,25 +1138,3 @@ class KVClient(object):
         """
         target[name][ID] = data
     
\ No newline at end of file
-
-    def _pull_handler(self, name, ID, target):
-        """Default handler for PULL operation.
-
-        On default, _pull_handler perform get operation for the tensor.
-
-        Parameters
-        ----------
-        name : str
-            data name
-        ID : tensor (mx.ndarray or torch.tensor)
-            a vector storing the ID list.
-        target : dict of data
-            self._data_store
-
-        Return
-        ------
-        tensor
-            a tensor with the same row size of ID.
-        """
-        return target[name][ID]
-    
\ No newline at end of file

python/dgl/network.py

@@ -327,3 +327,56 @@ def _clear_kv_msg(msg):
     F.sync()
     if msg.c_ptr is not None:
         _CAPI_DeleteKVMsg(msg.c_ptr)
+
+
+def _fast_pull(name, id_tensor,
+               machine_count, group_count, machine_id, client_id,
+               partition_book, g2l, local_data,
+               sender, receiver):
+    """ Pull message
+
+    Parameters
+    ----------
+    name : str
+        data name string
+    id_tensor : tensor
+        tensor of ID
+    machine_count : int
+        count of total machine
+    group_count : int
+        count of server group
+    machine_id : int
+        current machine id
+    client_id : int
+        current client ID
+    partition_book : tensor
+        tensor of partition book
+    g2l : tensor
+        tensor of global2local
+    local_data : tensor
+        tensor of local shared data
+    sender : ctypes.c_void_p
+        C Sender handle
+    receiver : ctypes.c_void_p
+        C Receiver handle
+
+    Return
+    ------
+    tensor
+        target tensor
+    """
+    if g2l is not None:
+        res_tensor = _CAPI_FastPull(name, machine_id, machine_count, group_count, client_id,
+                                    F.zerocopy_to_dgl_ndarray(id_tensor),
+                                    F.zerocopy_to_dgl_ndarray(partition_book),
+                                    F.zerocopy_to_dgl_ndarray(local_data),
+                                    sender, receiver, 'has_g2l',
+                                    F.zerocopy_to_dgl_ndarray(g2l))
+    else:
+        res_tensor = _CAPI_FastPull(name, machine_id, machine_count, group_count, client_id,
+                                    F.zerocopy_to_dgl_ndarray(id_tensor),
+                                    F.zerocopy_to_dgl_ndarray(partition_book),
+                                    F.zerocopy_to_dgl_ndarray(local_data),
+                                    sender, receiver, 'no_g2l')
+
+    return F.zerocopy_from_dgl_ndarray(res_tensor)

src/graph/network.cc

@@ -5,6 +5,8 @@
  */
 #include "./network.h"
 
+#include <stdlib.h>
+
 #include <dgl/runtime/container.h>
 #include <dgl/runtime/ndarray.h>
 #include <dgl/packed_func_ext.h>
@@ -21,6 +23,8 @@
 using dgl::network::StringPrintf;
 using namespace dgl::runtime;
 
+const bool AUTO_FREE = true;
+
 namespace dgl {
 namespace network {
 
@@ -34,7 +38,8 @@ static void NaiveDeleter(DLManagedTensor* managed_tensor) {
 NDArray CreateNDArrayFromRaw(std::vector<int64_t> shape,
                              DLDataType dtype,
                              DLContext ctx,
-                             void* raw) {
+                             void* raw,
+                             bool auto_free) {
   DLTensor tensor;
   tensor.ctx = ctx;
   tensor.ndim = static_cast<int>(shape.size());
@@ -53,7 +58,9 @@ NDArray CreateNDArrayFromRaw(std::vector<int64_t> shape,
   tensor.data = raw;
   DLManagedTensor *managed_tensor = new DLManagedTensor();
   managed_tensor->dl_tensor = tensor;
+  if (auto_free) {
     managed_tensor->deleter = NaiveDeleter;
+  }
   return NDArray::FromDLPack(managed_tensor);
 }
 
@@ -382,7 +389,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[1]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_0.data);
+        array_0.data,
+        AUTO_FREE);
       // edge_mapping
       Message array_1;
       CHECK_EQ(receiver->RecvFrom(&array_1, send_id), REMOVE_SUCCESS);
@@ -391,7 +399,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[3]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_1.data);
+        array_1.data,
+        AUTO_FREE);
       // layer_offset
       Message array_2;
       CHECK_EQ(receiver->RecvFrom(&array_2, send_id), REMOVE_SUCCESS);
@@ -400,7 +409,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[5]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_2.data);
+        array_2.data,
+        AUTO_FREE);
       // flow_offset
       Message array_3;
       CHECK_EQ(receiver->RecvFrom(&array_3, send_id), REMOVE_SUCCESS);
@@ -409,7 +419,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[7]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_3.data);
+        array_3.data,
+        AUTO_FREE);
       // CSR indptr
       Message array_4;
       CHECK_EQ(receiver->RecvFrom(&array_4, send_id), REMOVE_SUCCESS);
@@ -418,7 +429,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[9]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_4.data);
+        array_4.data,
+        AUTO_FREE);
       // CSR indice
       Message array_5;
       CHECK_EQ(receiver->RecvFrom(&array_5, send_id), REMOVE_SUCCESS);
@@ -427,7 +439,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[11]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_5.data);
+        array_5.data,
+        AUTO_FREE);
       // CSR edge_ids
       Message array_6;
       CHECK_EQ(receiver->RecvFrom(&array_6, send_id), REMOVE_SUCCESS);
@@ -436,7 +449,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
         {meta.data_shape_[13]},
         DLDataType{kDLInt, 64, 1},
         DLContext{kDLCPU, 0},
-        array_6.data);
+        array_6.data,
+        AUTO_FREE);
       // Create CSR
       CSRPtr csr(new CSR(indptr, indice, edge_ids));
       nf->graph = GraphPtr(new ImmutableGraph(csr, nullptr));
@@ -452,73 +466,60 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
 ////////////////////////// Distributed KVStore Components ////////////////////////////////
 
 
-DGL_REGISTER_GLOBAL("network._CAPI_SenderSendKVMsg")
-.set_body([] (DGLArgs args, DGLRetValue* rv) {
-    int args_count = 0;
-    CommunicatorHandle chandle = args[args_count++];
-    int recv_id = args[args_count++];
-    KVStoreMsg kv_msg;
-    kv_msg.msg_type = args[args_count++];
-    kv_msg.rank = args[args_count++];
-    network::Sender* sender = static_cast<network::Sender*>(chandle);
-    if (kv_msg.msg_type != kFinalMsg && kv_msg.msg_type != kBarrierMsg) {
-      std::string name = args[args_count++];
-      kv_msg.name = name;
-      if (kv_msg.msg_type != kIPIDMsg) {
-        kv_msg.id = args[args_count++];
-      }
-      if (kv_msg.msg_type != kPullMsg && kv_msg.msg_type != kIPIDMsg) {
-        kv_msg.data = args[args_count++];
-      }
-    }
+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);
+  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 &&
-        kv_msg.msg_type != kIPIDMsg) {
+  if (kv_msg->msg_type != kFinalMsg &&
+      kv_msg->msg_type != kBarrierMsg &&
+      kv_msg->msg_type != kIPIDMsg) {
     // Send ArrayMeta
-      ArrayMeta meta(kv_msg.msg_type);
-      meta.AddArray(kv_msg.id);
-      if (kv_msg.msg_type != kPullMsg) {
-        meta.AddArray(kv_msg.data);
+    ArrayMeta meta(kv_msg->msg_type);
+    meta.AddArray(kv_msg->id);
+    if (kv_msg->msg_type != kPullMsg) {
+      meta.AddArray(kv_msg->data);
     }
     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
     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;
+    send_id_msg.data = static_cast<char*>(kv_msg->id->data);
+    send_id_msg.size = kv_msg->id.GetSize();
+    NDArray id = kv_msg->id;
     send_id_msg.deallocator = [id](Message*) {};
     CHECK_EQ(sender->Send(send_id_msg, recv_id), ADD_SUCCESS);
     // Send data NDArray
-      if (kv_msg.msg_type != kPullMsg) {
+    if (kv_msg->msg_type != kPullMsg) {
       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;
+      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);
     }
   }
-  });
+}
 
-DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
-.set_body([] (DGLArgs args, DGLRetValue* rv) {
-    CommunicatorHandle chandle = args[0];
-    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
+static KVStoreMsg* recv_kv_message(network::Receiver* receiver) {
   KVStoreMsg *kv_msg = new KVStoreMsg();
   // Recv kv_Msg
   Message recv_kv_msg;
@@ -529,8 +530,7 @@ DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
   if (kv_msg->msg_type == kFinalMsg ||
       kv_msg->msg_type == kBarrierMsg ||
       kv_msg->msg_type == kIPIDMsg) {
-      *rv = kv_msg;
-      return;
+    return kv_msg;
   }
   // Recv ArrayMeta
   Message recv_meta_msg;
@@ -545,7 +545,8 @@ DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
     {meta.data_shape_[1]},
     DLDataType{kDLInt, 64, 1},
     DLContext{kDLCPU, 0},
-      recv_id_msg.data);
+    recv_id_msg.data,
+    AUTO_FREE);
   // Recv Data NDArray
   if (kv_msg->msg_type != kPullMsg) {
     Message recv_data_msg;
@@ -559,9 +560,39 @@ DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
       vec_shape,
       DLDataType{kDLFloat, 32, 1},
       DLContext{kDLCPU, 0},
-        recv_data_msg.data);
+      recv_data_msg.data,
+      AUTO_FREE);
   }
-    *rv = kv_msg;
+  return kv_msg;
+}
+
+DGL_REGISTER_GLOBAL("network._CAPI_SenderSendKVMsg")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    int args_count = 0;
+    CommunicatorHandle chandle = args[args_count++];
+    int recv_id = args[args_count++];
+    KVStoreMsg kv_msg;
+    kv_msg.msg_type = args[args_count++];
+    kv_msg.rank = args[args_count++];
+    network::Sender* sender = static_cast<network::Sender*>(chandle);
+    if (kv_msg.msg_type != kFinalMsg && kv_msg.msg_type != kBarrierMsg) {
+      std::string name = args[args_count++];
+      kv_msg.name = name;
+      if (kv_msg.msg_type != kIPIDMsg) {
+        kv_msg.id = args[args_count++];
+      }
+      if (kv_msg.msg_type != kPullMsg && kv_msg.msg_type != kIPIDMsg) {
+        kv_msg.data = args[args_count++];
+      }
+    }
+    send_kv_message(sender, &kv_msg, recv_id, AUTO_FREE);
+  });
+
+DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+    CommunicatorHandle chandle = args[0];
+    network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
+    *rv = recv_kv_message(receiver);
   });
 
 DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgType")
@@ -606,5 +637,120 @@ DGL_REGISTER_GLOBAL("network._CAPI_DeleteKVMsg")
     delete msg;
   });
 
+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];
+      }
+    }
+    row_size *= sizeof(float);
+    // 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];
+          local_ids.push_back(local_id);
+          local_ids_orginal.push_back(i);
+        } else {
+          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) {
+          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;
+    for (int i = 0; i < remote_ids.size(); ++i) {
+      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())},
+                                         DLDataType{kDLInt, 64, 1},
+                                         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);
+#endif
+        msg_count++;
+      }
+    }
+    char *return_data = new char[ID_size*row_size];
+    // Copy local data
+#pragma omp parallel for
+    for (int64_t i = 0; i < local_ids.size(); ++i) {
+      memcpy(return_data + local_ids_orginal[i] * row_size,
+             local_data_char + local_ids[i] * row_size,
+             row_size);
+    }
+    // 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);
+      for (size_t n = 0; n < id_size; ++n) {
+        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,
+                          DLDataType{kDLFloat, 32, 1},
+                          DLContext{kDLCPU, 0},
+                          return_data,
+                          AUTO_FREE);
+    *rv = res_tensor;
+  });
+
 }  // namespace network
 }  // namespace dgl