From 006b0ef7b0a7c022bdb0980168546cfa79b6ded9 Mon Sep 17 00:00:00 2001
From: Yurui Li <32978179+poohRui@users.noreply.github.com>
Date: Thu, 6 May 2021 12:39:15 +0800
Subject: [PATCH] Add GetRank and GetWordSize interface (#4809)
* add GetRank and GetWordSize interface
* minor fix
---
oneflow/api/python/env/env.cpp | 3 +++
oneflow/api/python/env/env.h | 3 +++
oneflow/api/python/env/env_api.h | 4 ++++
3 files changed, 10 insertions(+)
diff --git a/oneflow/api/python/env/env.cpp b/oneflow/api/python/env/env.cpp
index a6d0b7e0e..124ed6e78 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 a8e96ce78..26e3d7844 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 dfecc0bb4..fcc5953c8 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_
--
GitLab