[incubator-mxnet] branch master updated: Fix failing GPU test on single GPU host (kvstore) (#12726)

[cmeier]

This is an automated email from the ASF dual-hosted git repository.

cmeier pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/master by this push:
     new 836ba78  Fix failing GPU test on single GPU host (kvstore) (#12726)
836ba78 is described below

commit 836ba78a53aa695cbc8feec5df4881223fea15c7
Author: Pedro Larroy <928489+lar...@users.noreply.github.com>
AuthorDate: Sat Oct 6 08:27:27 2018 -0700

    Fix failing GPU test on single GPU host (kvstore) (#12726)
    
    Fixes #10977
---
 tests/python/gpu/test_kvstore_gpu.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/tests/python/gpu/test_kvstore_gpu.py 
b/tests/python/gpu/test_kvstore_gpu.py
index 5e9b120..4232a59 100644
--- a/tests/python/gpu/test_kvstore_gpu.py
+++ b/tests/python/gpu/test_kvstore_gpu.py
@@ -58,6 +58,7 @@ def init_kv_with_str(stype='default', kv_type='local'):
 # Test seed 89411477 (module seed 1829754103) resulted in a py3-gpu CI runner 
core dump.
 # Not reproducible, so this test is back on random seeds.
 @with_seed()
+@unittest.skipIf(mx.context.num_gpus() < 2, "test_rsp_push_pull needs more 
than 1 GPU")
 def test_rsp_push_pull():
     def check_rsp_push_pull(kv_type, sparse_pull, is_push_cpu=True):
         kv = init_kv_with_str('row_sparse', kv_type)
@@ -65,7 +66,8 @@ def test_rsp_push_pull():
         push_ctxs = [mx.cpu(i) if is_push_cpu else mx.gpu(i) for i in range(2)]
         kv.push('e', [mx.nd.ones(shape, ctx=context).tostype('row_sparse') for 
context in push_ctxs])
 
-        def check_rsp_pull(kv, count, ctxs, sparse_pull, is_same_rowid=False, 
use_slice=False):
+        def check_rsp_pull(kv, ctxs, sparse_pull, is_same_rowid=False, 
use_slice=False):
+            count = len(ctxs)
             num_rows = shape[0]
             row_ids = []
             all_row_ids = np.arange(num_rows)
@@ -100,14 +102,14 @@ def test_rsp_push_pull():
                     expected_val[:] = 2
                     assert_almost_equal(retained, expected_val)
 
-        check_rsp_pull(kv, 1, [mx.gpu(0)], sparse_pull)
-        check_rsp_pull(kv, 1, [mx.cpu(0)], sparse_pull)
-        check_rsp_pull(kv, 4, [mx.gpu(i//2) for i in range(4)], sparse_pull)
-        check_rsp_pull(kv, 4, [mx.gpu(i//2) for i in range(4)], sparse_pull, 
is_same_rowid=True)
-        check_rsp_pull(kv, 4, [mx.cpu(i) for i in range(4)], sparse_pull)
-        check_rsp_pull(kv, 4, [mx.cpu(i) for i in range(4)], sparse_pull, 
is_same_rowid=True)
-        check_rsp_pull(kv, 4, [mx.gpu(i//2) for i in range(4)], sparse_pull, 
use_slice=True)
-        check_rsp_pull(kv, 4, [mx.cpu(i) for i in range(4)], sparse_pull, 
use_slice=True)
+        check_rsp_pull(kv, [mx.gpu(0)], sparse_pull)
+        check_rsp_pull(kv, [mx.cpu(0)], sparse_pull)
+        check_rsp_pull(kv, [mx.gpu(i//2) for i in range(4)], sparse_pull)
+        check_rsp_pull(kv, [mx.gpu(i//2) for i in range(4)], sparse_pull, 
is_same_rowid=True)
+        check_rsp_pull(kv, [mx.cpu(i) for i in range(4)], sparse_pull)
+        check_rsp_pull(kv, [mx.cpu(i) for i in range(4)], sparse_pull, 
is_same_rowid=True)
+        check_rsp_pull(kv, [mx.gpu(i//2) for i in range(4)], sparse_pull, 
use_slice=True)
+        check_rsp_pull(kv, [mx.cpu(i) for i in range(4)], sparse_pull, 
use_slice=True)
 
     envs = ["","1"]
     key  = "MXNET_KVSTORE_USETREE"