[incubator-mxnet] branch master updated: Adds min cuda version assertion decorator for unit tests (#17054)

[zh...@apache.org]

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

zhasheng 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 3d38dbd  Adds min cuda version assertion decorator for unit tests (#17054)
3d38dbd is described below

commit 3d38dbde744954854015919d4faf56ac1aea16de
Author: perdasilva <pe...@gmail.com>
AuthorDate: Sat Dec 14 22:47:42 2019 +0100

    Adds min cuda version assertion decorator for unit tests (#17054)
---
 tests/python/gpu/test_operator_gpu.py |  4 ++-
 tests/python/unittest/common.py       | 54 +++++++++++++++++++++++++++++++----
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
index 7d23c2c..e548217 100644
--- a/tests/python/gpu/test_operator_gpu.py
+++ b/tests/python/gpu/test_operator_gpu.py
@@ -30,7 +30,7 @@ from mxnet import autograd
 
 curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
 sys.path.insert(0, os.path.join(curr_path, '../unittest'))
-from common import setup_module, with_seed, teardown, assert_raises_cudnn_not_satisfied
+from common import setup_module, with_seed, teardown, assert_raises_cudnn_not_satisfied, assert_raises_cuda_not_satisfied
 from common import run_in_spawned_process
 from test_operator import *
 from test_numpy_ndarray import *
@@ -2696,6 +2696,7 @@ def check_multihead_attention_selfatt(dtype):
         assert(grads_orig[k].shape == grads_opti[k].shape)
         assert_allclose(grads_orig[k], grads_opti[k], rtol=1e-2, atol=1e-3)
 
+@assert_raises_cuda_not_satisfied(min_version='9.1')
 def test_multihead_attention_selfatt():
     for dtype in ['float16', 'float32']:
         check_multihead_attention_selfatt(dtype=dtype)
@@ -2859,6 +2860,7 @@ def check_multihead_attention_encdec(dtype):
         assert(grads_orig[k].shape == grads_opti[k].shape)
         assert_allclose(grads_orig[k], grads_opti[k], rtol=1e-2, atol=1e-3)
 
+@assert_raises_cuda_not_satisfied(min_version='9.1')
 def test_multihead_attention_encdec():
     for dtype in ['float16', 'float32']:
         check_multihead_attention_encdec(dtype=dtype)
diff --git a/tests/python/unittest/common.py b/tests/python/unittest/common.py
index 816508f..5a46046 100644
--- a/tests/python/unittest/common.py
+++ b/tests/python/unittest/common.py
@@ -99,14 +99,42 @@ def random_seed(seed=None):
         random.seed(next_seed)
 
 
-def assert_raises_cudnn_not_satisfied(min_version):
+def _assert_raise_cuxx_version_not_satisfied(min_version, cfg):
+
+    def less_than(version_left, version_right):
+        """Compares two version strings in the format num(.[num])*"""
+        if not version_left or not version_right:
+            return False
+
+        left = version_left.split(".")
+        right = version_right.split(".")
+
+        # 0 pad shortest version - e.g. 
+        # less_than("9.1", "9.1.9") == less_than("9.1.0", "9.1.9")
+        longest = max(len(left), len(right))
+        left.extend([0] * (longest - len(left)))
+        right.extend([0] * (longest - len(right)))
+
+        # compare each of the version components
+        for l, r in zip(left, right):
+            if int(r) < int(l):
+                return False
+
+            # keep track of how many are the same
+            if int(r) == int(l):
+                longest = longest - 1
+
+        # longest = 0 mean version_left == version_right -> False
+        # longest > 0 version_left < version_right -> True
+        return longest > 0
+
     def test_helper(orig_test):
         @make_decorator(orig_test)
         def test_new(*args, **kwargs):
-            cudnn_off = os.getenv('CUDNN_OFF_TEST_ONLY') == 'true'
-            cudnn_env_version = os.getenv('CUDNN_VERSION', None if cudnn_off else '7.3.1')
-            cudnn_test_disabled = cudnn_off or cudnn_env_version < min_version
-            if not cudnn_test_disabled or mx.context.current_context().device_type == 'cpu':
+            cuxx_off = os.getenv(cfg['TEST_OFF_ENV_VAR']) == 'true'
+            cuxx_env_version = os.getenv(cfg['VERSION_ENV_VAR'], None if cuxx_off else cfg['DEFAULT_VERSION'])
+            cuxx_test_disabled = cuxx_off or less_than(cuxx_env_version, min_version)
+            if not cuxx_test_disabled or mx.context.current_context().device_type == 'cpu':
                 orig_test(*args, **kwargs)
             else:
                 assert_raises((MXNetError, RuntimeError), orig_test, *args, **kwargs)
@@ -114,6 +142,22 @@ def assert_raises_cudnn_not_satisfied(min_version):
     return test_helper
 
 
+def assert_raises_cudnn_not_satisfied(min_version):
+    return _assert_raise_cuxx_version_not_satisfied(min_version, {
+        'TEST_OFF_ENV_VAR': 'CUDNN_OFF_TEST_ONLY',
+        'VERSION_ENV_VAR': 'CUDNN_VERSION',
+        'DEFAULT_VERSION': '7.3.1'
+    })
+
+
+def assert_raises_cuda_not_satisfied(min_version):
+    return _assert_raise_cuxx_version_not_satisfied(min_version, {
+        'TEST_OFF_ENV_VAR': 'CUDA_OFF_TEST_ONLY',
+        'VERSION_ENV_VAR': 'CUDA_VERSION',
+        'DEFAULT_VERSION': '10.1'
+    })
+
+
 def with_seed(seed=None):
     """
     A decorator for nosetests test functions that manages rng seeds.