Add tensorflow basics test

Author: BichengYing

.pylintrc

@@ -236,7 +236,10 @@ contextmanager-decorators=contextlib.contextmanager
 # List of members which are set dynamically and missed by pylint inference
 # system, and so shouldn't trigger E1101 when accessed. Python regular
 # expressions are accepted.
-generated-members=numpy.*,torch.*,mpi4py.MPI.*,bluefog.torch.mpi_lib.*,tensorflow.*
+generated-members=numpy.*,torch.*,mpi4py.MPI.*,
+                  tensorflow.*,
+                  bluefog.torch.mpi_lib.*,
+                  bluefog.tensorflow.mpi_lib.*,
 
 # Tells whether missing members accessed in mixin class should be ignored. A
 # mixin class is detected if its name ends with "mixin" (case insensitive).
@@ -263,7 +266,7 @@ ignored-classes=optparse.Values,thread._local,_thread._local
 # (useful for modules/projects where namespaces are manipulated during runtime
 # and thus existing member attributes cannot be deduced by static analysis. It
 # supports qualified module names, as well as Unix pattern matching.
-ignored-modules=
+ignored-modules=tensorflow
 
 # Show a hint with possible names when a member name was not found. The aspect
 # of finding the hint is based on edit distance.

Makefile

@@ -9,6 +9,9 @@ test_torch_basic:
 test_torch_ops:
 	${MPIRUN} -np 4 ${PYTEST} ./test/torch_ops_test.py
 
+test_tensorflow_basic:
+	${PYTEST} ./test/tensorflow_basics_test.py && ${MPIRUN} -np 4 ${PYTEST} ./test/tensorflow_basics_test.py
+
 clean_build:
 	rm -R build
 

test/tensorflow_basics_test.py

@@ -0,0 +1,63 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import unittest
+import warnings
+import numpy as np
+import networkx as nx
+import pytest
+import tensorflow as tf
+
+from common import mpi_env_rank_and_size
+import bluefog.tensorflow as bf
+from bluefog.common.topology_util import PowerTwoRingGraph, BiRingGraph
+
+warnings.filterwarnings("ignore", message="numpy.dtype size changed")
+warnings.filterwarnings("ignore", message="numpy.ufunc size changed")
+
+
+class BasicsTests(tf.test.TestCase):
+    """
+    Tests for basics.py
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(BasicsTests, self).__init__(*args, **kwargs)
+        warnings.simplefilter("module")
+
+    def test_bluefog_rank(self):
+        """Test that the rank returned by bf.rank() is correct."""
+        true_rank, _ = mpi_env_rank_and_size()
+        bf.init()
+        rank = bf.rank()
+        # print("Rank: ", true_rank, rank)
+        assert true_rank == rank
+
+    def test_bluefog_size(self):
+        """Test that the size returned by bf.size() is correct."""
+        _, true_size = mpi_env_rank_and_size()
+        bf.init()
+        size = bf.size()
+        # print("Size: ", true_size, size)
+        assert true_size == size
+
+    def test_set_and_load_topology(self):
+        _, size = mpi_env_rank_and_size()
+        if size == 4:
+            expected_topology = nx.DiGraph(np.array(
+                [[0, 1, 1, 0], [0, 0, 1, 1], [1, 0, 0, 1], [1, 1, 0, 0]]
+            ))
+        elif size == 1:
+            expected_topology = nx.DiGraph(np.array([[0]]))
+        else:
+            expected_topology = PowerTwoRingGraph(size)
+        bf.init()
+        _, _, topology = bf.load_topology()
+        assert isinstance(topology, nx.DiGraph)
+        np.testing.assert_array_equal(
+            nx.to_numpy_array(expected_topology), nx.to_numpy_array(topology))
+
+
+if __name__ == "__main__":
+    tf.test.main()

test/torch_basics_test.py

@@ -15,6 +15,7 @@
 warnings.filterwarnings("ignore", message="numpy.dtype size changed")
 warnings.filterwarnings("ignore", message="numpy.ufunc size changed")
 
+
 class BasicsTests(unittest.TestCase):
     """
     Tests for basics.py