diff --git a/oneflow/api/python/env/env.cpp b/oneflow/api/python/env/env.cpp index a6d0b7e0e26cfadaf56720444003b5706b04e487..124ed6e7836084be5c37dba471090b3abde8f509 100644 --- a/oneflow/api/python/env/env.cpp +++ b/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); } diff --git a/oneflow/api/python/env/env.h b/oneflow/api/python/env/env.h index a8e96ce78c761a6e6ada9926495403f0114daba9..26e3d78446b34659a63b5a7c0c35ed46a90e5727 100644 --- a/oneflow/api/python/env/env.h +++ b/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_ diff --git a/oneflow/api/python/env/env_api.h b/oneflow/api/python/env/env_api.h index dfecc0bb49905c8f6d4b31cdd03a332cfaccbc85..fcc5953c800ded7e454acc954fffc38ac5505710 100644 --- a/oneflow/api/python/env/env_api.h +++ b/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_