remove WaitToRead in dist-kvstore (#7489)

* remove waittoread in dist-kv * update
apache · Aug 16, 2017 · c79d4e5 · c79d4e5
1 parent 75ee597
commit c79d4e5
Showing 1 changed file with 19 additions and 21 deletions.
diff --git a/src/kvstore/kvstore_dist.h b/src/kvstore/kvstore_dist.h
@@ -118,23 +118,22 @@ class KVStoreDist : public KVStoreLocal {
       if (recv_buf.is_none()) {
         // it may happen for the first time a no-rank-0 worker pull the weight.
         recv_buf = NDArray(
-          grouped_vals[i][0]->shape(), pinned_ctx_, false, grouped_vals[i][0]->dtype());
+          grouped_vals[i][0]->shape(), pinned_ctx_, true, grouped_vals[i][0]->dtype());
       }
-#if MKL_EXPERIMENTAL == 1
-      mkl_set_tblob_eager_mode(recv_buf.data());
-#endif
-      real_t* data = static_cast<real_t*>(recv_buf.data().dptr_);
-      size_t size = recv_buf.shape().Size();
-
-      auto pull_from_servers = [this, key, data, size](
+      auto pull_from_servers = [this, key, recv_buf](
           RunContext rctx, Engine::CallbackOnComplete cb) {
         // convert to ps keys
+        size_t size = recv_buf.shape().Size();
         PSKV& pskv = EncodeKey(key, size);
-
-        // issue pull, false means no delete
+#if MKL_EXPERIMENTAL == 1
+        mkl_set_tblob_eager_mode(recv_buf.data());
+#endif
+        real_t* data = static_cast<real_t*>(recv_buf.data().dptr_);
+        // false means not to delete data when SArray is deleted
         auto vals = new ps::SArray<real_t>(data, size, false);
+        // issue pull
         CHECK_NOTNULL(ps_worker_)->ZPull(
-        pskv.keys, vals, &pskv.lens, 0, [vals, cb](){ delete vals; cb(); });
+            pskv.keys, vals, &pskv.lens, 0, [vals, cb](){ delete vals; cb(); });
       };
 
       CHECK_NOTNULL(Engine::Get())->PushAsync(
@@ -226,27 +225,26 @@ class KVStoreDist : public KVStoreLocal {
         send_buf = merged;  // avoid memory copy
       } else {
         if (send_buf.is_none()) {
-          send_buf = NDArray(merged.shape(), pinned_ctx_, false, merged.dtype());
+          send_buf = NDArray(merged.shape(), pinned_ctx_, true, merged.dtype());
         }
         CopyFromTo(merged, &send_buf);
       }
 
       // push to servers
-      send_buf.WaitToRead();
-      size_t size = send_buf.shape().Size();
-#if MKL_EXPERIMENTAL == 1
-      mkl_set_tblob_eager_mode(send_buf.data());
-#endif
-      real_t* data = static_cast<real_t*>(send_buf.data().dptr_);
       auto push_to_servers =
-          [this, key, data, size](RunContext rctx, Engine::CallbackOnComplete cb) {
-         // convert to ps keys
+          [this, key, send_buf](RunContext rctx, Engine::CallbackOnComplete cb) {
+        // convert to ps keys
+        size_t size = send_buf.shape().Size();
         PSKV& pskv = EncodeKey(key, size);
 
+#if MKL_EXPERIMENTAL == 1
+        mkl_set_tblob_eager_mode(send_buf.data());
+#endif
+        real_t* data = static_cast<real_t*>(send_buf.data().dptr_);
         // do push. false means no delete
         ps::SArray<real_t> vals(data, size, false);
         CHECK_NOTNULL(ps_worker_)->ZPush(
-        pskv.keys, vals, pskv.lens, 0, [cb]() { cb(); });
+            pskv.keys, vals, pskv.lens, 0, [cb]() { cb(); });
       };
       Engine::Get()->PushAsync(
           push_to_servers,