Repository: tensorflow/addons

Issue:
Enable multiprocessing when testing with GPU and support distributed strategies in the tests.
**Describe the feature and the current behavior/state.**

Here I'm not going to discuss the bazel case as it's much more complicated to handle, and we currently advertize using pytest anyway to run the tests. We can of course make sure everything is compatible though.

This revamping of gpu testing has multiple objectives:

* The tests should behave the same weither the contributor has a gpu or not. Meaning we shouldn't run all the tests on a gpu just because a gpu is available, otherwise it hurts reproducibility.
* The test suite should be able to run with multiple workers in kokoro or when a user has multiple gpus. Pytest should use all gpus visible by the main process.
* We need to support testing with distributed strategies. Currently it doesn't work. A fix has been started in #1209 but we need to update it for pytest. 
* Making the whole thing simple to use and to maintain. Notably, we would get rid of this file: https://github.com/tensorflow/addons/blob/master/tools/testing/parallel_gpu_execute.sh which is quite hard to work on.


To do all that, here is my proposal:

Stuff to know: 
* Pytest-xdist uses multiprocessing to run the tests, not multithreading. 
* 2 environement variables are available in each of the workers to identify them. https://github.com/pytest-dev/pytest-xdist#identifying-the-worker-process-during-a-test


### Test workers
Suppose we have a machine with 10CPUs and 4 GPUs, 10 processes will start to run the test suite. Workers 0 to 3 will have ownership of one GPU each (we can use CUDA_VISIBLE_DEVICES to enforce that, but I'm not even sure that's needed with the proposed implementation). Workers 4 to 9 will have no gpu available.

### Virtual devices
Each of those processes, when starting, will split their physical device into 2 virtual device. Tests that just need to run on gpu will use the first of those virtual devices. Processes which need to test distributed strategies will use the two of them. We assume here that 2 virtual devices are enough to test distributed strategies.

### Impact on the contributors:
For this whole machinery to work, we need to know which test needs to run on CPU, GPU, or in distributed strategies. To do that we'll use pytest markers: `@pytest.mark.....`

* By default, if no marker is found, the test will run on CPU: `with device("CPU:0")`. It's equivalent to 
`@pytest.mark.run_on(["cpu"])`.
* To run with gpu only: `@pytest.mark.run_on(["gpu"])`.
* To run on the cpu and gpu: `@pytest.mark.run_on(["cpu", "gpu"])` (test runs twice)
* To run in within a distributed strategy `@pytest.mark.run_on(["distributed strategy"])`. (runs once here).
* To run with everything `@pytest.mark.run_on(["cpu", "gpu", "distributed strategy"])`
* To make crazy stuff, and not run the test in any device scope: `@pytest.mark.no_device_scope`. Then the contributor can do whatever he/she wants in the test.

Of course, if no gpu is available, we just skip the tests needing a distribution strategy or the gpu. Contributors who handle the devices manually have to make sure to skip manually the test if the gpu is used.

Since gpu are often the scarsest ressource (nb gpus << nb cpus), tests needing the gpu will also be marked with `@pytest.mark.tryfirst` to ensure that we don't have workers starvation at the end (to get maximum speed).

To implement that, we need first to convert all tests to pytest (as opposed to unittest) it's currently 80% done and thanks a lot @autoih for putting a LOT of work into that.



**Relevant information**
- Are you willing to contribute it (yes/no): yes
- Are you willing to maintain it going forward? (yes/no): yes
- Is there a relevant academic paper? (if so, where): no
- Is there already an implementation in another framework? (if so, where): no
- Was it part of tf.contrib? (if so, where): no

**Which API type would this fall under (layer, metric, optimizer, etc.)**

Testing

**Who will benefit with this feature?**

Contributors with gpu, CI.

**Any other info.**

I believe that the implementation will first go in tensorflow addons because we have 4 GPUs available in the CI. Later on when it's stable we can split it from tensorflow addons and make it a separate pytest plugin with a public API.

Comments welcome. Especially from @Squadrick , @hyang0129 , @seanpmorgan since I'm not a ninja of tf.device.


Discussion:

Resolving patch:
diff --git a/tensorflow_addons/conftest.py b/tensorflow_addons/conftest.py
--- a/tensorflow_addons/conftest.py
+++ b/tensorflow_addons/conftest.py
@@ -8,6 +8,7 @@
     pytest_configure,
     device,
     pytest_generate_tests,
+    pytest_collection_modifyitems,
 )
 
 # fixtures present in this file will be available


Test patch:
diff --git a/tensorflow_addons/optimizers/tests/conditional_gradient_test.py b/tensorflow_addons/optimizers/tests/conditional_gradient_test.py
--- a/tensorflow_addons/optimizers/tests/conditional_gradient_test.py
+++ b/tensorflow_addons/optimizers/tests/conditional_gradient_test.py
@@ -190,7 +190,7 @@ def test_basic_nuclear(use_resource):
     #       to address issue #36764
     for i, dtype in enumerate(
         _dtypes_with_checking_system(
-            use_gpu=tf.test.is_gpu_available(), system=platform.system()
+            use_gpu=test_utils.is_gpu_available(), system=platform.system()
         )
     ):
 
@@ -276,7 +276,7 @@ def test_minimize_sparse_resource_variable_nuclear():
     # TODO:
     #       to address issue #347 and #36764.
     for dtype in _dtypes_with_checking_system(
-        use_gpu=tf.test.is_gpu_available(), system=platform.system()
+        use_gpu=test_utils.is_gpu_available(), system=platform.system()
     ):
         var0 = tf.Variable([[1.0, 2.0]], dtype=dtype)
 
@@ -316,7 +316,7 @@ def loss():
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_tensor_learning_rate_and_conditional_gradient_nuclear():
     for dtype in _dtypes_with_checking_system(
-        use_gpu=tf.test.is_gpu_available(), system=platform.system()
+        use_gpu=test_utils.is_gpu_available(), system=platform.system()
     ):
         # TODO:
         # Based on issue #36764 in the following link,
@@ -845,7 +845,7 @@ def test_like_dist_belief_frobenius_cg01():
 def test_sparse_frobenius():
     # TODO:
     #       To address the issue #347.
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         var0 = tf.Variable(tf.zeros([4, 2], dtype=dtype))
         var1 = tf.Variable(tf.constant(1.0, dtype, [4, 2]))
         grads0 = tf.IndexedSlices(
@@ -1001,7 +1001,7 @@ def test_sharing_nuclear():
     # TODO:
     #       To address the issue #36764.
     for dtype in _dtypes_with_checking_system(
-        use_gpu=tf.test.is_gpu_available(), system=platform.system()
+        use_gpu=test_utils.is_gpu_available(), system=platform.system()
     ):
         var0 = tf.Variable([1.0, 2.0], dtype=dtype)
         var1 = tf.Variable([3.0, 4.0], dtype=dtype)
@@ -1340,7 +1340,7 @@ def test_sparse_nuclear():
     # TODO:
     #       To address the issue #347 and issue #36764.
     for dtype in _dtypes_with_checking_system(
-        use_gpu=tf.test.is_gpu_available(), system=platform.system()
+        use_gpu=test_utils.is_gpu_available(), system=platform.system()
     ):
         var0 = tf.Variable(tf.zeros([4, 2], dtype=dtype))
         var1 = tf.Variable(tf.constant(1.0, dtype, [4, 2]))
diff --git a/tensorflow_addons/optimizers/tests/lamb_test.py b/tensorflow_addons/optimizers/tests/lamb_test.py
--- a/tensorflow_addons/optimizers/tests/lamb_test.py
+++ b/tensorflow_addons/optimizers/tests/lamb_test.py
@@ -67,7 +67,7 @@ def get_beta_accumulators(opt, dtype):
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_sparse():
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize tf for numpy implementation.
         m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
         var0_np = np.array([1.0, 1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -113,7 +113,7 @@ def test_sparse():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_basic_with_learning_rate_decay():
-    for i, dtype in enumerate(_dtypes_to_test(use_gpu=tf.test.is_gpu_available())):
+    for i, dtype in enumerate(_dtypes_to_test(use_gpu=test_utils.is_gpu_available())):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -162,7 +162,7 @@ def test_basic_with_learning_rate_decay():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_basic_with_learning_rate_inverse_time_decay():
-    for i, dtype in enumerate(_dtypes_to_test(use_gpu=tf.test.is_gpu_available())):
+    for i, dtype in enumerate(_dtypes_to_test(use_gpu=test_utils.is_gpu_available())):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -204,7 +204,7 @@ def test_basic_with_learning_rate_inverse_time_decay():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_tensor_learning_rate():
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -239,7 +239,7 @@ def test_tensor_learning_rate():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_sharing():
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -293,7 +293,7 @@ def loss():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_resource():
-    for i, dtype in enumerate(_dtypes_to_test(use_gpu=tf.test.is_gpu_available())):
+    for i, dtype in enumerate(_dtypes_to_test(use_gpu=test_utils.is_gpu_available())):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
diff --git a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
--- a/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
+++ b/tensorflow_addons/optimizers/tests/weight_decay_optimizers_test.py
@@ -53,7 +53,7 @@ def do_test(
             the optimizer_params in the update_fn.
     """
     # TODO: Fix #347 issue
-    if do_sparse and tf.test.is_gpu_available():
+    if do_sparse and test_utils.is_gpu_available():
         pytest.skip("Wait #347 to be fixed")
 
     # Initialize variables for numpy implementation.
@@ -112,7 +112,7 @@ def do_test_sparse_repeated_indices(dtype, optimizer, **optimizer_kwargs):
             the optimizer_params in the update_fn.
     """
     # TODO: Fix #347 issue
-    if tf.test.is_gpu_available():
+    if test_utils.is_gpu_available():
         pytest.skip("Wait #347 to be fixed")
 
     repeated_index_update_var = tf.Variable([[1.0], [2.0]], dtype=dtype)
diff --git a/tensorflow_addons/optimizers/tests/yogi_test.py b/tensorflow_addons/optimizers/tests/yogi_test.py
--- a/tensorflow_addons/optimizers/tests/yogi_test.py
+++ b/tensorflow_addons/optimizers/tests/yogi_test.py
@@ -96,7 +96,7 @@ def _dtypes_to_test(use_gpu):
 
 
 def do_test_sparse(beta1=0.0, l1reg=0.0, l2reg=0.0):
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 1.0, 0.0, 1.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -169,7 +169,7 @@ def test_sparse_momentum_regularization():
 
 
 def test_sparse_repeated_indices():
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         repeated_index_update_var = tf.Variable([[1.0], [2.0]], dtype=dtype)
         aggregated_update_var = tf.Variable([[1.0], [2.0]], dtype=dtype)
         grad_repeated_index = tf.IndexedSlices(
@@ -199,7 +199,7 @@ def test_sparse_repeated_indices():
 
 
 def do_test_basic(beta1=0.0, l1reg=0.0, l2reg=0.0):
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 1.0, 0.0, 1.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -265,7 +265,7 @@ def test_basic_momentum_regularization():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_tensor_learning_rate():
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 1.0, 0.0, 1.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
@@ -301,7 +301,7 @@ def test_tensor_learning_rate():
 
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_sharing():
-    for dtype in _dtypes_to_test(use_gpu=tf.test.is_gpu_available()):
+    for dtype in _dtypes_to_test(use_gpu=test_utils.is_gpu_available()):
         # Initialize variables for numpy implementation.
         m0, v0, m1, v1 = 0.0, 1.0, 0.0, 1.0
         var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
diff --git a/tensorflow_addons/utils/test_utils.py b/tensorflow_addons/utils/test_utils.py
--- a/tensorflow_addons/utils/test_utils.py
+++ b/tensorflow_addons/utils/test_utils.py
@@ -31,6 +31,11 @@
 WORKER_ID = int(os.environ.get("PYTEST_XDIST_WORKER", "gw0")[2])
 NUMBER_OF_GPUS = len(tf.config.list_physical_devices("GPU"))
 
+
+def is_gpu_available():
+    return NUMBER_OF_GPUS >= 1
+
+
 # Some configuration before starting the tests.
 
 # we only need one core per worker.
@@ -39,18 +44,22 @@
 tf.config.threading.set_intra_op_parallelism_threads(1)
 tf.config.threading.set_inter_op_parallelism_threads(1)
 
-if NUMBER_OF_GPUS != 0:
+if is_gpu_available():
     # We use only the first gpu at the moment. That's enough for most use cases.
-    # split the first gpu into chunks of 100MB per pytest worker.
+    # split the first gpu into chunks of 100MB per virtual device.
     # It's the user's job to limit the amount of pytest workers depending
     # on the available memory.
     # In practice, each process takes a bit more memory.
     # There must be some kind of overhead but it's not very big (~200MB more)
+    # Each worker has two virtual devices.
+    # When running on gpu, only the first device is used. The other one is used
+    # in distributed strategies.
     first_gpu = tf.config.list_physical_devices("GPU")[0]
+    virtual_gpus = [
+        tf.config.LogicalDeviceConfiguration(memory_limit=100) for _ in range(2)
+    ]
 
-    tf.config.set_logical_device_configuration(
-        first_gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=100)],
-    )
+    tf.config.set_logical_device_configuration(first_gpu, virtual_gpus)
 
 
 def finalizer():
@@ -96,6 +105,20 @@ def pytest_addoption(parser):
     )
 
 
+def gpus_for_testing():
+    """For the moment it's very simple, but it might change in the future,
+    with multiple physical gpus for example. So it's better if this function
+    is called rather than hardcoding the gpu devices in the tests.
+    """
+    if not is_gpu_available():
+        raise SystemError(
+            "You are trying to get some gpus for testing but no gpu is available on "
+            "your system. \nDid you forget to use `@pytest.mark.needs_gpu` on your test"
+            " so that it's skipped automatically when no gpu is available?"
+        )
+    return ["gpu:0", "gpu:1"]
+
+
 @pytest.fixture(scope="session", autouse=True)
 def set_global_variables(request):
     if request.config.getoption("--skip-custom-ops"):
@@ -114,7 +137,11 @@ def device(request):
     requested_device = request.param
     if requested_device == "no_device":
         yield requested_device
-    else:
+    elif requested_device == tf.distribute.MirroredStrategy:
+        strategy = requested_device(gpus_for_testing())
+        with strategy.scope():
+            yield strategy
+    elif isinstance(requested_device, str):
         if requested_device in ["cpu", "gpu"]:
             # we use GPU:0 because the virtual device we created is the
             # only one in the first GPU (so first in the list of virtual devices).
@@ -126,13 +153,10 @@ def device(request):
 
 
 def get_marks(device_name):
-    marks = []
-    if device_name == "gpu":
-        marks.append(pytest.mark.needs_gpu)
-        if NUMBER_OF_GPUS == 0:
-            skip_message = "The gpu is not available."
-            marks.append(pytest.mark.skip(reason=skip_message))
-    return marks
+    if device_name == "gpu" or device_name == tf.distribute.MirroredStrategy:
+        return [pytest.mark.needs_gpu]
+    else:
+        return []
 
 
 def pytest_generate_tests(metafunc):
@@ -149,6 +173,13 @@ def pytest_generate_tests(metafunc):
     metafunc.parametrize("device", parameters, indirect=True)
 
 
+def pytest_collection_modifyitems(items):
+    for item in items:
+        if item.get_closest_marker("needs_gpu") is not None:
+            if not is_gpu_available():
+                item.add_marker(pytest.mark.skip("The gpu is not available."))
+
+
 def assert_allclose_according_to_type(
     a,
     b,
diff --git a/tensorflow_addons/utils/tests/test_utils_test.py b/tensorflow_addons/utils/tests/test_utils_test.py
--- a/tensorflow_addons/utils/tests/test_utils_test.py
+++ b/tensorflow_addons/utils/tests/test_utils_test.py
@@ -1,10 +1,45 @@
 import random
 
 import numpy as np
+import pytest
 import tensorflow as tf
+from tensorflow_addons.utils import test_utils
 
 
 def test_seed_is_set():
     assert random.randint(0, 10000) == 6311
     assert np.random.randint(0, 10000) == 2732
     assert tf.random.uniform([], 0, 10000, dtype=tf.int64).numpy() == 9457
+
+
+@pytest.mark.with_device(["cpu", "gpu", tf.distribute.MirroredStrategy])
+def test_all_scopes(device):
+    assert isinstance(device, str) or isinstance(device, tf.distribute.Strategy)
+
+
+def train_small_model():
+    model_input = tf.keras.layers.Input((3,))
+    model_output = tf.keras.layers.Dense(4)(model_input)
+    model = tf.keras.Model(model_input, model_output)
+    model.compile(loss="mse")
+
+    x = np.random.uniform(size=(5, 3))
+    y = np.random.uniform(size=(5, 4))
+    model.fit(x, y, epochs=1)
+
+
+@pytest.mark.with_device([tf.distribute.MirroredStrategy])
+def test_distributed_strategy(device):
+    assert isinstance(device, tf.distribute.Strategy)
+    train_small_model()
+
+
+@pytest.mark.with_device(["no_device"])
+@pytest.mark.needs_gpu
+def test_custom_device_placement():
+    with tf.device(test_utils.gpus_for_testing()[0]):
+        train_small_model()
+
+    strategy = tf.distribute.MirroredStrategy(test_utils.gpus_for_testing())
+    with strategy.scope():
+        train_small_model()