[BUGFIX] Fix race condition in kvstore.pushpull (#17007)

* add back gluon test

* fix typo

* change back gpu ctx

* also handle the case there some are pull and some are pushpull

* fix typo

Author: eric-haibin-lin

src/kvstore/kvstore_dist_server.h

@@ -364,21 +364,34 @@ class KVStoreDistServer {
       if (log_verbose_)  {
         LOG(INFO) << "sent response to " << update_buf->request.size() << " workers";
       }
+      /**
+       * Request can be for either push, pull or pushpull
+       * If pull flag is set, respond immediately with the updated values
+       * Otherwise, only send the notification
+       */
+      bool has_pull = false;
       for (const auto& req : update_buf->request) {
-        /**
-         * Request can be for either push, pull or pushpull
-         * If pull flag is set, respond immediately with the updated values
-         * Otherwise, only send the notification
-         */
-        if (req.pull) {
-          DefaultStorageResponse(type, key, req, req_data, server);
-        } else {
+        has_pull = has_pull || req.pull;
+      }
+      if (has_pull) {
+        // if there is a pull request, perform WaitToRead() once before DefaultStorageResponse
+        if (has_multi_precision_copy(type)) CopyFromTo(stored, store_[key]);
+        stored.WaitToRead();
+        for (const auto& req : update_buf->request) {
+          if (req.pull) {
+            DefaultStorageResponse(type, key, req, req_data, server);
+          }
+        }
+        update_buf->request.clear();
+      } else {
+        // otherwise, send response directly
+        for (const auto& req : update_buf->request) {
           server->Response(req);
         }
+        update_buf->request.clear();
+        if (has_multi_precision_copy(type)) CopyFromTo(stored, store_[key]);
+        stored.WaitToRead();
       }
-      update_buf->request.clear();
-      if (has_multi_precision_copy(type)) CopyFromTo(stored, store_[key]);
-      stored.WaitToRead();
     } else {
       update_buf->merged.WaitToRead();
     }

tests/nightly/dist_device_sync_kvstore.py

@@ -44,7 +44,10 @@ def check_diff_to_scalar(A, x, rank=None):
 def init_kv():
     # init kv dns keys
     kv.init(keys, [mx.nd.ones(shape)] * len(keys))
+    kv.init('9', mx.nd.ones(shape))
+    kv.init('10', mx.nd.ones(shape))
     kv.init('99', mx.nd.ones(big_shape))
+    kv.init('100', mx.nd.ones(big_shape))
     # worker info
     my_rank = kv.rank
     nworker = kv.num_workers
@@ -55,33 +58,30 @@ def init_kv():
 def test_sync_push_pull():
     kv, my_rank, nworker = init_kv()
     num_gpus = 2
-    def check_default_keys(kv, my_rank, nworker, nrepeat=3, offset=0, use_pushpull=False):
+    def check_default_keys(kv, my_rank, nworker, nrepeat=3):
         # checks pull after push in loop, because behavior during
         # consecutive pushes doesn't offer any guarantees
-        for i in range(offset, nrepeat):
+        for i in range(nrepeat):
             scale = my_rank + 1
             num = (nworker + 1) * nworker * rate * num_gpus / 2 * (i + 1) + 1
 
             arr = [mx.nd.ones(shape, ctx=mx.gpu(j)) * scale for j in range(num_gpus)]
             val = mx.nd.zeros(shape)
-            if use_pushpull:
-                kv.pushpull('3', arr, out=val)
-            else:
-                kv.push('3', arr)
-                kv.pull('3', out=val)
+            kv.push('9', arr)
+            kv.pull('9', out=val)
+            check_diff_to_scalar(val, num)
+            kv.pushpull('10', arr, out=val)
             check_diff_to_scalar(val, num)
 
             big_arr = [mx.nd.ones(big_shape, ctx=mx.gpu(j)) * scale for j in range(num_gpus)]
             big_val = mx.nd.zeros(big_shape)
-            if use_pushpull:
-                kv.pushpull('99', big_arr, out=big_val)
-            else:
-                kv.push('99', big_arr)
-                kv.pull('99', out=big_val)
+            kv.push('99', big_arr)
+            kv.pull('99', out=big_val)
+            check_diff_to_scalar(big_val, num)
+            kv.pushpull('100', big_arr, out=big_val)
             check_diff_to_scalar(big_val, num)
 
-    check_default_keys(kv, my_rank, nworker, nrepeat=3, offset=0, use_pushpull=False)
-    check_default_keys(kv, my_rank, nworker, nrepeat=3, offset=3, use_pushpull=True)
+    check_default_keys(kv, my_rank, nworker, nrepeat=3)
     print('worker ' + str(my_rank) + ' is done')
 
 def test_sync_init():
@@ -106,10 +106,12 @@ def check_trainer_kv_update(update_on_kv):
         x = params.get('x', shape=(10,1), lr_mult=1.0)
         params.initialize(ctx=[mx.cpu(0), mx.cpu(1)], init='zeros')
         try:
-            trainer = mx.gluon.Trainer(params, 'sgd', {'learning_rate': 0.1}, kvstore=kv, update_on_kvstore=update_on_kv)
+            trainer = mx.gluon.Trainer(params, 'sgd', {'learning_rate': 0.1},
+                                       kvstore=kv, update_on_kvstore=update_on_kv)
             trainer._init_kvstore()
             assert trainer._kv_initialized
-            assert trainer._update_on_kvstore is True
+            if update_on_kv is not None:
+                assert trainer._update_on_kvstore is update_on_kv
         except ValueError:
             assert update_on_kv is False
 
@@ -122,3 +124,4 @@ def check_trainer_kv_update(update_on_kv):
 if __name__ == "__main__":
     test_sync_init()
     test_sync_push_pull()
+    test_gluon_trainer_type()