Move framework.distribute to env

python/oneflow/__init__.py

@@ -125,9 +125,9 @@ def is_normal_exit(self):
 def atexit_hook(hook):
     if hook.is_normal_exit():
         if oneflow._oneflow_internal.IsEnvInited():
-            if oneflow.framework.distribute.is_multi_client():
+            if oneflow.env.is_multi_client():
                 oneflow._oneflow_internal.eager.multi_client.Sync()
-            elif oneflow.framework.distribute.get_rank() == 0:
+            elif oneflow.env.get_rank() == 0:
                 oneflow._oneflow_internal.eager.single_client.Sync()
     oneflow.framework.session_context.TryCloseDefaultSession()
     if hook.is_normal_exit():

python/oneflow/distributed/__init__.py

@@ -13,9 +13,3 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from oneflow.framework.distribute import (
-    get_local_rank,
-    get_rank,
-    get_world_size,
-    is_multi_client,
-)

python/oneflow/env.py

@@ -23,3 +23,33 @@
 from oneflow.framework.env_util import api_logbuflevel as logbuflevel
 from oneflow.framework.env_util import api_logtostderr as logtostderr
 from oneflow.framework.env_util import api_machine as machine
+
+import oneflow._oneflow_internal
+
+
+def get_local_rank():
+    return oneflow._oneflow_internal.GetLocalRank()
+
+
+def get_rank():
+    """Returns the rank of current process group.
+
+    Returns:
+        The rank of the process group.
+
+    """
+    return oneflow._oneflow_internal.GetRank()
+
+
+def get_world_size():
+    """Returns the number of processes in the current process group.
+
+    Returns:
+        The world size of the process group.
+
+    """
+    return oneflow._oneflow_internal.GetWorldSize()
+
+
+def is_multi_client():
+    return oneflow._oneflow_internal.IsMultiClient()

python/oneflow/framework/check_point_v2.py

@@ -105,7 +105,7 @@ def _LoadSingleVariable(
     path: Optional[str], consistent_src_rank: Optional[int] = None
 ) -> "flow.Tensor":
     if consistent_src_rank is not None:
-        rank = flow.framework.distribute.get_rank()
+        rank = flow.env.get_rank()
         if rank == consistent_src_rank:
             assert isinstance(path, str)
             file_backed_blob = FileBackendVariableBlob(path)
@@ -124,7 +124,7 @@ def _LoadSingleVariable(
 
 
 def _broadcast_py_object(obj, src: int = 0):
-    rank = flow.framework.distribute.get_rank()
+    rank = flow.env.get_rank()
     if src == rank:
         obj_bytes = pickle.dumps(obj)
         return pickle.loads(flow._oneflow_internal.cpu_broadcast(obj_bytes, src))
@@ -136,7 +136,7 @@ def Load(
     path: str, consistent_src_rank: Optional[int] = None,
 ) -> Dict[str, "flow.Tensor"]:
     assert os.path.isdir(path), "Directory {} doesn't exist!".format(path)
-    rank = flow.framework.distribute.get_rank()
+    rank = flow.env.get_rank()
     var_dict = {}
     if consistent_src_rank is None or rank == consistent_src_rank:
         all_files = os.listdir(path)
@@ -169,7 +169,7 @@ def save(
                 not var.is_consistent
             ), f"local tensor is needed, but {name} is a consistent tensor"
 
-    rank = flow.framework.distribute.get_rank()
+    rank = flow.env.get_rank()
     if consistent_mode and rank != consistent_dst_rank:
         return
 

python/oneflow/framework/distribute.py

@@ -19,34 +19,6 @@
 import oneflow._oneflow_internal
 
 
-def get_local_rank():
-    return oneflow._oneflow_internal.GetLocalRank()
-
-
-def get_rank():
-    """Returns the rank of current process group.
-
-    Returns:
-        The rank of the process group.
-
-    """
-    return oneflow._oneflow_internal.GetRank()
-
-
-def get_world_size():
-    """Returns the number of processes in the current process group.
-
-    Returns:
-        The world size of the process group.
-
-    """
-    return oneflow._oneflow_internal.GetWorldSize()
-
-
-def is_multi_client():
-    return oneflow._oneflow_internal.IsMultiClient()
-
-
 def split_sbp(axis: int) -> oneflow._oneflow_internal.sbp.sbp:
     """Generate a split scheme in which op will be splitted at `axis`.
 

python/oneflow/framework/unittest.py

@@ -113,7 +113,7 @@ def node_size():
 
 
 def has_world_size():
-    if oneflow.distributed.is_multi_client():
+    if oneflow.env.is_multi_client():
         return True
     if os.getenv("ONEFLOW_TEST_WORLD_SIZE"):
         assert os.getenv(
@@ -125,8 +125,8 @@ def has_world_size():
 
 
 def world_size():
-    if oneflow.distributed.is_multi_client():
-        return oneflow.distributed.get_world_size()
+    if oneflow.env.is_multi_client():
+        return oneflow.env.get_world_size()
     return int(os.getenv("ONEFLOW_TEST_WORLD_SIZE"))
 
 

python/oneflow/nn/graph/block.py

@@ -19,7 +19,7 @@
 
 import oneflow._oneflow_internal
 import oneflow.framework.graph_build_util as graph_build_util
-from oneflow.framework.distribute import get_rank
+from oneflow.env import get_rank
 from oneflow.framework.tensor import Tensor, TensorTuple
 from oneflow.nn.module import Module
 from oneflow.nn.parameter import Parameter

python/oneflow/nn/graph/graph.py

@@ -21,7 +21,7 @@
 import oneflow.framework.c_api_util as c_api_util
 import oneflow.framework.graph_build_util as graph_build_util
 import oneflow.framework.session_context as session_ctx
-from oneflow.framework.distribute import get_rank
+from oneflow.env import get_rank
 from oneflow.framework.tensor import Tensor, TensorTuple
 from oneflow.framework.multi_client_session import MultiClientSession
 from oneflow.framework.tensor_tuple_util import convert_to_tensor_tuple

python/oneflow/nn/modules/all_reduce.py

@@ -32,5 +32,5 @@ def __init__(self, parallel_conf_str: str):
 
     def forward(self, x):
         assert x.device.type == "cuda"
-        assert x.device.index == flow.framework.distribute.get_local_rank()
+        assert x.device.index == flow.env.get_local_rank()
         return self._op(x)[0]

python/oneflow/nn/parallel/ddp.py

@@ -43,7 +43,7 @@ def allreduce(grad):
 def DistributedDataParallel(
     module: "flow.nn.Module", *, broadcast_buffers: bool = True
 ):
-    world_size = flow.distributed.get_world_size()
+    world_size = flow.env.get_world_size()
     with flow.no_grad():
         for x in module.parameters():
             requires_grad = x.requires_grad

python/oneflow/test/graph/test_graph_asymmetric_io.py

@@ -55,8 +55,8 @@ def __init__(self):
                 flow.nn.init.constant_(self.linear2.weight, 2.3)
 
             def forward(self, x, y):
-                # print("local_x in rank : ", flow.distributed.get_rank(), " is : ", x)
-                # print("local_y in rank : ", flow.distributed.get_rank(), " is : ", y)
+                # print("local_x in rank : ", flow.env.get_rank(), " is : ", x)
+                # print("local_y in rank : ", flow.env.get_rank(), " is : ", y)
                 out0 = x + y
                 out1 = self.linear1(out0)
                 out2 = self.linear2(out1)
@@ -92,11 +92,11 @@ def build(self, x, y):
         graph_local_out = graph_out.to_local()
         # NOTE(chengcheng): MUST call for each rank sync correct input copy
         graph_local_out_np = graph_local_out.numpy()
-        # print("graph_local_out in rank ", flow.distributed.get_rank(),  " is : ", graph_local_out)
-        if flow.distributed.get_rank() == 0:
+        # print("graph_local_out in rank ", flow.env.get_rank(),  " is : ", graph_local_out)
+        if flow.env.get_rank() == 0:
             test_case.assertTrue(graph_local_out.shape.numel() == 0)
             test_case.assertTrue(graph_local_out_np.size == np.array([]).size)
-        elif flow.distributed.get_rank() == 1:
+        elif flow.env.get_rank() == 1:
             test_case.assertTrue(
                 np.allclose(
                     graph_local_out.numpy(), local_out.numpy(), atol=1e-4, rtol=1e-4

python/oneflow/test/graph/test_input_op_expr.py

@@ -31,7 +31,7 @@
 @flow.unittest.skip_unless_1n1d()
 class TestFeedInputTensor(unittest.TestCase):
     def test_feed_input_tensor(test_case):
-        test_case.assertTrue(oneflow.distributed.is_multi_client())
+        test_case.assertTrue(oneflow.env.is_multi_client())
         test_case.assertTrue(oneflow.framework.env_util.HasAllMultiClientEnvVars())
         x = flow.Tensor(1, 1, 10, 10)
         flow.nn.init.uniform_(x, a=-1.0, b=1.0)

python/oneflow/test/graph/test_multi_client_session.py

@@ -27,15 +27,15 @@
 @flow.unittest.skip_unless_1n1d()
 class TestMultiClientSession(unittest.TestCase):
     def test_case1(self):
-        self.assertTrue(flow.distributed.is_multi_client())
+        self.assertTrue(flow.env.is_multi_client())
         sess = session_ctx.GetDefaultSession()
         self.assertTrue(isinstance(sess, MultiClientSession))
         sess.TryInit()
         self.assertEqual(sess.status, sess.Status.INITED)
 
     def test_case2(self):
         print("test_case2")
-        self.assertTrue(flow.distributed.is_multi_client())
+        self.assertTrue(flow.env.is_multi_client())
         sess = session_ctx.GetDefaultSession()
         self.assertTrue(isinstance(sess, MultiClientSession))
         sess.TryInit()

python/oneflow/test/graph/test_output_op_expr.py

@@ -31,7 +31,7 @@
 @flow.unittest.skip_unless_1n1d()
 class TestFetchOutputTensor(unittest.TestCase):
     def test_fetch_output_tensor(test_case):
-        test_case.assertTrue(oneflow.distributed.is_multi_client())
+        test_case.assertTrue(oneflow.env.is_multi_client())
         test_case.assertTrue(oneflow.framework.env_util.HasAllMultiClientEnvVars())
         x = flow.Tensor(1, 1, 10, 10)
         flow.nn.init.uniform_(x, a=-1.0, b=1.0)

python/oneflow/test/graph/test_to_consistent.py

@@ -194,7 +194,7 @@ class ToConsistentGraphTestCase(oneflow.unittest.TestCase):
     def test_fwd_P2B(test_case):
         """ compare eager fwd and lazy bwd
         """
-        rank = flow.distributed.get_rank()
+        rank = flow.env.get_rank()
         # pid = os.getpid()
         # print(f"[{pid}][{rank}] ToConsistentGraphTestCase.test_fwd_P2B")
 
@@ -230,7 +230,7 @@ def test_fwd_P2B(test_case):
     def test_bwd_P2B(test_case):
         """ compare eager bwd and lazy bwd
         """
-        rank = flow.distributed.get_rank()
+        rank = flow.env.get_rank()
         # pid = os.getpid()
         # print(f"[{pid}][{rank}] ToConsistentGraphTestCase.test_bwd_P2B")
 
@@ -275,7 +275,7 @@ def test_bwd_P2B(test_case):
     def test_multi_graph(test_case):
         """ compare two lazy fwd
         """
-        rank = flow.distributed.get_rank()
+        rank = flow.env.get_rank()
         # pid = os.getpid()
         # print(f"[{pid}][{rank}] ToConsistentGraphTestCase.test_multi_graph")
 
@@ -370,7 +370,7 @@ def test_free_tensor_to_consistent(test_case):
 
     # @unittest.skipIf(True, "")
     def test_to_placement(test_case):
-        rank = flow.distributed.get_rank()
+        rank = flow.env.get_rank()
         # pid = os.getpid()
         # print(f"[{pid}][{rank}] ToConsistentGraphTestCase.test_to_placement")
 

python/oneflow/test/graph/test_user_op_expr.py

@@ -35,7 +35,7 @@ def _get_c_tensor(t):
 
 
 def _test_user_op_graph(test_case, is_cuda):
-    test_case.assertTrue(oneflow.distributed.is_multi_client())
+    test_case.assertTrue(oneflow.env.is_multi_client())
     test_case.assertTrue(oneflow.framework.env_util.HasAllMultiClientEnvVars())
 
     x0 = flow.tensor(np.random.rand(20, 30), dtype=flow.float32)

python/oneflow/test/graph/test_variable_op_expr.py

@@ -31,7 +31,7 @@
 @flow.unittest.skip_unless_1n1d()
 class TestFeedVariableTensor(unittest.TestCase):
     def test_feed_var_tensor(test_case):
-        test_case.assertTrue(oneflow.distributed.is_multi_client())
+        test_case.assertTrue(oneflow.env.is_multi_client())
         test_case.assertTrue(oneflow.framework.env_util.HasAllMultiClientEnvVars())
         x = flow.Tensor(1, 1, 10, 10)
         flow.nn.init.uniform_(x, a=-1.0, b=1.0)

python/oneflow/test/modules/test_allreduce.py

@@ -28,9 +28,9 @@ class TestAllReduce(flow.unittest.TestCase):
     def test_all_reduce(test_case):
         arr_rank1 = np.array([1, 2])
         arr_rank2 = np.array([3, 4])
-        if flow.distributed.get_rank() == 0:
+        if flow.env.get_rank() == 0:
             x = flow.Tensor(arr_rank1)
-        elif flow.distributed.get_rank() == 1:
+        elif flow.env.get_rank() == 1:
             x = flow.Tensor(arr_rank2)
         else:
             raise ValueError
@@ -41,7 +41,7 @@ def test_all_reduce(test_case):
     @flow.unittest.skip_unless_2n2d()
     def test_all_reduce_2nodes(test_case):
         np_arr = np.array([1, 2])
-        x = flow.Tensor(np_arr * (flow.distributed.get_rank() + 1))
+        x = flow.Tensor(np_arr * (flow.env.get_rank() + 1))
         x = x.to("cuda")
         y = flow.F.all_reduce(x)
         test_case.assertTrue(np.allclose(y.numpy(), np_arr * 10))

python/oneflow/test/modules/test_coco_reader.py

@@ -69,7 +69,7 @@ def build(self):
 @flow.unittest.skip_unless_1n2d()
 class COCODataLoaderDistributedTestCase(oneflow.unittest.TestCase):
     def test_case1(test_case):
-        rank = flow.distributed.get_rank()
+        rank = flow.env.get_rank()
         # pid = os.getpid()
         # print(f"[{pid}][{rank}] COCODataLoaderDistributedTestCase.test_case1")