Add GetRank and GetWordSize interface (#4809)

* add GetRank and GetWordSize interface

* minor fix

oneflow/api/python/env/env.cpp

@@ -30,4 +30,7 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
   m.def("DestroyEnv", &DestroyEnv, py::call_guard<py::gil_scoped_release>());
 
   m.def("CurrentMachineId", &CurrentMachineId);
+
+  m.def("GetRank", &GetRank);
+  m.def("GetWordSize", &GetWordSize);
 }

oneflow/api/python/env/env.h

@@ -91,6 +91,9 @@ inline Maybe<void> InitEnv(const std::string& env_proto_str) {
 
 inline Maybe<long long> CurrentMachineId() { return GlobalProcessCtx::Rank(); }
 
+inline Maybe<int64_t> GetRank() { return GlobalProcessCtx::Rank(); }
+inline Maybe<size_t> GetWordSize() { return GlobalProcessCtx::WorldSize(); }
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_API_PYTHON_ENV_ENV_H_

oneflow/api/python/env/env_api.h

@@ -40,4 +40,8 @@ inline void DestroyEnv() { return oneflow::DestroyEnv().GetOrThrow(); }
 
 inline long long CurrentMachineId() { return oneflow::CurrentMachineId().GetOrThrow(); }
 
+inline int64_t GetRank() { return oneflow::GetRank().GetOrThrow(); }
+
+inline size_t GetWordSize() { return oneflow::GetWordSize().GetOrThrow(); }
+
 #endif  // ONEFLOW_API_PYTHON_ENV_ENV_API_H_