print table stat info for pslib (#21296)

* print table stat
test=develop

* notes
test=develop

* notes
test=develop

paddle/fluid/framework/fleet/fleet_wrapper.cc

@@ -560,6 +560,19 @@ void FleetWrapper::SaveModel(const std::string& path, const int mode) {
 #endif
 }
 
+void FleetWrapper::PrintTableStat(const uint64_t table_id) {
+#ifdef PADDLE_WITH_PSLIB
+  auto ret = pslib_ptr_->_worker_ptr->print_table_stat(table_id);
+  ret.wait();
+  int32_t err_code = ret.get();
+  if (err_code == -1) {
+    LOG(ERROR) << "print table stat failed";
+  }
+#else
+  VLOG(0) << "FleetWrapper::PrintTableStat does nothing when no pslib";
+#endif
+}
+
 double FleetWrapper::GetCacheThreshold(int table_id) {
 #ifdef PADDLE_WITH_PSLIB
   double cache_threshold = 0.0;

paddle/fluid/framework/fleet/fleet_wrapper.h

@@ -167,6 +167,8 @@ class FleetWrapper {
                            std::string model_path, std::string model_proto_file,
                            std::vector<std::string> table_var_list,
                            bool load_combine);
+
+  void PrintTableStat(const uint64_t table_id);
   // mode = 0, load all feature
   // mode = 1, laod delta feature, which means load diff
   void LoadModel(const std::string& path, const int mode);

paddle/fluid/pybind/fleet_wrapper_py.cc

@@ -63,6 +63,7 @@ void BindFleetWrapper(py::module* m) {
            &framework::FleetWrapper::CreateClient2ClientConnection)
       .def("shrink_sparse_table", &framework::FleetWrapper::ShrinkSparseTable)
       .def("shrink_dense_table", &framework::FleetWrapper::ShrinkDenseTable)
+      .def("print_table_stat", &framework::FleetWrapper::PrintTableStat)
       .def("client_flush", &framework::FleetWrapper::ClientFlush)
       .def("load_from_paddle_model",
            &framework::FleetWrapper::LoadFromPaddleModel)

python/paddle/fluid/incubate/fleet/parameter_server/pslib/__init__.py

@@ -238,6 +238,25 @@ class PSLib(Fleet):
         """
         self._fleet_ptr.save_model(dirname)
 
+    def print_table_stat(self, table_id):
+        """
+        print stat info of table_id,
+        format: tableid, feasign size, mf size
+
+        Args:
+            table_id(int): the id of table
+
+        Example:
+            .. code-block:: python
+
+              fleet.print_table_stat(0)
+
+        """
+        self._role_maker._barrier_worker()
+        if self._role_maker.is_first_worker():
+            self._fleet_ptr.print_table_stat(table_id)
+        self._role_maker._barrier_worker()
+
     def save_persistables(self, executor, dirname, main_program=None, **kwargs):
         """
         save presistable parameters,