[BugFix] Fix a small bug on kvstore (#2058)

* small fix

* update

python/dgl/distributed/rpc.py

@@ -991,11 +991,10 @@ def register_ctrl_c():
     """
     _CAPI_DGLRPCHandleCtrlC()
 
-def copy_data_to_shared_memory(source, dst):
+def copy_data_to_shared_memory(dst, source):
     """Copy tensor data to shared-memory tensor
     """
-    _CAPI_DGLCopyDataToSharedMemory(F.zerocopy_to_dgl_ndarray(source),
-                                    F.zerocopy_to_dgl_ndarray(dst))
+    F.zerocopy_to_dgl_ndarray(dst).copyfrom(F.zerocopy_to_dgl_ndarray(source))
 
 ############### Some basic services will be defined here #############
 

src/rpc/rpc.cc

@@ -476,15 +476,5 @@ DGL_REGISTER_GLOBAL("distributed.rpc._CAPI_DGLRPCFastPull")
   *rv = res_tensor;
 });
 
-DGL_REGISTER_GLOBAL("distributed.rpc._CAPI_DGLCopyDataToSharedMemory")
-.set_body([] (DGLArgs args, DGLRetValue* rv) {
-  NDArray source = args[0];
-  NDArray dst = args[1];
-  CHECK_EQ(source.GetSize(), dst.GetSize());
-  char* src_ptr = static_cast<char*>(source->data);
-  char* dst_ptr = static_cast<char*>(dst->data);
-  memcpy(src_ptr, dst_ptr, dst.GetSize());
-});
-
 }  // namespace rpc
 }  // namespace dgl