Addd a wrapper to skip flaky unit tests.

apex/testing/common_utils.py

@@ -10,6 +10,7 @@ import unittest
 
 
 TEST_WITH_ROCM = os.getenv('APEX_TEST_WITH_ROCM', '0') == '1'
+SKIP_FLAKY_TEST = os.getenv('APEX_SKIP_FLAKY_TEST', '0') == '1'
 
 ## Wrapper to skip the unit tests.
 def skipIfRocm(fn):
@@ -20,3 +21,13 @@ def skipIfRocm(fn):
         else:
             fn(*args, **kwargs)
     return wrapper
+
+## Wrapper to skip the flaky unit tests.
+def skipFlakyTest(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        if SKIP_FLAKY_TEST:
+            raise unittest.SkipTest("Test is flaky.")
+        else:
+            fn(*args, **kwargs)
+    return wrapper

tests/L0/run_rocm.sh

 #!/bin/bash
-APEX_TEST_WITH_ROCM=1 python run_test.py
+APEX_TEST_WITH_ROCM=1 APEX_SKIP_FLAKY_TEST=1 python run_test.py

tests/L0/run_test.py

@@ -2,6 +2,7 @@ import unittest
 import sys
 
 from apex.testing.common_utils import TEST_WITH_ROCM
+from apex.testing.common_utils import SKIP_FLAKY_TEST
 
 test_dirs = ["run_amp", "run_fp16util", "run_optimizers", "run_fused_layer_norm", "run_pyprof_nvtx", "run_pyprof_data", "run_mlp"]
 
@@ -15,7 +16,7 @@ runner = unittest.TextTestRunner(verbosity=2)
 errcode = 0
 
 for test_dir in test_dirs:
-    if (test_dir in ROCM_BLACKLIST) and TEST_WITH_ROCM:
+    if (test_dir in ROCM_BLACKLIST) and TEST_WITH_ROCM and SKIP_FLAKY_TEST:
         continue
     suite = unittest.TestLoader().discover(test_dir)