From b3bcebbeb1debeae72be94907b45ff8c8df5101d Mon Sep 17 00:00:00 2001
From: Thunderbrook <52529258+Thunderbrook@users.noreply.github.com>
Date: Thu, 7 Apr 2022 23:37:07 +0800
Subject: [PATCH] [GPUPS] bind afs wrpper (#41227)

* afs wrapper

* format

* format

* macro
---
 .../fluid/framework/fleet/ps_gpu_wrapper.cc   | 37 +++++++++++++++++++
 paddle/fluid/framework/fleet/ps_gpu_wrapper.h | 21 +++++++++++
 paddle/fluid/pybind/ps_gpu_wrapper_py.cc      | 21 +++++++++++
 paddle/fluid/pybind/ps_gpu_wrapper_py.h       |  3 ++
 paddle/fluid/pybind/pybind.cc                 |  3 ++
 5 files changed, 85 insertions(+)

diff --git a/paddle/fluid/framework/fleet/ps_gpu_wrapper.cc b/paddle/fluid/framework/fleet/ps_gpu_wrapper.cc
index 75f5c24af5a..c7852de00a1 100755
--- a/paddle/fluid/framework/fleet/ps_gpu_wrapper.cc
+++ b/paddle/fluid/framework/fleet/ps_gpu_wrapper.cc
@@ -37,6 +37,43 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
+#ifdef PADDLE_WITH_PSLIB
+void AfsWrapper::init(const std::string& fs_name, const std::string& fs_user,
+                      const std::string& pass_wd, const std::string& conf) {
+  int ret = afs_handler_.init(fs_name.c_str(), fs_user.c_str(), pass_wd.c_str(),
+                              conf.c_str());
+  if (ret != 0) {
+    LOG(ERROR) << "AFS Init Error";
+  }
+}
+
+int AfsWrapper::remove(const std::string& path) {
+  return afs_handler_.remove(path);
+}
+
+int AfsWrapper::mkdir(const std::string& path) {
+  return afs_handler_.mkdir(path);
+}
+
+std::vector<std::string> AfsWrapper::list(const std::string& path) {
+  return afs_handler_.list(path);
+}
+
+int AfsWrapper::exist(const std::string& path) {
+  return afs_handler_.exist(path);
+}
+
+int AfsWrapper::upload(const std::string& local_file,
+                       const std::string& afs_file) {
+  return afs_handler_.upload_file(local_file, afs_file);
+}
+
+int AfsWrapper::download(const std::string& local_file,
+                         const std::string& afs_file) {
+  return afs_handler_.download_file(local_file, afs_file);
+}
+#endif
+
 std::shared_ptr<PSGPUWrapper> PSGPUWrapper::s_instance_ = NULL;
 bool PSGPUWrapper::is_initialized_ = false;
 #ifdef PADDLE_WITH_PSLIB
diff --git a/paddle/fluid/framework/fleet/ps_gpu_wrapper.h b/paddle/fluid/framework/fleet/ps_gpu_wrapper.h
index d9d29cc072d..9b7d6de082d 100755
--- a/paddle/fluid/framework/fleet/ps_gpu_wrapper.h
+++ b/paddle/fluid/framework/fleet/ps_gpu_wrapper.h
@@ -55,6 +55,27 @@ namespace framework {
 #define TYPEALIGN(ALIGNVAL, LEN) \
   (((uint64_t)(LEN) + ((ALIGNVAL)-1)) & ~((uint64_t)((ALIGNVAL)-1)))
 
+#ifdef PADDLE_WITH_PSLIB
+class AfsWrapper {
+ public:
+  AfsWrapper() {}
+  virtual ~AfsWrapper() {}
+  void init(const std::string& fs_name, const std::string& fs_user,
+            const std::string& pass_wd, const std::string& conf);
+  int remove(const std::string& path);
+  int mkdir(const std::string& path);
+  std::vector<std::string> list(const std::string& path);
+
+  int exist(const std::string& path);
+  int upload(const std::string& local_file, const std::string& afs_file);
+
+  int download(const std::string& local_file, const std::string& afs_file);
+
+ private:
+  paddle::ps::AfsApiWrapper afs_handler_;
+};
+#endif
+
 class PSGPUWrapper {
  public:
   virtual ~PSGPUWrapper() { delete HeterPs_; }
diff --git a/paddle/fluid/pybind/ps_gpu_wrapper_py.cc b/paddle/fluid/pybind/ps_gpu_wrapper_py.cc
index fe1f27226ba..79529fca7d1 100644
--- a/paddle/fluid/pybind/ps_gpu_wrapper_py.cc
+++ b/paddle/fluid/pybind/ps_gpu_wrapper_py.cc
@@ -63,6 +63,27 @@ void BindPSGPUWrapper(py::module* m) {
       .def("finalize", &framework::PSGPUWrapper::Finalize,
            py::call_guard<py::gil_scoped_release>());
 }  // end PSGPUWrapper
+#ifdef PADDLE_WITH_PSLIB
+void BindAfsWrapper(py::module* m) {
+  py::class_<framework::AfsWrapper, std::shared_ptr<framework::AfsWrapper>>(
+      *m, "AfsWrapper")
+      .def(py::init([]() { return std::make_shared<framework::AfsWrapper>(); }))
+      .def("init", &framework::AfsWrapper::init,
+           py::call_guard<py::gil_scoped_release>())
+      .def("list", &framework::AfsWrapper::list,
+           py::call_guard<py::gil_scoped_release>())
+      .def("mkdir", &framework::AfsWrapper::mkdir,
+           py::call_guard<py::gil_scoped_release>())
+      .def("exist", &framework::AfsWrapper::exist,
+           py::call_guard<py::gil_scoped_release>())
+      .def("download", &framework::AfsWrapper::download,
+           py::call_guard<py::gil_scoped_release>())
+      .def("upload", &framework::AfsWrapper::upload,
+           py::call_guard<py::gil_scoped_release>())
+      .def("remove", &framework::AfsWrapper::remove,
+           py::call_guard<py::gil_scoped_release>());
+}
+#endif
 #endif
 }  // end namespace pybind
 }  // end namespace paddle
diff --git a/paddle/fluid/pybind/ps_gpu_wrapper_py.h b/paddle/fluid/pybind/ps_gpu_wrapper_py.h
index ba4f146389e..22cd5ef0fd1 100644
--- a/paddle/fluid/pybind/ps_gpu_wrapper_py.h
+++ b/paddle/fluid/pybind/ps_gpu_wrapper_py.h
@@ -24,6 +24,9 @@ namespace pybind {
 
 #ifdef PADDLE_WITH_HETERPS
 void BindPSGPUWrapper(py::module* m);
+#ifdef PADDLE_WITH_PSLIB
+void BindAfsWrapper(py::module* m);
+#endif
 #endif
 }  // namespace pybind
 }  // namespace paddle
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
index 44abf3357d6..c9e304e696d 100644
--- a/paddle/fluid/pybind/pybind.cc
+++ b/paddle/fluid/pybind/pybind.cc
@@ -4458,6 +4458,9 @@ All parameter, weight, gradient are variables in Paddle.
 #endif
 #ifdef PADDLE_WITH_HETERPS
   BindPSGPUWrapper(&m);
+#ifdef PADDLE_WITH_PSLIB
+  BindAfsWrapper(&m);
+#endif
 #endif
   BindGlooWrapper(&m);
   BindBoxHelper(&m);
-- 
GitLab