add clear one table (#23089)

* add clear_one_table
* test=develop

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) {

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

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)

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.

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: