chore: import upstream snapshot with attribution

test/ps/CMakeLists.txt

@@ -0,0 +1,11 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP})
+  list(APPEND TEST_OPS ${TEST_OP})
+  set_tests_properties(${TEST_OP} PROPERTIES TIMEOUT 120)
+endforeach()

test/ps/__init__.py

@@ -0,0 +1,17 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.p
+
+# Note: On Windows, import form subdirectories such as dirA()->dirB(), current directory
+# will still be dirA(), But is should be dirB(). So it will ModuleNotFoundError
+# please refer to https://stackoverflow.com/questions/8953844/import-module-from-subfolder

test/ps/config_gpubox.yaml

@@ -0,0 +1,55 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# workspace
+#workspace: "models/rank/dnn"
+
+
+runner:
+  train_data_dir: "data/sample_data/train"
+  train_reader_path: "criteo_reader" # importlib format
+  use_gpu: True
+  use_auc: False
+  train_batch_size: 32
+  epochs: 3
+  print_interval: 10
+  model_save_path: "output_model_dnn_queue"
+
+  sync_mode: "gpubox"
+  thread_num: 30
+  reader_type: "InmemoryDataset"  # DataLoader / QueueDataset / RecDataset / InmemoryDataset
+  pipe_command: "python dataset_generator_criteo.py"
+  dataset_debug: False
+  split_file_list: False
+
+  infer_batch_size: 2
+  infer_reader_path: "criteo_reader" # importlib format
+  test_data_dir: "data/sample_data/train"
+  infer_load_path: "output_model_dnn_queue"
+  infer_start_epoch: 0
+  infer_end_epoch: 3
+# hyper parameters of user-defined network
+hyper_parameters:
+  # optimizer config
+  optimizer:
+    class: Adam
+    learning_rate: 0.001
+    strategy: async
+  # user-defined <key, value> pairs
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1024
+  sparse_feature_dim: 9
+  dense_input_dim: 13
+  fc_sizes: [512, 256, 128, 32]
+  distributed_embedding: 0

test/ps/cpu_async_ps_config.yaml

@@ -0,0 +1,33 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# refer to PaddleRec/models/rank/dnn/benchmark.yaml
+
+hyper_parameters:
+  optimizer:
+    class: Adam
+    learning_rate: 0.0001
+    adam_lazy_mode: True
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1000001
+  sparse_feature_dim: 10
+  dense_input_dim: 13
+  fc_sizes: [400, 400, 400]
+
+runner:
+  sync_mode: "async"  # sync / async / geo / heter
+  thread_num: 16
+  use_gpu: 0
+
+  model_path: "../ps_dnn_model.py"

test/ps/cpu_geo_ps_config.yaml

@@ -0,0 +1,34 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# refer to PaddleRec/models/rank/dnn/benchmark.yaml
+
+hyper_parameters:
+  optimizer:
+    class: Adam
+    learning_rate: 0.0001
+    adam_lazy_mode: True
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1000001
+  sparse_feature_dim: 10
+  dense_input_dim: 13
+  fc_sizes: [400, 400, 400]
+
+runner:
+  geo_step: 400
+  sync_mode: "geo"
+  thread_num: 16
+  use_gpu: 0
+
+  model_path: "../ps_dnn_model.py"

test/ps/cpu_sync_ps_config.yaml

@@ -0,0 +1,33 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# refer to PaddleRec/models/rank/dnn/benchmark.yaml
+
+hyper_parameters:
+  optimizer:
+    class: Adam
+    learning_rate: 0.0001
+    adam_lazy_mode: True
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1000001
+  sparse_feature_dim: 10
+  dense_input_dim: 13
+  fc_sizes: [400, 400, 400]
+
+runner:
+  sync_mode: "sync"
+  thread_num: 16
+  use_gpu: 0
+
+  model_path: "../ps_dnn_model.py"

test/ps/dataset_generator_A.py

@@ -0,0 +1,50 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.distributed import fleet
+
+cont_min_ = [0, -3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+cont_max_ = [20, 600, 100, 50, 64000, 500, 100, 50, 500, 10, 10, 10, 50]
+cont_diff_ = [20, 603, 100, 50, 64000, 500, 100, 50, 500, 10, 10, 10, 50]
+hash_dim_ = 1000001
+continuous_range_ = range(1, 14)
+categorical_range_ = range(14, 40)
+
+
+class CriteoDataset(fleet.MultiSlotDataGenerator):
+    def generate_sample(self, line):
+        """
+        Read the data line by line and process it as a dictionary
+        """
+
+        def reader():
+            """
+            This function needs to be implemented by the user, based on data format
+            """
+            features = line.rstrip('\n').split('\t')
+            feature_name = []
+            sparse_feature = []
+            for idx in categorical_range_:
+                sparse_feature.append(
+                    [hash(str(idx) + features[idx]) % hash_dim_]
+                )
+            for idx in categorical_range_:
+                feature_name.append("C" + str(idx - 13))
+            yield list(zip(feature_name, sparse_feature))
+
+        return reader
+
+
+d = CriteoDataset()
+d.run_from_stdin()

test/ps/dataset_generator_B.py

@@ -0,0 +1,54 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.distributed import fleet
+
+cont_min_ = [0, -3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+cont_max_ = [20, 600, 100, 50, 64000, 500, 100, 50, 500, 10, 10, 10, 50]
+cont_diff_ = [20, 603, 100, 50, 64000, 500, 100, 50, 500, 10, 10, 10, 50]
+hash_dim_ = 1000001
+continuous_range_ = range(1, 14)
+categorical_range_ = range(14, 40)
+
+
+class CriteoDataset(fleet.MultiSlotDataGenerator):
+    def generate_sample(self, line):
+        """
+        Read the data line by line and process it as a dictionary
+        """
+
+        def reader():
+            """
+            This function needs to be implemented by the user, based on data format
+            """
+            features = line.rstrip('\n').split('\t')
+            dense_feature = []
+            for idx in continuous_range_:
+                if features[idx] == "":
+                    dense_feature.append(0.0)
+                else:
+                    dense_feature.append(
+                        (float(features[idx]) - cont_min_[idx - 1])
+                        / cont_diff_[idx - 1]
+                    )
+            label = [int(features[0])]
+            feature_name = ["dense_feature"]
+            feature_name.append("label")
+            yield list(zip(feature_name, [label, dense_feature]))
+
+        return reader
+
+
+d = CriteoDataset()
+d.run_from_stdin()

test/ps/dataset_generator_criteo.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+
+from paddle.distributed import fleet
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+class Reader(fleet.MultiSlotDataGenerator):
+    def init(self):
+        padding = 0
+        sparse_slots = "click 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26"
+        self.sparse_slots = sparse_slots.strip().split(" ")
+        self.dense_slots = ["dense_feature"]
+        self.dense_slots_shape = [13]
+        self.slots = self.sparse_slots + self.dense_slots
+        self.slot2index = {}
+        self.visit = {}
+        for i in range(len(self.slots)):
+            self.slot2index[self.slots[i]] = i
+            self.visit[self.slots[i]] = False
+        self.padding = padding
+        logger.info("pipe init success")
+
+    def line_process(self, line):
+        line = line.strip().split(" ")
+        output = [(i, []) for i in self.slots]
+        for i in line:
+            slot_feasign = i.split(":")
+            slot = slot_feasign[0]
+            if slot not in self.slots:
+                continue
+            if slot in self.sparse_slots:
+                feasign = int(slot_feasign[1])
+            else:
+                feasign = float(slot_feasign[1])
+            output[self.slot2index[slot]][1].append(feasign)
+            self.visit[slot] = True
+        for i in self.visit:
+            slot = i
+            if not self.visit[slot]:
+                if i in self.dense_slots:
+                    output[self.slot2index[i]][1].extend(
+                        [self.padding]
+                        * self.dense_slots_shape[self.slot2index[i]]
+                    )
+                else:
+                    output[self.slot2index[i]][1].extend([self.padding])
+            else:
+                self.visit[slot] = False
+
+        return output
+        # return [label] + sparse_feature + [dense_feature]
+
+    def generate_sample(self, line):
+        r"Dataset Generator"
+
+        def reader():
+            output_dict = self.line_process(line)
+            # {key, value} dict format: {'labels': [1], 'sparse_slot1': [2, 3], 'sparse_slot2': [4, 5, 6, 8], 'dense_slot': [1,2,3,4]}
+            # dict must match static_model.create_feed()
+            yield output_dict
+
+        return reader
+
+
+if __name__ == "__main__":
+    r = Reader()
+    r.init()
+    r.run_from_stdin()

test/ps/download_criteo_data.sh

@@ -0,0 +1,17 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+wget --no-check-certificate https://paddlerec.bj.bcebos.com/benchmark/sample_train.txt
+mkdir train_data
+mv sample_train.txt train_data/

test/ps/download_data.sh

@@ -0,0 +1,27 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+wget --no-check-certificate https://fleet.bj.bcebos.com/ctr_data.tar.gz
+tar -zxvf ctr_data.tar.gz
+mv ./raw_data ./train_data_full
+mkdir train_data && cd train_data
+cp ../train_data_full/part-0 ../train_data_full/part-1 ./ && cd ..
+mv ./test_data ./test_data_full
+mkdir test_data && cd test_data
+cp ../test_data_full/part-220 ./  && cd ..
+echo "Complete data download."
+echo "Full Train data stored in ./train_data_full "
+echo "Full Test data stored in ./test_data_full "
+echo "Rapid Verification train data stored in ./train_data "
+echo "Rapid Verification test data stored in ./test_data "

test/ps/fl_async_ps_config.yaml

@@ -0,0 +1,37 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# refer to PaddleRec/models/rank/dnn/benchmark.yaml
+
+hyper_parameters:
+  optimizer:
+    class: Adam
+    learning_rate: 0.0001
+    adam_lazy_mode: True
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1000001
+  sparse_feature_dim: 10
+  dense_input_dim: 13
+  fc_sizes: []
+
+runner:
+  sync_mode: "async"  # sync / async / geo / heter
+  is_fl_ps_mode: 1
+  reader_thread_num: 16
+  use_gpu: 0
+  batch_size: 2
+  train_files_path: "./train_data"
+  epoch_num: 4
+
+  model_path: "../ps_dnn_model.py"

test/ps/fl_ps_trainer.py

@@ -0,0 +1,162 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+
+import paddle
+from paddle import base
+from paddle.distributed import fleet
+
+
+def get_dataset(inputs, config, pipe_cmd, role="worker"):
+    dataset = base.DatasetFactory().create_dataset()
+    dataset.set_use_var(inputs)
+    dataset.set_pipe_command(pipe_cmd)
+    dataset.set_batch_size(config.get('runner.batch_size'))
+    reader_thread_num = int(config.get('runner.reader_thread_num'))
+    dataset.set_thread(reader_thread_num)
+    train_files_path = config.get('runner.train_files_path')
+    print(f'train_data_files:{train_files_path}')
+    file_list = [
+        os.path.join(train_files_path, x) for x in os.listdir(train_files_path)
+    ]
+    if role == "worker":
+        file_list = fleet.util.get_file_shard(file_list)
+        print(f"worker file list: {file_list}")
+    elif role == "heter_worker":
+        file_list = fleet.util.get_heter_file_shard(file_list)
+        print(f"heter worker file list: {file_list}")
+
+    return dataset, file_list
+
+
+def fl_ps_train():
+    # 0. get role
+    from paddle.distributed.fleet.base import role_maker
+
+    role_maker = role_maker.PaddleCloudRoleMaker()
+    role_maker._generate_role()
+    fleet.util._set_role_maker(role_maker)
+
+    # 1. load yaml-config to dict-config
+    from ps_dnn_trainer import (
+        StaticModel,
+        YamlHelper,
+        get_user_defined_strategy,
+    )
+
+    yaml_helper = YamlHelper()
+    config_yaml_path = '../ps/fl_async_ps_config.yaml'
+    config = yaml_helper.load_yaml(config_yaml_path)
+    # yaml_helper.print_yaml(config)
+
+    # 2. get static model
+    paddle.enable_static()
+    model = StaticModel(config)
+    feeds_list = model.create_feeds()
+    metrics = model.fl_net(feeds_list)
+    loss = model._cost
+
+    # 3. compile time - build program_desc
+    user_defined_strategy = get_user_defined_strategy(config)
+    a_sync_configs = user_defined_strategy.a_sync_configs
+    a_sync_configs["launch_barrier"] = True
+    user_defined_strategy.a_sync_configs = a_sync_configs
+    print(
+        "launch_barrier: ",
+        user_defined_strategy.a_sync_configs["launch_barrier"],
+    )
+    learning_rate = config.get("hyper_parameters.optimizer.learning_rate")
+    inner_optimizer = paddle.optimizer.Adam(learning_rate, lazy_mode=True)
+    from paddle.distributed.fleet.meta_optimizers.ps_optimizer import (
+        ParameterServerOptimizer,
+    )
+
+    ps_optimizer = ParameterServerOptimizer(inner_optimizer)
+    ps_optimizer._set_basic_info(
+        loss, role_maker, inner_optimizer, user_defined_strategy
+    )
+    ps_optimizer.minimize_impl(loss)
+
+    # 4. runtime
+    from paddle.distributed.ps.the_one_ps import TheOnePSRuntime
+
+    _runtime_handle = TheOnePSRuntime()  # ps 目录下重构版的 TheOnePSRuntime
+    _runtime_handle._set_basic_info(ps_optimizer.pass_ctx._attrs)
+    epoch_num = int(config.get('runner.epoch_num'))
+    # 4.1 run server - build fleet_desc
+    if role_maker._is_server():
+        _runtime_handle._init_server()
+        _runtime_handle._run_server()
+    # 4.2 run worker
+    elif role_maker._is_worker():
+        place = base.CPUPlace()
+        exe = base.Executor(place)
+        exe.run(base.default_startup_program())
+        _runtime_handle._init_worker()
+        print('trainer get dataset')
+        inputs = feeds_list[1:-1]
+        dataset, file_list = get_dataset(
+            inputs, config, "python dataset_generator_A.py"
+        )
+        print(
+            f"base.default_main_program: {base.default_main_program()._heter_pipeline_opt}"
+        )
+        for epoch in range(epoch_num):
+            # A 方和 B 方如果要以文件粒度 shuffle 时，则需要固定同一个种子
+            dataset.set_filelist(file_list)
+            start_time = time.time()
+            exe.train_from_dataset(
+                program=base.default_main_program(),
+                dataset=dataset,
+                print_period=2,
+                debug=False,
+            )
+            end_time = time.time()
+            print(
+                f"trainer epoch {epoch} finished, use time={end_time - start_time}\n"
+            )
+        exe.close()
+        _runtime_handle._stop_worker()
+        print("Fl partyA Trainer Success!")
+    else:
+        exe = base.Executor()
+        exe.run(base.default_startup_program())
+        _runtime_handle._init_worker()
+        inputs = [
+            feeds_list[0],
+            feeds_list[-1],
+        ]  # 顺序务必要和 dataset_generator_B.py 中保持一致
+        dataset, file_list = get_dataset(
+            inputs, config, "python dataset_generator_B.py", "heter_worker"
+        )
+        print(
+            f"base.default_main_program: {base.default_main_program()._heter_pipeline_opt}"
+        )
+        for epoch in range(epoch_num):
+            dataset.set_filelist(file_list)
+            exe.train_from_dataset(
+                program=base.default_main_program(),
+                dataset=dataset,
+                print_period=2,
+                debug=False,
+            )
+        exe.close()
+        _runtime_handle._stop_worker()
+        print("Fl partB Trainer Success!")
+
+
+if __name__ == '__main__':
+    fl_ps_train()

test/ps/gpu_ps_config.yaml

@@ -0,0 +1,35 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# refer to PaddleRec/models/rank/dnn/config_gpubox.yaml
+
+hyper_parameters:
+  optimizer:
+    class: Adam
+    learning_rate: 0.001
+    strategy: async
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1024
+  sparse_feature_dim: 11
+  dense_input_dim: 13
+  fc_sizes: [512, 256, 128, 32]
+  distributed_embedding: 0
+
+runner:
+  geo_step: 400
+  sync_mode: "gpubox"
+  thread_num: 16
+  use_gpu: 1
+
+  model_path: "../ps_dnn_model.py"

test/ps/gpubox_run.sh

@@ -0,0 +1,64 @@
+# !/bin/bash
+
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+if [ ! -d "./log" ]; then
+  mkdir ./log
+  echo "Create log floder for store running log"
+fi
+
+export FLAGS_LAUNCH_BARRIER=0
+export PADDLE_TRAINER_ID=0
+export PADDLE_PSERVER_NUMS=1
+export PADDLE_TRAINERS=1
+export PADDLE_TRAINERS_NUM=${PADDLE_TRAINERS}
+export POD_IP=127.0.0.1
+
+# set free port if 29011 is occupied
+export PADDLE_PSERVERS_IP_PORT_LIST="127.0.0.1:29011"
+export PADDLE_PSERVER_PORT_ARRAY=(29011)
+
+# set gpu numbers according to your device
+#export FLAGS_selected_gpus="0,1,2,3,4,5,6,7"
+export FLAGS_selected_gpus="0,1"
+
+# set your model yaml
+#SC="gpubox_ps_trainer.py"
+SC="static_gpubox_trainer.py"
+
+# run pserver
+export TRAINING_ROLE=PSERVER
+for((i=0;i<$PADDLE_PSERVER_NUMS;i++))
+do
+    cur_port=${PADDLE_PSERVER_PORT_ARRAY[$i]}
+    echo "PADDLE WILL START PSERVER "$cur_port
+    export PADDLE_PORT=${cur_port}
+    python -u $SC &> ./log/pserver.$i.log &
+done
+
+# run trainer
+export TRAINING_ROLE=TRAINER
+for((i=0;i<$PADDLE_TRAINERS;i++))
+do
+    echo "PADDLE WILL START Trainer "$i
+    export PADDLE_TRAINER_ID=$i
+    python -u $SC &> ./log/worker.$i.log
+done
+
+if [ $? -eq 0 ];then
+echo "Training log stored in ./log/"
+else
+exit 1
+fi

test/ps/heter_ps_config.yaml

@@ -0,0 +1,34 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+hyper_parameters:
+  optimizer:
+    class: Adam
+    learning_rate: 0.0001
+    strategy: async    # 有用
+  sparse_inputs_slots: 27
+  sparse_feature_number: 1024
+  sparse_feature_dim: 11
+  dense_input_dim: 13
+  fc_sizes: [512, 256, 128, 32]
+  distributed_embedding: 0
+
+runner:
+  sync_mode: "heter"
+  thread_num: 8
+  micro_num: 8  # micro batch num for each thread
+  pipeline: True
+
+  model_path: "../ps_dnn_model.py"

test/ps/ps_dnn_model.py

@@ -0,0 +1,389 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import paddle
+from paddle import nn
+
+
+class DNNLayer(nn.Layer):
+    def __init__(
+        self,
+        sparse_feature_number,
+        sparse_feature_dim,
+        dense_feature_dim,
+        num_field,
+        layer_sizes,
+        sync_mode=None,
+    ):
+        super().__init__()
+        self.sync_mode = sync_mode
+        self.sparse_feature_number = sparse_feature_number
+        self.sparse_feature_dim = sparse_feature_dim
+        self.dense_feature_dim = dense_feature_dim
+        self.num_field = num_field
+        self.layer_sizes = layer_sizes
+
+        self.embedding = paddle.nn.Embedding(
+            self.sparse_feature_number,
+            self.sparse_feature_dim,
+            sparse=True,
+            weight_attr=paddle.ParamAttr(
+                name="SparseFeatFactors",
+                initializer=paddle.nn.initializer.Uniform(),
+            ),
+        )
+
+        sizes = [
+            sparse_feature_dim * num_field + dense_feature_dim,
+            *self.layer_sizes,
+            2,
+        ]
+        acts = ["relu" for _ in range(len(self.layer_sizes))] + [None]
+        self._mlp_layers = []
+        for i in range(len(layer_sizes) + 1):
+            linear = paddle.nn.Linear(
+                in_features=sizes[i],
+                out_features=sizes[i + 1],
+                weight_attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Normal(
+                        std=1.0 / math.sqrt(sizes[i])
+                    )
+                ),
+            )
+            self.add_sublayer(f'linear_{i}', linear)
+            self._mlp_layers.append(linear)
+            if acts[i] == 'relu':
+                act = paddle.nn.ReLU()
+                self.add_sublayer(f'act_{i}', act)
+                self._mlp_layers.append(act)
+
+    def forward(self, sparse_inputs, dense_inputs):
+        sparse_embs = []
+        for s_input in sparse_inputs:
+            if self.sync_mode == "gpubox":
+                emb = paddle.static.nn.sparse_embedding(
+                    input=s_input,
+                    size=[self.sparse_feature_number, self.sparse_feature_dim],
+                    param_attr=paddle.ParamAttr(name="embedding"),
+                )
+            else:
+                emb = self.embedding(s_input)
+            emb = paddle.reshape(emb, shape=[-1, self.sparse_feature_dim])
+            # emb.stop_gradient = True
+            sparse_embs.append(emb)
+
+        y_dnn = paddle.concat(x=[*sparse_embs, dense_inputs], axis=1)
+
+        if self.sync_mode == 'heter':
+            with paddle.base.device_guard('gpu'):
+                for n_layer in self._mlp_layers:
+                    y_dnn = n_layer(y_dnn)
+        else:
+            for n_layer in self._mlp_layers:
+                y_dnn = n_layer(y_dnn)
+
+        return y_dnn
+
+
+class FlDNNLayer(nn.Layer):
+    def __init__(
+        self,
+        sparse_feature_number,
+        sparse_feature_dim,
+        dense_feature_dim,
+        sparse_number,
+        sync_mode=None,
+    ):
+        super().__init__()
+
+        self.PART_A_DEVICE_FlAG = 'gpu:0'
+        self.PART_A_JOINT_OP_DEVICE_FlAG = 'gpu:2'
+        self.PART_B_DEVICE_FlAG = 'gpu:1'
+        self.PART_B_JOINT_OP_DEVICE_FlAG = 'gpu:3'
+
+        self.sync_mode = sync_mode
+        self.sparse_feature_number = sparse_feature_number
+        self.sparse_feature_dim = sparse_feature_dim
+        self.slot_num = sparse_number
+        self.dense_feature_dim = dense_feature_dim
+
+        layer_sizes_a = [
+            self.slot_num * self.sparse_feature_dim,
+            5,
+            7,
+        ]  # for test
+        layer_sizes_b = [self.dense_feature_dim, 6, 7]
+        layer_sizes_top = [7, 2]
+
+        self.embedding = paddle.nn.Embedding(
+            self.sparse_feature_number,
+            self.sparse_feature_dim,
+            sparse=True,
+            weight_attr=paddle.ParamAttr(
+                name="SparseFeatFactors",
+                initializer=paddle.nn.initializer.Uniform(),
+            ),
+        )
+
+        # part_a fc
+        acts = ["relu" for _ in range(len(layer_sizes_a))]
+        self._mlp_layers_a = []
+        for i in range(len(layer_sizes_a) - 1):
+            linear = paddle.nn.Linear(
+                in_features=layer_sizes_a[i],
+                out_features=layer_sizes_a[i + 1],
+                weight_attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Normal(
+                        std=1.0 / math.sqrt(layer_sizes_a[i])
+                    )
+                ),
+            )
+            self.add_sublayer(f'linear_{i}', linear)
+            self._mlp_layers_a.append(linear)
+            act = paddle.nn.ReLU()
+            self.add_sublayer(f'act_{i}', act)
+            self._mlp_layers_a.append(act)
+
+        # part_b fc
+        acts = ["relu" for _ in range(len(layer_sizes_b))]
+        self._mlp_layers_b = []
+        for i in range(len(layer_sizes_b) - 1):
+            linear = paddle.nn.Linear(
+                in_features=layer_sizes_b[i],
+                out_features=layer_sizes_b[i + 1],
+                weight_attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Normal(
+                        std=1.0 / math.sqrt(layer_sizes_b[i])
+                    )
+                ),
+            )
+            self.add_sublayer(f'linear_{i}', linear)
+            self._mlp_layers_b.append(linear)
+            act = paddle.nn.ReLU()
+            self.add_sublayer(f'act_{i}', act)
+            self._mlp_layers_b.append(act)
+
+        # top fc
+        acts = ["relu" for _ in range(len(layer_sizes_top))]
+        self._mlp_layers_top = []
+        for i in range(len(layer_sizes_top) - 1):
+            linear = paddle.nn.Linear(
+                in_features=layer_sizes_top[i],
+                out_features=layer_sizes_top[i + 1],
+                weight_attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Normal(
+                        std=1.0 / math.sqrt(layer_sizes_top[i])
+                    )
+                ),
+            )
+            self.add_sublayer(f'linear_{i}', linear)
+            self._mlp_layers_top.append(linear)
+            act = paddle.nn.ReLU()
+            self.add_sublayer(f'act_{i}', act)
+            self._mlp_layers_top.append(act)
+
+    def bottom_a_layer(self, sparse_inputs):
+        with paddle.base.device_guard(self.PART_A_DEVICE_FlAG):
+            sparse_embs = []
+            for s_input in sparse_inputs:
+                emb = self.embedding(s_input)
+                emb = paddle.reshape(emb, shape=[-1, self.sparse_feature_dim])
+                sparse_embs.append(emb)
+
+            y = paddle.concat(x=sparse_embs, axis=1)
+            y = self._mlp_layers_a[0](y)
+            y = self._mlp_layers_a[1](y)
+
+            y = self._mlp_layers_a[2](y)
+        with paddle.base.device_guard(
+            self.PART_A_JOINT_OP_DEVICE_FlAG
+        ):  # joint point
+            bottom_a = self._mlp_layers_a[3](y)
+
+        return bottom_a
+
+    def bottom_b_layer(self, dense_inputs):
+        with paddle.base.device_guard(self.PART_B_DEVICE_FlAG):
+            y = self._mlp_layers_b[0](dense_inputs)
+            y = self._mlp_layers_b[1](y)
+
+            y = self._mlp_layers_b[2](y)
+            bottom_b = self._mlp_layers_b[3](y)
+
+        return bottom_b
+
+    def interactive_layer(self, bottom_a, bottom_b):
+        with paddle.base.device_guard(
+            self.PART_B_JOINT_OP_DEVICE_FlAG
+        ):  # joint point
+            interactive = paddle.add(bottom_a, bottom_b)
+        return interactive
+
+    def top_layer(self, interactive, label_input):
+        with paddle.base.device_guard(self.PART_B_DEVICE_FlAG):
+            y = self._mlp_layers_top[0](interactive)
+            y_top = self._mlp_layers_top[1](y)
+            predict_2d = paddle.nn.functional.softmax(y_top)
+            (
+                auc,
+                batch_auc,
+                [
+                    self.batch_stat_pos,
+                    self.batch_stat_neg,
+                    self.stat_pos,
+                    self.stat_neg,
+                ],
+            ) = paddle.static.auc(
+                input=predict_2d,
+                label=label_input,
+                num_thresholds=2**12,
+                slide_steps=20,
+            )
+
+            cost = paddle.nn.functional.cross_entropy(
+                input=y_top, label=label_input
+            )
+            avg_cost = paddle.mean(x=cost)
+
+        return auc, avg_cost
+
+    def forward(self, sparse_inputs, dense_inputs, label_input):
+        bottom_a = self.bottom_a_layer(sparse_inputs)
+
+        bottom_b = self.bottom_b_layer(dense_inputs)
+
+        interactive = self.interactive_layer(bottom_a, bottom_b)
+
+        auc, avg_cost = self.top_layer(interactive, label_input)
+
+        return auc, avg_cost
+
+
+class StaticModel:
+    def __init__(self, config):
+        self.cost = None
+        self.infer_target_var = None
+        self.config = config
+        self._init_hyper_parameters()
+        self.sync_mode = config.get("runner.sync_mode")
+
+    def _init_hyper_parameters(self):
+        self.is_distributed = False
+        self.distributed_embedding = False
+
+        if self.config.get("hyper_parameters.distributed_embedding", 0) == 1:
+            self.distributed_embedding = True
+
+        self.sparse_feature_number = self.config.get(
+            "hyper_parameters.sparse_feature_number"
+        )
+        self.sparse_feature_dim = self.config.get(
+            "hyper_parameters.sparse_feature_dim"
+        )
+        self.sparse_inputs_slots = self.config.get(
+            "hyper_parameters.sparse_inputs_slots"
+        )
+        self.dense_input_dim = self.config.get(
+            "hyper_parameters.dense_input_dim"
+        )
+        self.learning_rate = self.config.get(
+            "hyper_parameters.optimizer.learning_rate"
+        )
+        self.fc_sizes = self.config.get("hyper_parameters.fc_sizes")
+
+    def create_feeds(self, is_infer=False):
+        dense_input = paddle.static.data(
+            name="dense_input",
+            shape=[None, self.dense_input_dim],
+            dtype="float32",
+        )
+
+        sparse_input_ids = [
+            paddle.static.data(name=str(i), shape=[None, 1], dtype="int64")
+            for i in range(1, self.sparse_inputs_slots)
+        ]
+
+        label = paddle.static.data(name="label", shape=[None, 1], dtype="int64")
+
+        feeds_list = [label, *sparse_input_ids, dense_input]
+        return feeds_list
+
+    def net(self, input, is_infer=False):
+        self.label_input = input[0]
+        self.sparse_inputs = input[1 : self.sparse_inputs_slots]
+        self.dense_input = input[-1]
+        sparse_number = self.sparse_inputs_slots - 1
+
+        dnn_model = DNNLayer(
+            self.sparse_feature_number,
+            self.sparse_feature_dim,
+            self.dense_input_dim,
+            sparse_number,
+            self.fc_sizes,
+            sync_mode=self.sync_mode,
+        )
+        raw_predict_2d = dnn_model.forward(self.sparse_inputs, self.dense_input)
+        predict_2d = paddle.nn.functional.softmax(raw_predict_2d)
+        self.predict = predict_2d
+        (
+            auc,
+            batch_auc,
+            [
+                self.batch_stat_pos,
+                self.batch_stat_neg,
+                self.stat_pos,
+                self.stat_neg,
+            ],
+        ) = paddle.static.auc(
+            input=self.predict,
+            label=self.label_input,
+            num_thresholds=2**12,
+            slide_steps=20,
+        )
+        self.inference_target_var = auc
+        if is_infer:
+            fetch_dict = {'auc': auc}
+            return fetch_dict
+
+        cost = paddle.nn.functional.cross_entropy(
+            input=raw_predict_2d, label=self.label_input
+        )
+        avg_cost = paddle.mean(x=cost)
+        self._cost = avg_cost
+
+        fetch_dict = {'cost': avg_cost, 'auc': auc}
+        return fetch_dict
+
+    def fl_net(self, input, is_infer=False):
+        self.label_input = input[0]
+        self.sparse_inputs = input[1 : self.sparse_inputs_slots]
+        self.dense_input = input[-1]
+        self.sparse_number = self.sparse_inputs_slots - 1
+
+        fl_dnn_model = FlDNNLayer(
+            self.sparse_feature_number,
+            self.sparse_feature_dim,
+            self.dense_input_dim,
+            self.sparse_number,
+            sync_mode=self.sync_mode,
+        )
+
+        auc, avg_cost = fl_dnn_model.forward(
+            self.sparse_inputs, self.dense_input, self.label_input
+        )
+        fetch_dict = {'cost': avg_cost, 'auc': auc}
+        self._cost = avg_cost
+        return fetch_dict

test/ps/ps_dnn_trainer.py

@@ -0,0 +1,600 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import ast
+import copy
+import os
+import struct
+import sys
+
+import numpy as np
+import yaml
+from ps_dnn_model import StaticModel
+
+import paddle
+from paddle.distributed import fleet
+from paddle.distributed.fleet.base import role_maker
+from paddle.distributed.ps.utils.ps_program_builder import (
+    debug_program,
+    logger,
+    new_pass,
+    ps_log_root_dir,
+)
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
+
+
+def is_distributed_env():
+    node_role = os.getenv("TRAINING_ROLE")
+    print(f"-- Role: {node_role} --")
+    if node_role is None:
+        return False
+    else:
+        return True
+
+
+class YamlHelper:
+    def load_yaml(self, yaml_file, other_part=None):
+        part_list = ["runner", "hyper_parameters"]
+        if other_part:
+            part_list += other_part
+        running_config = self.get_all_inters_from_yaml(yaml_file, part_list)
+        running_config = self.workspace_adapter(running_config)
+        return running_config
+
+    def print_yaml(self, config):
+        print(self.pretty_print_envs(config))
+
+    def parse_yaml(self, config):
+        vs = [int(i) for i in yaml.__version__.split(".")]
+        if vs[0] < 5:
+            use_full_loader = False
+        elif vs[0] > 5:
+            use_full_loader = True
+        else:
+            if vs[1] >= 1:
+                use_full_loader = True
+            else:
+                use_full_loader = False
+
+        if os.path.isfile(config):
+            with open(config, 'r', encoding="utf-8") as rb:
+                if use_full_loader:
+                    _config = yaml.load(rb.read(), Loader=yaml.FullLoader)
+                else:
+                    _config = yaml.load(rb.read())
+                return _config
+        else:
+            raise ValueError(f"config {config} can not be supported")
+
+    def get_all_inters_from_yaml(self, file, filters):
+        _envs = self.parse_yaml(file)
+        all_flattens = {}
+
+        def fatten_env_namespace(namespace_nests, local_envs):
+            for k, v in local_envs.items():
+                if isinstance(v, dict):
+                    nests = copy.deepcopy(namespace_nests)
+                    nests.append(k)
+                    fatten_env_namespace(nests, v)
+                else:
+                    global_k = ".".join([*namespace_nests, k])
+                    all_flattens[global_k] = v
+
+        fatten_env_namespace([], _envs)
+        ret = {}
+        for k, v in all_flattens.items():
+            for f in filters:
+                if k.startswith(f):
+                    ret[k] = v
+        return ret
+
+    def workspace_adapter(self, config):
+        workspace = config.get("workspace")
+        for k, v in config.items():
+            if isinstance(v, str) and "{workspace}" in v:
+                config[k] = v.replace("{workspace}", workspace)
+        return config
+
+    def pretty_print_envs(self, envs, header=None):
+        spacing = 2
+        max_k = 40
+        max_v = 45
+
+        for k, v in envs.items():
+            max_k = max(max_k, len(k))
+
+        h_format = "    " + "|{{:>{}s}}{}{{:^{}s}}|\n".format(
+            max_k, " " * spacing, max_v
+        )
+        l_format = "    " + f"|{{:>{max_k}s}}{{}}{{:^{max_v}s}}|\n"
+        length = max_k + max_v + spacing
+
+        border = "    +" + "".join(["="] * length) + "+"
+        line = "    +" + "".join(["-"] * length) + "+"
+
+        draws = ""
+        draws += border + "\n"
+
+        if header:
+            draws += h_format.format(header[0], header[1])
+        else:
+            draws += h_format.format("Ps Benchmark Envs", "Value")
+
+        draws += line + "\n"
+
+        for k, v in sorted(envs.items()):
+            if isinstance(v, str) and len(v) >= max_v:
+                str_v = "... " + v[-41:]
+            else:
+                str_v = v
+
+            draws += l_format.format(k, " " * spacing, str(str_v))
+
+        draws += border
+
+        _str = f"\n{draws}\n"
+        return _str
+
+
+def get_user_defined_strategy(config):
+    if not is_distributed_env():
+        logger.warning(
+            "Not Find Distributed env, Change To local train mode. If you want train with fleet, please use [fleetrun] command."
+        )
+        # return None
+    sync_mode = config.get("runner.sync_mode")
+    assert sync_mode in ["async", "sync", "geo", "heter", "gpubox"]
+    if sync_mode == "sync":
+        strategy = paddle.distributed.fleet.DistributedStrategy()
+        strategy.a_sync = False
+    elif sync_mode == "async":
+        strategy = paddle.distributed.fleet.DistributedStrategy()
+        strategy.a_sync = True
+        strategy.is_fl_ps_mode = (
+            True if config.get("runner.is_fl_ps_mode") == 1 else False
+        )
+        if strategy.is_fl_ps_mode:
+            strategy.pipeline = False
+            micro_num = 1
+            strategy.pipeline_configs = {
+                "accumulate_steps": micro_num
+            }  # num_microbatches
+    elif sync_mode == "geo":
+        strategy = paddle.distributed.fleet.DistributedStrategy()
+        strategy.a_sync = True
+        strategy.a_sync_configs = {"k_steps": config.get("runner.geo_step")}
+    elif sync_mode == "heter":
+        strategy = paddle.distributed.fleet.DistributedStrategy()
+        strategy.a_sync = True
+        strategy.a_sync_configs = {"heter_worker_device_guard": "gpu"}
+        strategy.pipeline = True
+        strategy.pipeline_configs = {
+            "accumulate_steps": config.get('runner.micro_num')
+        }
+    elif sync_mode == "gpubox":
+        print(f"sync_mode = {sync_mode}")
+        strategy = paddle.distributed.fleet.DistributedStrategy()
+        strategy.a_sync = True
+        strategy.a_sync_configs = {"use_ps_gpu": 1}
+
+    strategy.trainer_desc_configs = {
+        "dump_fields_path": config.get("runner.dump_fields_path", ""),
+        "dump_fields": config.get("runner.dump_fields", []),
+        "dump_param": config.get("runner.dump_param", []),
+        "stat_var_names": config.get("stat_var_names", []),
+        "local_sparse": config.get("runner.local_sparse", []),
+        "remote_sparse": config.get("runner.remote_sparse", []),
+    }
+    print("strategy:", strategy.trainer_desc_configs)
+
+    if config.get("runner.fs_client.uri") is not None:
+        strategy.fs_client_param = {
+            "uri": config.get("runner.fs_client.uri", ""),
+            "user": config.get("runner.fs_client.user", ""),
+            "passwd": config.get("runner.fs_client.passwd", ""),
+            "hadoop_bin": config.get("runner.fs_client.hadoop_bin", "hadoop"),
+        }
+    print("strategy:", strategy.fs_client_param)
+
+    strategy.adam_d2sum = config.get("hyper_parameters.adam_d2sum", True)
+    table_config = {}
+    for x in config:
+        if x.startswith("table_parameters"):
+            table_name = x.split('.')[1]
+            if table_name not in table_config:
+                table_config[table_name] = {}
+            table_config[table_name][x] = config[x]
+    print("table_config:", table_config)
+    strategy.sparse_table_configs = table_config
+    print("strategy table config:", strategy.sparse_table_configs)
+    a_sync_configs = strategy.a_sync_configs
+    a_sync_configs["launch_barrier"] = False
+    # a_sync_configs["launch_barrier"] = True
+    strategy.a_sync_configs = a_sync_configs
+    print("launch_barrier: ", strategy.a_sync_configs["launch_barrier"])
+
+    return strategy
+
+
+def get_distributed_strategy(user_defined_strategy):  # pslib
+    from paddle.incubate.distributed.fleet.parameter_server.distribute_transpiler.distributed_strategy import (
+        StrategyFactory,
+    )
+
+    k_steps = user_defined_strategy.a_sync_configs["k_steps"]
+    strategy = None
+
+    if not user_defined_strategy.a_sync and k_steps == 0:
+        strategy = StrategyFactory.create_sync_strategy()
+
+    if user_defined_strategy.a_sync and k_steps == 0:
+        strategy = StrategyFactory.create_async_strategy()
+
+    if user_defined_strategy.a_sync and k_steps > 0:
+        strategy = StrategyFactory.create_geo_strategy(k_steps)
+
+    if not strategy:
+        raise ValueError("k_steps must be invalid value, please check")
+
+    return strategy
+
+
+def get_model(config):
+    abs_dir = config['config_abs_dir']
+    sys.path.append(abs_dir)
+    static_model = StaticModel(config)
+    return static_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("PsTest train script")
+    parser.add_argument(
+        '-m', '--config_yaml', type=str, required=True, help='config file path'
+    )
+    parser.add_argument(
+        '-bf16',
+        '--pure_bf16',
+        type=ast.literal_eval,
+        default=False,
+        help="whether use bf16",
+    )
+
+    parser.add_argument(
+        '--run_minimize', type=int, default=0, help="test single pass"
+    )
+    parser.add_argument(
+        '--run_single_pass', type=int, default=0, help="test single pass"
+    )
+    parser.add_argument(
+        '--run_the_one_ps', type=int, default=0, help="test the_one_ps"
+    )
+    parser.add_argument(
+        '--debug_new_minimize', type=int, default=0, help="test single pass"
+    )
+    parser.add_argument(
+        '--debug_new_pass', type=int, default=0, help="test single pass"
+    )
+    parser.add_argument(
+        '--applied_pass_name', type=str, default="", help="test single pass"
+    )
+    parser.add_argument(
+        '--debug_the_one_ps', type=int, default=0, help="test the_one_ps"
+    )
+
+    args = parser.parse_args()
+    args.abs_dir = os.path.dirname(os.path.abspath(args.config_yaml))
+    yaml_helper = YamlHelper()
+    config = yaml_helper.load_yaml(args.config_yaml)
+    config["yaml_path"] = args.config_yaml
+    config["config_abs_dir"] = args.abs_dir
+    config["pure_bf16"] = args.pure_bf16
+    config['run_minimize'] = args.run_minimize
+    config['run_single_pass'] = args.run_single_pass
+    config['run_the_one_ps'] = args.run_the_one_ps
+    config['debug_new_minimize'] = args.debug_new_minimize
+    config['debug_new_pass'] = args.debug_new_pass
+    config['applied_pass_name'] = args.applied_pass_name
+    config['debug_the_one_ps'] = args.debug_the_one_ps
+    yaml_helper.print_yaml(config)
+    return config
+
+
+def bf16_to_fp32(val):
+    return np.float32(struct.unpack('<f', struct.pack('<I', val << 16))[0])
+
+
+class DnnTrainer:
+    def __init__(self, config):
+        self.metrics = {}
+        self.config = config
+        self.input_data = None
+        self.reader = None
+        self.exe = None
+        self.train_result_dict = {}
+        self.train_result_dict["speed"] = []
+        self.model = None
+        self.pure_bf16 = self.config['pure_bf16']
+        self.role_maker = role_maker.PaddleCloudRoleMaker()
+
+    def init_fleet_with_gloo(self, use_gloo=False):
+        if use_gloo:
+            os.environ["PADDLE_WITH_GLOO"] = "1"
+            fleet.init(self.role_maker)
+        else:
+            fleet.init()
+
+        if fleet.is_server():
+            print(f"server: {fleet.server_index()} started")
+        else:
+            print(f"worker: {fleet.worker_index()} started")
+
+    def run_minimize(self):
+        self.init_fleet_with_gloo()
+        self.model = get_model(self.config)
+        print("cpu_num: {}".format(os.getenv("CPU_NUM")))
+        self.input_data = self.model.create_feeds()
+        self.metrics = self.model.net(self.input_data)
+        loss = self.model._cost
+        user_defined_strategy = get_user_defined_strategy(self.config)
+        learning_rate = self.config.get(
+            "hyper_parameters.optimizer.learning_rate"
+        )
+        sync_mode = self.config.get("runner.sync_mode")
+        inner_optimizer = paddle.optimizer.Adam(learning_rate, lazy_mode=True)
+
+        self.role_maker._generate_role()  # 必要
+        if self.config['debug_new_minimize'] == 1:
+            print("entering run_minimize -- new")
+            from paddle.distributed.fleet.meta_optimizers.ps_optimizer import (
+                ParameterServerOptimizer,
+            )
+
+            ps_optimizer = ParameterServerOptimizer(inner_optimizer)
+            ps_optimizer._set_basic_info(
+                loss, self.role_maker, inner_optimizer, user_defined_strategy
+            )
+            ps_optimizer.minimize_impl(loss)
+        else:
+            print("entering run_minimize -- old")
+            fleet_obj = fleet.distributed_optimizer(
+                inner_optimizer, user_defined_strategy
+            )  # Fleet object
+            fleet_obj.minimize(loss)
+
+        if fleet.is_server():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + '_run_minimize'
+                + '_debug:_'
+                + str(self.config['debug_new_minimize'])
+                + '_server_main.prototxt'
+            )
+            debug_program(_main_file, loss.block.program)
+        elif fleet.is_worker():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + '_run_minimize'
+                + '_debug:_'
+                + str(self.config['debug_new_minimize'])
+                + '_worker_main.prototxt'
+            )
+            debug_program(_main_file, loss.block.program)
+        elif self.role_maker._is_heter_worker():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + '_run_minimize'
+                + '_debug:_'
+                + str(self.config['debug_new_minimize'])
+                + '_heter_worker_main.prototxt'
+            )
+            debug_program(_main_file, loss.block.program)
+
+    def run_single_pass(self):
+        self.init_fleet_with_gloo()
+        self.model = get_model(config)
+        input_data = self.model.create_feeds()
+        metrics = self.model.net(input_data)
+        loss = self.model._cost
+        user_defined_strategy = get_user_defined_strategy(config)
+        learning_rate = config.get("hyper_parameters.optimizer.learning_rate")
+        sync_mode = self.config.get("runner.sync_mode")
+        inner_optimizer = paddle.optimizer.Adam(learning_rate, lazy_mode=True)
+        startup_program = paddle.static.default_startup_program()
+        inner_optimizer.minimize(loss, startup_program)
+        if self.config['debug_new_pass'] == 1:
+            print(
+                "entering run {} - new".format(str(config["applied_pass_name"]))
+            )
+            from paddle.distributed.fleet.meta_optimizers.ps_optimizer import (
+                ParameterServerOptimizer,
+            )
+
+            ps_optimizer = ParameterServerOptimizer(inner_optimizer)
+            ps_optimizer._set_basic_info(
+                loss, self.role_maker, inner_optimizer, user_defined_strategy
+            )
+            ps_optimizer._set_origin_programs([loss])
+            ps_optimizer._init_ps_pass_context(loss, startup_program)
+            _main = ps_optimizer.pass_ctx._attrs['cloned_main']
+
+            append_send_ops_pass = new_pass(
+                config["applied_pass_name"], ps_optimizer.pass_ctx._attrs
+            )
+            append_send_ops_pass.apply([_main], [None], ps_optimizer.pass_ctx)
+        else:
+            print(
+                "entering run {} - old".format(str(config["applied_pass_name"]))
+            )
+            from paddle.incubate.distributed.fleet.parameter_server.ir import (
+                public,
+            )
+
+            dist_strategy = get_distributed_strategy(user_defined_strategy)
+            compiled_config = public.CompileTimeStrategy(
+                loss.block.program,
+                startup_program,
+                dist_strategy,
+                self.role_maker,
+            )
+
+            _main = compiled_config.origin_main_program.clone()
+            _startup = compiled_config.origin_startup_program.clone()
+            from paddle.incubate.distributed.fleet.parameter_server.ir import (
+                trainer_pass as worker,
+            )
+
+            _main = worker.append_send_ops_pass(_main, compiled_config)
+
+        if fleet.is_server():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + "_"
+                + str(config["applied_pass_name"])
+                + '_debug:_'
+                + str(self.config['debug_new_pass'])
+                + '_server_main.prototxt'
+            )
+            debug_program(_main_file, _main)
+        elif fleet.is_worker():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + "_"
+                + str(config["applied_pass_name"])
+                + '_debug:_'
+                + str(self.config['debug_new_pass'])
+                + '_worker_main.prototxt'
+            )
+            debug_program(_main_file, _main)
+
+    def run_the_one_ps(self):
+        self.init_fleet_with_gloo()
+        self.model = get_model(self.config)
+        self.input_data = self.model.create_feeds()
+        self.metrics = self.model.net(self.input_data)
+        loss = self.model._cost
+        user_defined_strategy = get_user_defined_strategy(self.config)
+        learning_rate = self.config.get(
+            "hyper_parameters.optimizer.learning_rate"
+        )
+        sync_mode = self.config.get("runner.sync_mode")
+        inner_optimizer = paddle.optimizer.Adam(learning_rate, lazy_mode=True)
+
+        self.role_maker._generate_role()  # 必要
+        if self.config['debug_the_one_ps'] == 1:
+            print("entering run_the_one_ps -- new")
+
+            from paddle.distributed.fleet.meta_optimizers.ps_optimizer import (
+                ParameterServerOptimizer,
+            )
+
+            ps_optimizer = ParameterServerOptimizer(inner_optimizer)
+            ps_optimizer._set_basic_info(
+                loss, self.role_maker, inner_optimizer, user_defined_strategy
+            )
+            ps_optimizer.minimize_impl(loss)
+
+            from paddle.distributed.ps.the_one_ps import TheOnePSRuntime
+
+            _runtime_handle = (
+                TheOnePSRuntime()
+            )  # ps 目录下重构版的 TheOnePSRuntime
+            _runtime_handle._set_basic_info(ps_optimizer.pass_ctx._attrs)
+            if fleet.is_worker():
+                worker_desc = (
+                    _runtime_handle.ps_desc_builder.build_worker_desc()
+                )
+                with open(
+                    ps_log_root_dir + sync_mode + '_' + 'new_worker_ps_desc',
+                    'w',
+                ) as f:
+                    f.write(worker_desc)
+            if fleet.is_server():
+                server_desc = (
+                    _runtime_handle.ps_desc_builder.build_server_desc()
+                )
+                with open(
+                    ps_log_root_dir + sync_mode + '_' + 'new_server_ps_desc',
+                    'w',
+                ) as f:
+                    f.write(server_desc)
+
+        else:
+            pass
+        '''
+            print("entering run_the_one_ps -- old")
+            fleet_obj = fleet.distributed_optimizer(
+                inner_optimizer, user_defined_strategy)
+            fleet_obj.minimize(loss)
+            if fleet.is_worker():
+                worker_desc = fleet_obj._runtime_handle._get_fleet_proto(is_server=False, is_sync=False)
+                server_desc = fleet_obj._runtime_handle._get_fleet_proto(is_server=True, is_sync=False)
+                with open(ps_log_root_dir + sync_mode + '_' + 'worker_ps_desc', 'w') as f:
+                    f.write(str(worker_desc) + str(server_desc))
+            if fleet.is_server():
+                server_desc = fleet_obj._runtime_handle._get_fleet_proto(is_server=True, is_sync=False)
+                with open(ps_log_root_dir + sync_mode + '_' + 'server_ps_desc', 'w') as f:
+                    f.write(str(server_desc) + str(fleet_obj._runtime_handle._get_fs_client_desc().to_string()))
+        '''
+        if fleet.is_server():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + '_run_the_one_ps'
+                + '_debug:_'
+                + str(self.config['debug_the_one_ps'])
+                + '_server_main.prototxt'
+            )
+            debug_program(_main_file, loss.block.program)
+        elif fleet.is_worker():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + '_run_the_one_ps'
+                + '_debug:_'
+                + str(self.config['debug_the_one_ps'])
+                + '_worker_main.prototxt'
+            )
+            debug_program(_main_file, loss.block.program)
+        elif self.role_maker._is_heter_worker():
+            _main_file = (
+                ps_log_root_dir
+                + sync_mode
+                + '_run_the_one_ps'
+                + '_debug:_'
+                + str(self.config['debug_the_one_ps'])
+                + '_heter_worker_main.prototxt'
+            )
+            debug_program(_main_file, loss.block.program)
+
+
+if __name__ == "__main__":
+    paddle.enable_static()
+    config = parse_args()
+    print(">>>>>>>>>> python process started")
+    os.environ["CPU_NUM"] = str(config.get("runner.thread_num"))
+    benchmark_main = DnnTrainer(config)
+    if config['run_single_pass'] == 1:
+        benchmark_main.run_single_pass()
+    elif config['run_minimize'] == 1:
+        benchmark_main.run_minimize()
+    elif config['run_the_one_ps'] == 1:
+        benchmark_main.run_the_one_ps()

test/ps/static_gpubox_trainer.py

@@ -0,0 +1,198 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+import sys
+import time
+
+import paddle
+from paddle.distributed import fleet
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+
+def get_dataset(inputs, config):
+    dataset = paddle.distributed.InMemoryDataset()
+    dataset._set_use_ps_gpu(config.get('runner.use_gpu'))
+    pipe_cmd = config.get('runner.pipe_command')
+    dataset.init(
+        use_var=inputs,
+        pipe_command=pipe_cmd,
+        batch_size=32,
+        thread_num=int(config.get('runner.thread_num')),
+        fs_name=config.get("runner.fs_name", ""),
+        fs_ugi=config.get("runner.fs_ugi", ""),
+    )
+    dataset.set_filelist(["train_data/sample_train.txt"])
+    dataset.update_settings(
+        parse_ins_id=config.get("runner.parse_ins_id", False),
+        parse_content=config.get("runner.parse_content", False),
+    )
+
+    return dataset
+
+
+class Main:
+    def __init__(self):
+        self.metrics = {}
+        self.input_data = None
+        self.reader = None
+        self.exe = None
+        self.model = None
+        self.PSGPU = None
+        self.train_result_dict = {}
+        self.train_result_dict["speed"] = []
+        self.train_result_dict["auc"] = []
+
+    def run(self):
+        from ps_dnn_trainer import YamlHelper
+
+        yaml_helper = YamlHelper()
+        config_yaml_path = 'config_gpubox.yaml'
+        self.config = yaml_helper.load_yaml(config_yaml_path)
+
+        os.environ["CPU_NUM"] = str(self.config.get("runner.thread_num"))
+        fleet.init()
+        self.network()
+        if fleet.is_server():
+            self.run_server()
+        elif fleet.is_worker():
+            self.run_worker()
+            fleet.stop_worker()
+        logger.info("Run Success, Exit.")
+        logger.info("-" * 100)
+
+    def network(self):
+        from ps_dnn_trainer import StaticModel, get_user_defined_strategy
+
+        # self.model = get_model(self.config)
+        self.model = StaticModel(self.config)
+        self.input_data = self.model.create_feeds()
+        self.init_reader()
+        self.metrics = self.model.net(self.input_data)
+        self.inference_target_var = self.model.inference_target_var
+        logger.info("cpu_num: {}".format(os.getenv("CPU_NUM")))
+        # self.model.create_optimizer(get_strategy(self.config)
+        user_defined_strategy = get_user_defined_strategy(self.config)
+        optimizer = paddle.optimizer.Adam(0.01, lazy_mode=True)
+        optimizer = fleet.distributed_optimizer(
+            optimizer, user_defined_strategy
+        )
+        optimizer.minimize(self.model._cost)
+        logger.info("end network.....")
+
+    def run_server(self):
+        logger.info("Run Server Begin")
+        fleet.init_server(self.config.get("runner.warmup_model_path"))
+        fleet.run_server()
+
+    def run_worker(self):
+        logger.info("Run Worker Begin")
+        use_cuda = int(self.config.get("runner.use_gpu"))
+        use_auc = self.config.get("runner.use_auc", False)
+        place = paddle.CUDAPlace(0) if use_cuda else paddle.CPUPlace()
+        self.exe = paddle.static.Executor(place)
+        '''
+        with open("./{}_worker_main_program.prototxt".format(
+                fleet.worker_index()), 'w+') as f:
+            f.write(str(paddle.static.default_main_program()))
+        with open("./{}_worker_startup_program.prototxt".format(
+                fleet.worker_index()), 'w+') as f:
+            f.write(str(paddle.static.default_startup_program()))
+        '''
+        self.exe.run(paddle.static.default_startup_program())
+        fleet.init_worker()
+        '''
+        save_model_path = self.config.get("runner.model_save_path")
+        if save_model_path and (not os.path.exists(save_model_path)):
+            os.makedirs(save_model_path)
+        '''
+        reader_type = self.config.get("runner.reader_type", None)
+        epochs = int(self.config.get("runner.epochs"))
+        sync_mode = self.config.get("runner.sync_mode")
+
+        gpus_env = os.getenv("FLAGS_selected_gpus")
+        self.PSGPU = paddle.framework.core.PSGPU()
+        gpuslot = [int(i) for i in range(1, self.model.sparse_inputs_slots)]
+        gpu_mf_sizes = [self.model.sparse_feature_dim - 1] * (
+            self.model.sparse_inputs_slots - 1
+        )
+        self.PSGPU.set_slot_vector(gpuslot)
+        self.PSGPU.set_slot_dim_vector(gpu_mf_sizes)
+        self.PSGPU.init_gpu_ps([int(s) for s in gpus_env.split(",")])
+        gpu_num = len(gpus_env.split(","))
+        opt_info = paddle.static.default_main_program()._fleet_opt
+        if use_auc is True:
+            opt_info['stat_var_names'] = [
+                self.model.stat_pos.name,
+                self.model.stat_neg.name,
+            ]
+        else:
+            opt_info['stat_var_names'] = []
+
+        for epoch in range(epochs):
+            epoch_start_time = time.time()
+
+            self.dataset_train_loop(epoch)
+
+            epoch_time = time.time() - epoch_start_time
+
+            self.PSGPU.end_pass()
+
+            fleet.barrier_worker()
+            self.reader.release_memory()
+            logger.info(f"finish {epoch} epoch training....")
+        self.PSGPU.finalize()
+
+    def init_reader(self):
+        if fleet.is_server():
+            return
+        # self.reader, self.file_list = get_reader(self.input_data, config)
+        self.reader = get_dataset(self.input_data, self.config)
+
+    def dataset_train_loop(self, epoch):
+        start_time = time.time()
+        self.reader.load_into_memory()
+        print(
+            f"self.reader.load_into_memory cost :{time.time() - start_time} seconds"
+        )
+
+        begin_pass_time = time.time()
+        self.PSGPU.begin_pass()
+        print(f"begin_pass cost:{time.time() - begin_pass_time} seconds")
+
+        logger.info(f"Epoch: {epoch}, Running Dataset Begin.")
+        fetch_info = [
+            f"Epoch {epoch} Var {var_name}" for var_name in self.metrics
+        ]
+        print_step = int(self.config.get("runner.print_interval"))
+        self.exe.train_from_dataset(
+            program=paddle.static.default_main_program(),
+            dataset=self.reader,
+            debug=self.config.get("runner.dataset_debug"),
+        )
+
+
+if __name__ == "__main__":
+    paddle.enable_static()
+    benchmark_main = Main()
+    benchmark_main.run()

test/ps/test_fl_ps.py

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shlex  # noqa: F401
+import sys
+import unittest
+
+sys.path.append("../distributed_passes")
+from dist_pass_test_base import remove_path_if_exists
+
+
+class FlPsTest(unittest.TestCase):
+    def test_launch_fl_ps(self):
+        '''
+        cmd = [
+            'python', '-m', 'paddle.distributed.fleet.launch', '--log_dir',
+            '/ps_log/fl_ps', '--servers', "127.0.0.1:8070", '--workers',
+            "127.0.0.1:8080,127.0.0.1:8081", '--heter_workers',
+            "127.0.0.1:8090,127.0.0.1:8091", '--heter_devices', "cpu",
+            '--worker_num', "2", '--heter_worker_num', "2", 'fl_ps_trainer.py'
+        ]
+        cmd = [shlex.quote(c) for c in cmd]
+        prepare_python_path_and_return_module(__file__)
+        exitcode = os.system(' '.join(cmd))
+        '''
+
+
+if __name__ == '__main__':
+    remove_path_if_exists('/ps_log')
+    remove_path_if_exists('/ps_usr_print_log')
+    if not os.path.exists('./train_data'):
+        os.system('sh download_data.sh')
+        os.system('rm -rf ctr_data.tar.gz')
+        os.system('rm -rf train_data_full')
+        os.system('rm -rf test_data_full')
+    unittest.main()
+    if os.path.exists('./train_data'):
+        os.system('rm -rf train_data')
+        os.system('rm -rf test_data')

test/ps/test_gpubox_ps.py

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import shlex  # noqa: F401
+import unittest
+
+
+class GpuBoxTest(unittest.TestCase):
+    def test_gpubox(self):
+        if not os.path.exists('./train_data'):
+            os.system('bash download_criteo_data.sh')
+
+        exitcode = os.system('bash gpubox_run.sh')
+        if os.path.exists('./train_data'):
+            os.system('rm -rf train_data')
+        if exitcode:
+            os.system('cat ./log/worker.0.log')
+
+        assert exitcode == 0
+
+
+if __name__ == '__main__':
+    unittest.main()

test/ps/test_the_one_ps.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import unittest
+
+from google.protobuf import text_format  # noqa: F401
+
+sys.path.append("../distributed_passes")
+from ps_pass_test_base import PsPassTestBase, remove_path_if_exists
+
+import paddle.distributed.fleet.proto.the_one_ps_pb2 as ps_pb2  # noqa: F401
+from paddle.distributed.ps.utils.public import logger, ps_log_root_dir
+
+
+class TestTheOnePs(PsPassTestBase):
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def check(self, file1, file2):
+        '''
+        f = open(file1, "rb")
+        ps_desc_1 = ps_pb2.PSParameter()
+        text_format.Parse(f.read(), ps_desc_1)
+        f.close()
+
+        f = open(file2, "rb")
+        ps_desc_2 = ps_pb2.PSParameter()
+        text_format.Parse(f.read(), ps_desc_2)
+        f.close()
+        str1 = text_format.MessageToString(ps_desc_1)
+        str2 = text_format.MessageToString(ps_desc_2)
+        #logger.info('### msg10: {}'.format(str1))
+        #logger.info('### msg20: {}'.format(str2))
+        if str1 == str2:
+            return True
+        else:
+            return False
+        '''
+        pass
+
+    def test_ps_cpu_async(self):
+        self.init()
+        self.config['ps_mode_config'] = "../ps/cpu_async_ps_config.yaml"
+        self.config['run_the_one_ps'] = '1'
+
+        self.config['debug_the_one_ps'] = '0'
+        self.config['log_dir'] = (
+            ps_log_root_dir + "async_cpu_log_old_the_one_ps"
+        )
+        remove_path_if_exists(self.config['log_dir'])
+        self.ps_launch()
+
+        self.config['debug_the_one_ps'] = '1'
+        self.config['log_dir'] = (
+            ps_log_root_dir + "async_cpu_log_new_the_one_ps"
+        )
+        remove_path_if_exists(self.config['log_dir'])
+        self.ps_launch()
+
+        desc1 = '/ps_desc_baseline/async_worker_ps_desc'
+        desc2 = '/ps_log/async_new_worker_ps_desc'
+        desc3 = '/ps_desc_baseline/async_server_ps_desc'
+        desc4 = '/ps_log/async_new_server_ps_desc'
+        if self.check(desc1, desc2):
+            logger.info('test_ps_cpu_async ps_desc: worker passed!')
+        else:
+            logger.info('test_ps_cpu_async ps_desc: worker failed!')
+        if self.check(desc3, desc4):
+            logger.info('test_ps_cpu_async ps_desc: server passed!')
+        else:
+            logger.info('test_ps_cpu_async ps_desc: server failed!')
+
+    def test_ps_cpu_geo(self):
+        self.init()
+        self.config['ps_mode_config'] = "../ps/cpu_geo_ps_config.yaml"
+        self.config['run_the_one_ps'] = '1'
+
+        self.config['debug_the_one_ps'] = '0'
+        self.config['log_dir'] = ps_log_root_dir + "geo_cpu_log_old_the_one_ps"
+        remove_path_if_exists(self.config['log_dir'])
+        self.ps_launch()
+
+        self.config['debug_the_one_ps'] = '1'
+        self.config['log_dir'] = ps_log_root_dir + "geo_cpu_log_new_the_one_ps"
+        remove_path_if_exists(self.config['log_dir'])
+        self.ps_launch()
+
+        desc1 = '/ps_desc_baseline/geo_worker_ps_desc'
+        desc2 = '/ps_log/geo_new_worker_ps_desc'
+        desc3 = '/ps_desc_baseline/geo_server_ps_desc'
+        desc4 = '/ps_log/geo_new_server_ps_desc'
+        if self.check(desc1, desc2):
+            logger.info('test_ps_cpu_geo ps_desc: worker passed!')
+        else:
+            logger.info('test_ps_cpu_geo ps_desc: worker failed!')
+        if self.check(desc3, desc4):
+            logger.info('test_ps_cpu_geo ps_desc: server passed!')
+        else:
+            logger.info('test_ps_cpu_geo ps_desc: server failed!')
+
+
+if __name__ == '__main__':
+    unittest.main()