code: def test_pushpull():
      num_gpus = 2
  
      def check_default_keys(nrepeat=3):
          # init kv dns keys
          kv.broadcast(
              "3",
              mx.nd.ones(shape, ctx=mx.gpu()),
              mx.nd.ones(shape, ctx=mx.gpu()),
          )
          kv.broadcast(
              "99",
              mx.nd.ones(big_shape, ctx=mx.gpu()),
              mx.nd.ones(big_shape, ctx=mx.gpu()),
          )
          for _ in range(nrepeat):
              scale = my_rank + 1
              num = (my_num_workers + 1) * my_num_workers * num_gpus / 2
  
              arrs = [
                  mx.nd.ones(shape, ctx=mx.gpu(j)) * scale
                  for j in range(num_gpus)
              ]
              # inplace
              kv.pushpull("3", arrs)
              for arr in arrs:
                  check_diff_to_scalar(arr, num)
  
              big_arrs = [
                  mx.nd.ones(big_shape, ctx=mx.gpu(j)) * scale
                  for j in range(num_gpus)
              ]
              # inplace
              kv.pushpull("99", big_arrs)
              for big_arr in big_arrs:
                  check_diff_to_scalar(big_arr, num)
  
      check_default_keys(nrepeat=3)
      print("worker " + str(my_rank) + " is done")
  
file_path: mxnet/tests/nightly/dist_device_sync_kvstore_custom.py
imports: import sys
  import mxnet
  import numpy
  import numpy.random
  import time
  import argparse
label_desc: Use image node_id 2 for calling a global function with example usage: test_pushpull() without return types