From 68ea1ad55b7f80ec8cb0271615020d75f26174bc Mon Sep 17 00:00:00 2001
From: xujiaqi01 <173596896@qq.com>
Date: Thu, 26 Mar 2020 18:27:58 +0800
Subject: [PATCH] add clear one table (#23089)

* add clear_one_table
* test=develop
---
 paddle/fluid/framework/fleet/fleet_wrapper.cc   |  9 +++++++++
 paddle/fluid/framework/fleet/fleet_wrapper.h    |  2 ++
 paddle/fluid/pybind/fleet_wrapper_py.cc         |  1 +
 .../fleet/parameter_server/pslib/__init__.py    | 17 +++++++++++++++++
 .../tests/unittests/test_fleet_rolemaker.py     |  1 +
 5 files changed, 30 insertions(+)

diff --git a/paddle/fluid/framework/fleet/fleet_wrapper.cc b/paddle/fluid/framework/fleet/fleet_wrapper.cc
index fc52e1a4c93..371a5507f1f 100644
--- a/paddle/fluid/framework/fleet/fleet_wrapper.cc
+++ b/paddle/fluid/framework/fleet/fleet_wrapper.cc
@@ -804,6 +804,15 @@ void FleetWrapper::ClearModel() {
 #endif
 }
 
+void FleetWrapper::ClearOneTable(const uint64_t table_id) {
+#ifdef PADDLE_WITH_PSLIB
+  auto ret = pslib_ptr_->_worker_ptr->clear(table_id);
+  ret.wait();
+#else
+  VLOG(0) << "FleetWrapper::ClearOneTable does nothing when no pslib";
+#endif
+}
+
 void FleetWrapper::ShrinkDenseTable(int table_id, Scope* scope,
                                     std::vector<std::string> var_list,
                                     float decay, int emb_dim) {
diff --git a/paddle/fluid/framework/fleet/fleet_wrapper.h b/paddle/fluid/framework/fleet/fleet_wrapper.h
index 5d831f31c7f..a54aea034d2 100644
--- a/paddle/fluid/framework/fleet/fleet_wrapper.h
+++ b/paddle/fluid/framework/fleet/fleet_wrapper.h
@@ -220,6 +220,8 @@ class FleetWrapper {
                              const std::vector<uint64_t>& feasign_list);
   // clear all models, release their memory
   void ClearModel();
+  // clear one table
+  void ClearOneTable(const uint64_t table_id);
   // shrink sparse table
   void ShrinkSparseTable(int table_id);
   // shrink dense table
diff --git a/paddle/fluid/pybind/fleet_wrapper_py.cc b/paddle/fluid/pybind/fleet_wrapper_py.cc
index fac6de452ae..3b4505c611b 100644
--- a/paddle/fluid/pybind/fleet_wrapper_py.cc
+++ b/paddle/fluid/pybind/fleet_wrapper_py.cc
@@ -59,6 +59,7 @@ void BindFleetWrapper(py::module* m) {
       .def("save_cache", &framework::FleetWrapper::SaveCache)
       .def("load_model", &framework::FleetWrapper::LoadModel)
       .def("clear_model", &framework::FleetWrapper::ClearModel)
+      .def("clear_one_table", &framework::FleetWrapper::ClearOneTable)
       .def("stop_server", &framework::FleetWrapper::StopServer)
       .def("finalize_worker", &framework::FleetWrapper::FinalizeWorker)
       .def("gather_servers", &framework::FleetWrapper::GatherServers)
diff --git a/python/paddle/fluid/incubate/fleet/parameter_server/pslib/__init__.py b/python/paddle/fluid/incubate/fleet/parameter_server/pslib/__init__.py
index 40337110cfe..27fecf495ad 100644
--- a/python/paddle/fluid/incubate/fleet/parameter_server/pslib/__init__.py
+++ b/python/paddle/fluid/incubate/fleet/parameter_server/pslib/__init__.py
@@ -402,6 +402,23 @@ class PSLib(Fleet):
                                                        var_list, decay, emb_dim)
         self._role_maker._barrier_worker()
 
+    def clear_one_table(self, table_id):
+        """
+        clear_one_table() will be called by user. It will clear one table.
+
+        Args:
+            table_id(int): table id
+
+        Examples:
+            .. code-block:: python
+
+              fleet.clear_one_table(0)
+        """
+        self._role_maker._barrier_worker()
+        if self._role_maker.is_first_worker():
+            self._fleet_ptr.clear_one_table(table_id)
+        self._role_maker._barrier_worker()
+
     def clear_model(self):
         """
         clear_model() will be called by user. It will clear sparse model.
diff --git a/python/paddle/fluid/tests/unittests/test_fleet_rolemaker.py b/python/paddle/fluid/tests/unittests/test_fleet_rolemaker.py
index 7322891338f..47aeee95921 100644
--- a/python/paddle/fluid/tests/unittests/test_fleet_rolemaker.py
+++ b/python/paddle/fluid/tests/unittests/test_fleet_rolemaker.py
@@ -99,6 +99,7 @@ class TestCloudRoleMaker(unittest.TestCase):
         except:
             print("do not support pslib test, skip")
             return
+        fleet.clear_one_table(0)
         from paddle.fluid.incubate.fleet.base.role_maker import \
             MPISymetricRoleMaker
         try:
-- 
GitLab