caffe2 - support flaky operator tests for caffe2 build (#18155)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18155

- Make a python decorator caffe2_flaky for caffe2 operator unit tests.
- The environment variable CAFFE2_RUN_FLAKY_TESTS are now used to mark flaky test mode

During test run,
- If flaky tests mode are on, only flaky tests are run
- If flaky tests mode are off, only non-flaky tests are run

Mark ctc_beam_search_decoder_op_test as flaky

Reviewed By: ezyang, salexspb

Differential Revision: D14468816

fbshipit-source-id: dceb4a48daeb5437ad9cc714bef3343e9761f3a4

diff --git a/caffe2/python/operator_test/ctc_beam_search_decoder_op_test.py b/caffe2/python/operator_test/ctc_beam_search_decoder_op_test.py
index 4deef35..1ec7232 100644 (file)
--- a/caffe2/python/operator_test/ctc_beam_search_decoder_op_test.py
+++ b/caffe2/python/operator_test/ctc_beam_search_decoder_op_test.py
@@ -4,6 +4,7 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 from caffe2.python import core
+from caffe2.python.test_util import caffe2_flaky
 from collections import defaultdict, Counter
 from hypothesis import given
 import caffe2.python.hypothesis_test_util as hu
@@ -18,7 +19,7 @@ DEFAULT_PRUNE_THRESHOLD = 0.001
 
 
 class TestCTCBeamSearchDecoderOp(serial.SerializedTestCase):
-
+    @caffe2_flaky
     @serial.given(
         batch=st.sampled_from([1, 2, 4]),
         max_time=st.sampled_from([1, 8, 64]),

diff --git a/caffe2/python/test_util.py b/caffe2/python/test_util.py
index 1bebe50..a2cf3ac 100644 (file)
--- a/caffe2/python/test_util.py
+++ b/caffe2/python/test_util.py
@@ -10,6 +10,7 @@ from caffe2.python import core, workspace
 import unittest
 import os
 
+
 def rand_array(*dims):
     # np.random.rand() returns float instead of 0-dim array, that's why need to
     # do some tricks
@@ -51,6 +52,22 @@ def get_default_test_flags():
     ]
 
 
+def caffe2_flaky(test_method):
+    # This decorator is used to mark a test method as flaky.
+    # This is used in conjunction with the environment variable
+    # CAFFE2_RUN_FLAKY_TESTS that specifies "flaky tests" mode
+    # If flaky tests mode are on, only flaky tests are run
+    # If flaky tests mode are off, only non-flaky tests are run
+    # NOTE: the decorator should be applied as the top-level decorator
+    # in a test method.
+    test_method.__caffe2_flaky__ = True
+    return test_method
+
+
+def is_flaky_test_mode():
+    return os.getenv('CAFFE2_RUN_FLAKY_TESTS', '0') == '1'
+
+
 class TestCase(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -60,6 +77,15 @@ class TestCase(unittest.TestCase):
         core.SetEnginePref({}, {})
 
     def setUp(self):
+        # Skip tests based on whether we're in flaky test mode and
+        # the test is decorated as a flaky test.
+        test_method = getattr(self, self._testMethodName)
+        is_flaky_test = getattr(test_method, "__caffe2_flaky__", False)
+        if (is_flaky_test_mode() and not is_flaky_test):
+            raise unittest.SkipTest("Non-flaky tests are skipped in flaky test mode")
+        elif (not is_flaky_test_mode() and is_flaky_test):
+            raise unittest.SkipTest("Flaky tests are skipped in regular test mode")
+
         self.ws = workspace.C.Workspace()
         workspace.ResetWorkspace()