【HeterPS】fix hdfs and fleet_util for supporting save/load/infer (#33903)

* fix hdfs and fleet_util for supporting save/load infer;test=develop

Author: danleifeng

paddle/fluid/framework/device_worker_factory.cc

@@ -69,11 +69,6 @@ REGISTER_DEVICE_WORKER_CLASS(DownpourWorkerOpt);
 REGISTER_DEVICE_WORKER_CLASS(HeterCpuWorker);
 #endif
 
-#if (defined PADDLE_WITH_NCCL || defined PADDLE_WITH_RCCL) && \
-    (defined PADDLE_WITH_PSLIB)
-REGISTER_DEVICE_WORKER_CLASS(HeterBoxWorker);
-#endif
-
 #if (defined PADDLE_WITH_NCCL || defined PADDLE_WITH_RCCL) && \
     (defined PADDLE_WITH_PSLIB)
 REGISTER_DEVICE_WORKER_CLASS(PSGPUWorker);

paddle/fluid/framework/fleet/heter_ps/CMakeLists.txt

@@ -8,11 +8,11 @@ IF(WITH_GPU)
         SET(HETERPS_DEPS ${HETERPS_DEPS} ${RPC_DEPS})
     endif()
     nv_library(heter_comm SRCS heter_comm.h feature_value.h heter_resource.cc heter_resource.h hashtable.h DEPS ${HETERPS_DEPS})
-    nv_test(test_heter_comm SRCS test_heter_comm.cu feature_value.h DEPS heter_comm)
+    nv_test(test_heter_comm SRCS feature_value.h DEPS heter_comm)
     nv_library(heter_ps SRCS heter_ps.cu DEPS heter_comm)
 ENDIF()
 IF(WITH_ROCM)
     hip_library(heter_comm SRCS heter_comm.h feature_value.h heter_resource.cc heter_resource.h hashtable.h DEPS cub device_context)
-    hip_test(test_heter_comm SRCS test_heter_comm.cu feature_value.h DEPS heter_comm)
+    hip_test(test_heter_comm SRCS feature_value.h DEPS heter_comm)
     hip_library(heter_ps SRCS heter_ps.cu DEPS heter_comm)
 ENDIF()

paddle/fluid/framework/fleet/heter_ps/cudf/concurrent_unordered_map.cuh.h

@@ -765,7 +765,7 @@ x.second );
   unsigned long long get_num_collisions() const { return m_collisions; }
 
   void print() {
-    for (size_type i = 0; i < 10; ++i) {
+    for (size_type i = 0; i < 5; ++i) {
       std::cout << i << ": " << m_hashtbl_values[i].first << ","
                 << m_hashtbl_values[i].second << std::endl;
     }

paddle/fluid/framework/fleet/heter_ps/heter_comm_inl.h

@@ -115,7 +115,7 @@ void HeterComm<KeyType, ValType, GradType>::init_path() {
   path_.resize(total_gpu);
 
   if (!topo_aware_) {
-    VLOG(1) << "init path without topo aware";
+    VLOG(3) << "init path without topo aware";
     for (int i = 0; i < total_gpu; ++i) {
       path_[i].resize(total_gpu);
       for (int j = 0; j < total_gpu; ++j) {
@@ -130,7 +130,7 @@ void HeterComm<KeyType, ValType, GradType>::init_path() {
       }
     }
   } else {
-    VLOG(1) << "init path with topo aware";
+    VLOG(3) << "init path with topo aware";
     for (int i = 0; i < total_gpu; ++i) {
       path_[i].resize(total_gpu);
       for (int j = 0; j < total_gpu; ++j) {

paddle/fluid/framework/fleet/ps_gpu_wrapper.cc

@@ -68,8 +68,6 @@ void PSGPUWrapper::BuildTask(std::shared_ptr<HeterContext> gpu_task,
   thread_keys_.resize(thread_keys_thread_num_);
   for (int i = 0; i < thread_keys_thread_num_; i++) {
     thread_keys_[i].resize(thread_keys_shard_num_);
-    for (int j = 0; j < thread_keys_shard_num_; j++) {
-    }
   }
   const std::deque<Record>& vec_data = input_channel->GetData();
   size_t total_len = vec_data.size();
@@ -255,7 +253,7 @@ void PSGPUWrapper::BuildTask(std::shared_ptr<HeterContext> gpu_task,
         }
       }
 #endif
-      VLOG(1) << "GpuPs build hbmps done";
+      VLOG(3) << "GpuPs build hbmps done";
 
       device_mutex[dev]->unlock();
     }
@@ -295,7 +293,7 @@ void PSGPUWrapper::BuildGPUPS(uint64_t table_id, int feature_dim) {
   HeterPs_ = HeterPsBase::get_instance(size_max, resource_);
   HeterPs_->set_nccl_comm_and_size(inner_comms_, inter_comms_, node_size_);
   auto build_func = [this, &gpu_task, &feature_keys_count](int i) {
-    std::cout << "building table: " << i << std::endl;
+    VLOG(3) << "building table: " << i;
     this->HeterPs_->build_ps(i, gpu_task->device_keys_[i].data(),
                              gpu_task->device_values_[i].data(),
                              feature_keys_count[i], 500000, 2);

python/paddle/distributed/fleet/utils/fs.py

@@ -111,6 +111,10 @@ def list_dirs(self, fs_path):
     def touch(self, fs_path, exist_ok=True):
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def cat(self, fs_path=None):
+        raise NotImplementedError
+
 
 class LocalFS(FS):
     """
@@ -676,14 +680,35 @@ def is_exist(self, fs_path):
 
         return True
 
+    def upload_dir(self, local_dir, dest_dir, overwrite=False):
+        """
+        upload dir to hdfs
+        Args:
+            local_dir(str): local dir
+            dest_dir(str): hdfs dest dir
+            overwrite(bool): is overwrite
+        Returns:
+            return code
+        """
+        local_dir = local_dir.rstrip("/")
+        dest_dir = dest_dir.rstrip("/")
+        local_basename = os.path.basename(local_dir)
+        if self.is_exist(dest_dir + "/" + local_basename) and overwrite:
+            self.delete(dest_dir + "/" + local_basename)
+        if not self.is_exist(dest_dir):
+            self.mkdirs(dest_dir)
+        self._try_upload(local_dir, dest_dir)
+
     # can't retry
-    def upload(self, local_path, fs_path):
+    def upload(self, local_path, fs_path, multi_processes=1, overwrite=False):
         """
         Upload the local path to remote HDFS.
 
         Args:
             local_path(str): The local path.
             fs_path(str): The HDFS path.
+            multi_processes(int|1): the upload data process at the same time, default=5
+            overwrite(bool|False): will overwrite file on HDFS or not
 
         Examples:
 
@@ -700,35 +725,83 @@ def upload(self, local_path, fs_path):
                 client = HDFSClient(hadoop_home, configs)
                 client.upload("test_hdfs_client", "hdfs:/test_hdfs_client")
         """
-        if self.is_exist(fs_path):
-            raise FSFileExistsError("{} exists".format(fs_path))
+
+        def __subprocess_upload(hdfs_path_single, datas):
+            for data in datas:
+                self._try_upload(data, hdfs_path_single)
+
+        def get_local_files(path):
+            """
+            get local files
+            Args:
+                path(str): local path
+            Returns:
+                list of local files
+            """
+            rlist = []
+
+            if not os.path.exists(path):
+                return rlist
+
+            if os.path.isdir(path):
+                for file in os.listdir(path):
+                    t = os.path.join(path, file)
+                    rlist.append(t)
+            else:
+                rlist.append(path)
+            return rlist
 
         local = LocalFS()
         if not local.is_exist(local_path):
             raise FSFileNotExistsError("{} not exists".format(local_path))
+        # upload_dir
+        if local.is_dir(local_path):
+            self.upload_dir(local_path, fs_path, overwrite=overwrite)
+            return
+        # upload files
+        all_files = get_local_files(local_path)
+        if not all_files:
+            print("there are nothing need to upload, function exit")
+            return
+
+        if self.is_exist(fs_path) and overwrite:
+            self.delete(fs_path)
+            self.mkdirs(fs_path)
+
+        procs = []
+        for i in range(multi_processes):
+            process_datas = self._split_files(all_files, i, multi_processes)
+            p = multiprocessing.Process(
+                target=__subprocess_upload, args=(fs_path, process_datas))
+            procs.append(p)
+            p.start()
 
-        return self._try_upload(local_path, fs_path)
+        # complete the processes
+        for proc in procs:
+            proc.join()
 
     @_handle_errors()
     def _try_upload(self, local_path, fs_path):
         cmd = "put {} {}".format(local_path, fs_path)
         ret = 0
         try:
-            ret, lines = self._run_cmd(cmd)
+            ret, _ = self._run_cmd(cmd)
             if ret != 0:
                 raise ExecuteError(cmd)
         except Exception as e:
             self.delete(fs_path)
             raise e
 
     # can't retry
-    def download(self, fs_path, local_path):
+    def download(self, fs_path, local_path, multi_processes=1, overwrite=False):
         """
         Download remote HDFS path to the local.
 
         Args:
             fs_path(str):  The HDFS path.
             local_path(str): The local path.
+            multi_processes(int|1): the download data process at the same time, default=1
+            overwrite(bool): is overwrite
 
         Examples:
 
@@ -745,17 +818,43 @@ def download(self, fs_path, local_path):
                 client = HDFSClient(hadoop_home, configs)
                 client.download("hdfs:/test_hdfs_client", "./")
         """
+
+        def __subprocess_download(local_path, datas):
+            """
+            download file from HDFS
+            Args:
+                local_path(str): the local file path
+                datas(str): the hdfs file path list
+            """
+            for data in datas:
+                self._try_download(data, local_path)
+
         if not self.is_exist(fs_path):
             raise FSFileNotExistsError("{} not exits".format(fs_path))
-
-        return self._try_download(fs_path, local_path)
+        # download file
+        if self.is_file(fs_path):
+            return self._try_download(fs_path, local_path)
+        # download dir
+        _, all_files = self.ls_dir(fs_path)
+
+        procs = []
+        for i in range(multi_processes):
+            process_datas = self._split_files(all_files, i, multi_processes)
+            p = multiprocessing.Process(
+                target=__subprocess_download, args=(local_path, process_datas))
+            procs.append(p)
+            p.start()
+
+        # complete the processes
+        for proc in procs:
+            proc.join()
 
     @_handle_errors()
     def _try_download(self, fs_path, local_path):
         cmd = "get {} {}".format(fs_path, local_path)
         ret = 0
         try:
-            ret, lines = self._run_cmd(cmd)
+            ret, _ = self._run_cmd(cmd)
             if ret != 0:
                 raise ExecuteError(cmd)
         except Exception as e:
@@ -803,7 +902,7 @@ def mkdirs(self, fs_path):
 
         if out_hdfs and not self.is_exist(fs_path):
             cmd = "mkdir -p {}".format(fs_path)
-            ret, lines = self._run_cmd(cmd)
+            ret, _ = self._run_cmd(cmd)
             if ret != 0:
                 raise ExecuteError(cmd)
 
@@ -939,7 +1038,71 @@ def _touchz(self, fs_path):
         cmd = "touchz {}".format(fs_path)
         ret, _ = self._run_cmd(cmd)
         if ret != 0:
-            raise ExecuteError
+            raise ExecuteError(cmd)
 
     def need_upload_download(self):
         return True
+
+    def cat(self, fs_path=None):
+        """
+        Cat a remote HDFS file.
+
+        Args:
+            fs_path(str): The HDFS file path.
+
+        Returns:
+            file content
+
+        Examples:
+
+            .. code-block:: text
+
+                from paddle.distributed.fleet.utils import HDFSClient
+
+                hadoop_home = "/home/client/hadoop-client/hadoop/"
+                configs = {
+                    "fs.default.name": "hdfs://xxx.hadoop.com:54310",
+                    "hadoop.job.ugi": "hello,hello123"
+                }
+
+                client = HDFSClient(hadoop_home, configs)
+                client.cat("hdfs:/test_hdfs_client")
+        """
+        if self.is_file(fs_path):
+            output = self._try_cat(fs_path)
+            return "\n".join(output)
+        else:
+            return ""
+
+    @_handle_errors()
+    def _try_cat(self, fs_path):
+        cmd = "cat {}".format(fs_path)
+        ret, output = self._run_cmd(cmd)
+        if ret != 0:
+            raise ExecuteError(cmd)
+        return output
+
+    def _split_files(self, files, trainer_id, trainers):
+        """
+        split file list
+        Args:
+            files(list): file list
+            trainer_id(int): trainer mpi rank id
+            trainers(int): all trainers num
+        Returns:
+            fileist(list): file list of current trainer
+        """
+        remainder = len(files) % trainers
+        blocksize = len(files) // trainers
+
+        blocks = [blocksize] * trainers
+        for i in range(remainder):
+            blocks[i] += 1
+
+        trainer_files = [[]] * trainers
+        begin = 0
+        for i in range(trainers):
+            trainer_files[i] = files[begin:begin + blocks[i]]
+            begin += blocks[i]
+
+        return trainer_files[trainer_id]

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

@@ -13,6 +13,8 @@
 """Defination of Server and Worker."""
 
 from . import ps_pb2 as pslib
+# NOTE: reduce removed in fuctools in python3
+from functools import reduce
 
 
 class Server(object):