[tune] TensorFlow Distributed Trainable (#11876)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>

python/ray/tune/tests/test_tensorflow_trainable.py

@@ -0,0 +1,61 @@
+import pytest
+import ray
+from ray.tune.integration.tensorflow import DistributedTrainableCreator
+from ray.tune.examples.tf_distributed_keras_example import train_mnist
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_connect_cluster():
+    try:
+        address_info = ray.init(address="auto")
+    except Exception as e:
+        pytest.skip(str(e))
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_single_step(ray_start_2_cpus):  # noqa: F811
+    trainable_cls = DistributedTrainableCreator(train_mnist)
+    trainer = trainable_cls()
+    trainer.train()
+    trainer.stop()
+
+
+def test_step_after_completion(ray_start_2_cpus):  # noqa: F811
+    trainable_cls = DistributedTrainableCreator(train_mnist, num_workers=2)
+    trainer = trainable_cls(config={"epochs": 1})
+    with pytest.raises(RuntimeError):
+        for i in range(10):
+            trainer.train()
+
+
+def test_validation(ray_start_2_cpus):  # noqa: F811
+    def bad_func(a, b, c):
+        return 1
+
+    t_cls = DistributedTrainableCreator(bad_func)
+    with pytest.raises(ValueError):
+        t_cls()
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))